mdarray-linalg 0.2.0

use approx::assert_relative_eq;
use mdarray::{DArray, Dyn};
use num_complex::{Complex, ComplexFloat};
use rand::Rng;

use super::common::{assert_complex_matrix_eq, assert_matrix_eq, naive_matmul};
use crate::{
    svd::{SVD, SVDDecomp},
    utils::pretty_print,
};

// These reconstruction tests assume that the backend represents singular
// values with the same scalar type as the input matrix. Backends that use a
// different `SingularValue` type will need a conversion-aware test helper.
fn test_svd_reconstruction<T>(
    bd: &impl SVD<T, Dyn, SingularValue = T>,
    a: &DArray<T, 2>,
    debug_print: bool,
)
where
    T: ComplexFloat<Real = f64>
        + Default
        + Copy
        + std::fmt::Debug
        + approx::AbsDiffEq<Epsilon = T::Real>
        + std::fmt::Display
        + approx::RelativeEq,
    T::Real: std::fmt::Display,
{
    let (m, n) = (a.shape().0, a.shape().1);
    let min_dim = m.min(n);

    let SVDDecomp { s, u, vt } = bd.svd(&mut a.clone()).expect("SVD failed");

    let mut sigma = DArray::<T, 2>::zeros([m, n]);
    for i in 0..min_dim {
        sigma[[i, i]] = s[i];
    }

    // let mut sigma = diag(&s);

    if debug_print {
        println!("=== A original ===");
        pretty_print(a);
        println!("=== Σ (Sigma) ===");
        pretty_print(&sigma);
        println!("=== U ===");
        pretty_print(&u);
        println!("=== Vᵀ ===");
        pretty_print(&vt);
    }

    let us = naive_matmul(&u, &sigma);
    if debug_print {
        println!("=== U × Σ ===");
        pretty_print(&us);
    }

    let usvt = naive_matmul(&us, &vt);
    if debug_print {
        println!("=== U × Σ × Vᵀ  ===");
        pretty_print(&usvt);
        println!("=== A original ===");
        pretty_print(a);
    }

    assert_matrix_eq!(*a, usvt);
}

pub fn test_svd_square_matrix(bd: &impl SVD<f64, Dyn, SingularValue = f64>) {
    let n = 3;
    let a = DArray::<f64, 2>::from_fn([n, n], |i| (i[0] * i[1]) as f64);
    test_svd_reconstruction(bd, &a, true);
}

pub fn test_svd_rectangular_m_gt_n(bd: &impl SVD<f64, Dyn, SingularValue = f64>) {
    let (m, n) = (4, 3);
    let a = DArray::<f64, 2>::from_fn([m, n], |i| (i[0] * i[1]) as f64);
    test_svd_reconstruction(bd, &a, true);
}

pub fn test_svd_rectangular_n_gt_m(bd: &impl SVD<f64, Dyn, SingularValue = f64>) {
    let (m, n) = (3, 4);
    let a = DArray::<f64, 2>::from_fn([m, n], |i| (i[0] * i[1]) as f64);
    test_svd_reconstruction(bd, &a, true);
}

pub fn test_svd_big_square_matrix(bd: &impl SVD<f64, Dyn, SingularValue = f64>) {
    let n = 200;
    let a = DArray::<f64, 2>::from_fn([n, n], |i| (i[0] * i[1]) as f64);
    test_svd_reconstruction(bd, &a, false);
}

pub fn test_svd_random_matrix(bd: &impl SVD<f64, Dyn, SingularValue = f64>) {
    let mut rng = rand::rng();
    let n = 4;
    let a = DArray::<f64, 2>::from_fn([n, n], |_| rng.random::<f64>());
    test_svd_reconstruction(bd, &a, true);
}

pub fn test_svd_cplx_square_matrix(
    bd: &impl SVD<Complex<f64>, Dyn, SingularValue = Complex<f64>>,
) {
    let n = 3;
    let a = DArray::<Complex<f64>, 2>::from_fn([n, n], |i| {
        Complex::new((i[0] * i[1]) as f64, i[1] as f64)
    });

    let SVDDecomp { s, u, vt } = bd.svd(&mut a.clone()).expect("SVD failed");

    // assert_eq!(*s.shape(), (n,));
    // assert_eq!(*u.shape(), (n, n));
    // assert_eq!(*vt.shape(), (n, n));

    let mut sigma = DArray::<Complex<f64>, 2>::zeros([n, n]);
    for i in 0..n {
        sigma[[i, i]] = s[i];
    }

    println!("=== Σ (Sigma) ===");
    pretty_print(&sigma);
    println!("=== U ===");
    pretty_print(&u);
    println!("=== Vᵀ ===");
    pretty_print(&vt);

    let us = naive_matmul(&u, &sigma);
    println!("=== U × Σ ===");
    pretty_print(&us);
    let usvt = naive_matmul(&us, &vt);
    println!("=== U × Σ × Vᵀ  ===");
    pretty_print(&usvt);
    println!("=== A original ===");
    pretty_print(&a);

    assert_complex_matrix_eq!(a, usvt);
}

/// Test complex SVD with random matrix having significant imaginary parts.
/// This test is specifically designed to catch the V^T vs V^H bug.
pub fn test_svd_cplx_random_matrix(
    bd: &impl SVD<Complex<f64>, Dyn, SingularValue = Complex<f64>>,
) {
    let mut rng = rand::rng();
    let n = 5;

    // Create random complex matrix with significant imaginary parts
    let a = DArray::<Complex<f64>, 2>::from_fn([n, n], |_| {
        Complex::new(
            rng.random::<f64>() * 2.0 - 1.0,
            rng.random::<f64>() * 2.0 - 1.0,
        )
    });

    let SVDDecomp { s, u, vt } = bd.svd(&mut a.clone()).expect("SVD failed");

    // Build sigma matrix
    let mut sigma = DArray::<Complex<f64>, 2>::zeros([n, n]);
    for i in 0..n {
        sigma[[i, i]] = s[i];
    }

    // Reconstruct: A = U * Σ * V^H (vt should be V^H)
    let us = naive_matmul(&u, &sigma);
    let usvt = naive_matmul(&us, &vt);

    assert_complex_matrix_eq!(a, usvt);
}