geographdb_core/algorithms/

mpo.rs

use crate::algorithms::four_d::{GraphNode4D, GraphProperties, TemporalEdge};

/// Matrix Product Operator (MPO) for weight matrix compression.
///
/// An MPO represents a linear map W: ℝ^n_in → ℝ^n_out as a chain of rank-4
/// tensors. Each site has shape [chi_left, d_out, d_in, chi_right].
/// Forward pass costs O(n · d² · χ³) instead of O(n_out · n_in).
///
/// Construction: reshape W as a rank-2k tensor with interleaved output/input
/// indices, then decompose left-to-right via sequential SVD.
/// One site in an MPO chain. Data is row-major [chi_left, d_out, d_in, chi_right].
#[derive(Debug, Clone)]
pub struct MpoSite {
    pub chi_left: usize,
    pub d_out: usize,
    pub d_in: usize,
    pub chi_right: usize,
    pub data: Vec<f32>,
}

impl MpoSite {
    fn get(&self, il: usize, io: usize, ii: usize, ir: usize) -> f32 {
        self.data[il * self.d_out * self.d_in * self.chi_right
            + io * self.d_in * self.chi_right
            + ii * self.chi_right
            + ir]
    }
}

/// A Matrix Product Operator representing a compressed weight matrix.
#[derive(Debug, Clone)]
pub struct Mpo {
    pub sites: Vec<MpoSite>,
}

// ── Linear algebra helpers ────────────────────────────────────────────────────

fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {
    let mut c = vec![0.0f32; m * n];
    for i in 0..m {
        for p in 0..k {
            let aip = a[i * k + p];
            for j in 0..n {
                c[i * n + j] += aip * b[p * n + j];
            }
        }
    }
    c
}

fn transpose(a: &[f32], m: usize, n: usize) -> Vec<f32> {
    let mut t = vec![0.0f32; n * m];
    for i in 0..m {
        for j in 0..n {
            t[j * m + i] = a[i * n + j];
        }
    }
    t
}

/// Thin SVD of A (m×n). Returns (U m×k, sigma k, Vt k×n) where k = min(m,n).
/// Singular values are descending. U columns and Vt rows are orthonormal.
pub fn svd_thin(a: &[f32], m: usize, n: usize) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
    if m < n {
        // Work on A^T (n×m, n >= m), then swap back.
        let at = transpose(a, m, n);
        let (v, sigma, ut) = svd_thin_tall(&at, n, m);
        // A = U * S * Vt  where U = Vt_t^T (m×m) and Vt = U_t^T (m×n)
        let u = transpose(&ut, m, m);
        let vt = transpose(&v, n, m);
        return (u, sigma, vt);
    }
    svd_thin_tall(a, m, n)
}

/// Thin SVD of A (m×n, m ≥ n) via Jacobi iteration on A^T A.
fn svd_thin_tall(a: &[f32], m: usize, n: usize) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
    assert!(m >= n, "svd_thin_tall requires m >= n, got m={m} n={n}");

    // C = A^T · A  (n×n symmetric PSD)
    let at = transpose(a, m, n);
    let mut c = matmul(&at, a, n, m, n);

    // One-sided Jacobi: accumulate eigenvectors in v, converge diagonal of c to eigenvalues.
    let mut v = vec![0.0f32; n * n];
    for i in 0..n {
        v[i * n + i] = 1.0;
    }

    // Jacobi sweep iteration
    for _ in 0..100 {
        let mut converged = true;
        for p in 0..n {
            for q in (p + 1)..n {
                let cpq = c[p * n + q];
                if cpq.abs() > 1e-15 {
                    converged = false;
                    let cpp = c[p * n + p];
                    let cqq = c[q * n + q];
                    let theta = (cqq - cpp) / (2.0 * cpq);
                    let t = 1.0 / (theta.abs() + (1.0 + theta * theta).sqrt());
                    let t = if theta < 0.0 { -t } else { t };
                    let cos = 1.0 / (1.0 + t * t).sqrt();
                    let sin = t * cos;
                    let tau = sin / (1.0 + cos);

                    // Update symmetric c elements
                    c[p * n + p] = cpp - t * cpq;
                    c[q * n + q] = cqq + t * cpq;
                    c[p * n + q] = 0.0;
                    c[q * n + p] = 0.0;

                    for r in 0..n {
                        if r != p && r != q {
                            let crp = c[r * n + p];
                            let crq = c[r * n + q];
                            c[r * n + p] = crp - sin * (crq + tau * crp);
                            c[r * n + q] = crq + sin * (crp - tau * crq);
                            // Symmetrize
                            c[p * n + r] = c[r * n + p];
                            c[q * n + r] = c[r * n + q];
                        }
                    }

                    // Update V matrix (right singular vectors/eigenvectors)
                    for i in 0..n {
                        let vip = v[i * n + p];
                        let viq = v[i * n + q];
                        v[i * n + p] = vip - sin * (viq + tau * vip);
                        v[i * n + q] = viq + sin * (vip - tau * viq);
                    }
                }
            }
        }
        if converged {
            break;
        }
    }

    // Eigenvalues sit on the diagonal of c.
    let mut sigma: Vec<f32> = (0..n).map(|i| c[i * n + i].max(0.0).sqrt()).collect();

    // Sort descending.
    let mut order: Vec<usize> = (0..n).collect();
    order.sort_by(|&a, &b| sigma[b].partial_cmp(&sigma[a]).unwrap());
    sigma = order.iter().map(|&i| sigma[i]).collect();
    let mut v_sorted = vec![0.0f32; n * n];
    for (new_col, &old_col) in order.iter().enumerate() {
        for row in 0..n {
            v_sorted[row * n + new_col] = v[row * n + old_col];
        }
    }

    // U = A · V · diag(1/sigma)
    let av = matmul(a, &v_sorted, m, n, n);
    let mut u = vec![0.0f32; m * n];
    for col in 0..n {
        let inv_s = if sigma[col] > 1e-10 {
            1.0 / sigma[col]
        } else {
            0.0
        };
        for row in 0..m {
            u[row * n + col] = av[row * n + col] * inv_s;
        }
    }

    let vt = transpose(&v_sorted, n, n);
    (u, sigma, vt)
}

fn nth_root(n: usize, k: usize) -> usize {
    if k == 1 {
        return n;
    }
    let r = (n as f64).powf(1.0 / k as f64).round() as usize;
    for candidate in r.saturating_sub(2)..=r + 2 {
        if candidate.pow(k as u32) == n {
            return candidate;
        }
    }
    r
}

/// Factor `d` into a list of small integer factors (each ≥ 2).
///
/// Prefers 4, then 2, then odd factors. Returns `vec![d]` for primes.
pub fn factor_dimension(d: usize) -> Vec<usize> {
    if d < 2 {
        return vec![d];
    }
    let mut remaining = d;
    let mut factors = Vec::new();
    while remaining.is_multiple_of(4) {
        factors.push(4);
        remaining /= 4;
    }
    while remaining.is_multiple_of(2) {
        factors.push(2);
        remaining /= 2;
    }
    let mut odd = 3;
    while odd * odd <= remaining {
        while remaining.is_multiple_of(odd) {
            factors.push(odd);
            remaining /= odd;
        }
        odd += 2;
    }
    if remaining > 1 {
        factors.push(remaining);
    }
    factors
}

/// Pad a factorization to exactly `target_len` elements by prepending 1s.
fn pad_factors(factors: &mut Vec<usize>, target_len: usize) {
    while factors.len() < target_len {
        factors.insert(0, 1);
    }
}

// ── MPO construction ──────────────────────────────────────────────────────────

/// Compress a dense weight matrix into an MPO via sequential SVD.
///
/// `weights`: row-major n_out × n_in matrix.
/// `n_sites`: number of sites. Requires n_in = d_in^n_sites, n_out = d_out^n_sites.
/// `chi_max`: maximum bond dimension.
pub fn compress_matrix_to_mpo(
    weights: &[f32],
    n_out: usize,
    n_in: usize,
    n_sites: usize,
    chi_max: usize,
) -> Mpo {
    assert_eq!(weights.len(), n_out * n_in);
    assert!(n_sites >= 1);
    let d_in = nth_root(n_in, n_sites);
    let d_out = nth_root(n_out, n_sites);
    assert_eq!(d_in.pow(n_sites as u32), n_in);
    assert_eq!(d_out.pow(n_sites as u32), n_out);

    // Reshape W (n_out × n_in) into W_T with interleaved indices:
    // W_T[io_0, ii_0, io_1, ii_1, ..., io_{n-1}, ii_{n-1}]
    let phys = d_out * d_in; // local Hilbert space size per site
    let do_strides: Vec<usize> = (0..n_sites)
        .map(|k| d_out.pow((n_sites - 1 - k) as u32))
        .collect();
    let di_strides: Vec<usize> = (0..n_sites)
        .map(|k| d_in.pow((n_sites - 1 - k) as u32))
        .collect();

    let mut w_t = vec![0.0f32; n_out * n_in];
    for row in 0..n_out {
        for col in 0..n_in {
            let mut r = row;
            let mut c = col;
            let mut interleaved = 0usize;
            let mut pair_stride = phys.pow((n_sites - 1) as u32);
            for k in 0..n_sites {
                let io_k = r / do_strides[k];
                r %= do_strides[k];
                let ii_k = c / di_strides[k];
                c %= di_strides[k];
                interleaved += (io_k * d_in + ii_k) * pair_stride;
                if k + 1 < n_sites {
                    pair_stride /= phys;
                }
            }
            w_t[interleaved] = weights[row * n_in + col];
        }
    }

    // Sequential left-to-right SVD decomposition.
    // At step k: current has shape (chi_left, phys^(n_sites-k))
    let mut current = w_t;
    let mut chi_left = 1usize;
    let mut sites = Vec::with_capacity(n_sites);

    for k in 0..n_sites {
        let total_right = phys.pow((n_sites - k) as u32);
        let m_svd = chi_left * phys;
        let n_svd = total_right / phys; // = phys^(n_sites-k-1)

        // Unfold: (chi_left, phys, n_svd) → (chi_left*phys, n_svd)
        let mut unfolded = vec![0.0f32; m_svd * n_svd];
        for il in 0..chi_left {
            for ip in 0..phys {
                let row = il * phys + ip;
                for col in 0..n_svd {
                    unfolded[row * n_svd + col] = current[il * total_right + ip * n_svd + col];
                }
            }
        }

        // SVD: always orient so m >= n for svd_thin.
        let (u, sigma, vt, k_svd) = if m_svd >= n_svd {
            let (u, s, vt) = svd_thin(&unfolded, m_svd, n_svd);
            let k = s.len();
            (u, s, vt, k)
        } else {
            let ut = transpose(&unfolded, m_svd, n_svd);
            let (v, s, ut2) = svd_thin(&ut, n_svd, m_svd);
            let k = s.len();
            let u2 = transpose(&ut2, m_svd, m_svd);
            let vt2 = transpose(&v, n_svd, m_svd);
            (u2, s, vt2, k)
        };

        const SVD_EPS: f32 = 1e-9;
        let chi_new = sigma
            .iter()
            .filter(|&&s| s > SVD_EPS)
            .count()
            .min(chi_max)
            .max(1);
        let chi_right = if k == n_sites - 1 { 1 } else { chi_new };

        // Site tensor: M[il, io, ii, ir] = U[il*phys + io*d_in + ii, ir] * sigma[ir]
        let mut data = vec![0.0f32; chi_left * d_out * d_in * chi_right];
        for il in 0..chi_left {
            for io in 0..d_out {
                for ii in 0..d_in {
                    let u_row = il * phys + io * d_in + ii;
                    for ir in 0..chi_right {
                        let u_val = if ir < k_svd {
                            u[u_row * k_svd + ir]
                        } else {
                            0.0
                        };
                        let s_val = if ir < k_svd { sigma[ir] } else { 0.0 };
                        data[il * d_out * d_in * chi_right
                            + io * d_in * chi_right
                            + ii * chi_right
                            + ir] = u_val * s_val;
                    }
                }
            }
        }
        sites.push(MpoSite {
            chi_left,
            d_out,
            d_in,
            chi_right,
            data,
        });

        // Remainder: Vt[:chi_right, :] becomes current (chi_right × n_svd).
        if k < n_sites - 1 {
            current = (0..chi_right * n_svd)
                .map(|idx| {
                    let ir = idx / n_svd;
                    let col = idx % n_svd;
                    if ir < k_svd {
                        vt[ir * n_svd + col]
                    } else {
                        0.0
                    }
                })
                .collect();
        }
        chi_left = chi_right;
    }

    Mpo { sites }
}

// ── MPO forward pass ──────────────────────────────────────────────────────────

/// Apply the MPO as a forward pass: y = W_mpo · x.
pub fn mpo_apply(mpo: &Mpo, x: &[f32]) -> Vec<f32> {
    let n_sites = mpo.sites.len();
    if n_sites == 0 {
        return Vec::new();
    }

    let d_out_dims: Vec<usize> = mpo.sites.iter().map(|s| s.d_out).collect();
    let d_in_dims: Vec<usize> = mpo.sites.iter().map(|s| s.d_in).collect();
    let n_out: usize = d_out_dims.iter().product();
    let n_in: usize = d_in_dims.iter().product();
    assert_eq!(x.len(), n_in);

    // Strides for decoding flat output/input indices into per-site indices.
    let out_strides: Vec<usize> = {
        let mut s = vec![1usize; n_sites];
        for k in (0..n_sites - 1).rev() {
            s[k] = s[k + 1] * d_out_dims[k + 1];
        }
        s
    };
    let in_strides: Vec<usize> = {
        let mut s = vec![1usize; n_sites];
        for k in (0..n_sites - 1).rev() {
            s[k] = s[k + 1] * d_in_dims[k + 1];
        }
        s
    };

    let mut y = vec![0.0f32; n_out];

    // y[io_0,...,io_{n-1}] = Σ_{ii, bonds} [∏_k M_k[bond_{k-1}, io_k, ii_k, bond_k]] · x[ii]
    // Contract bond chain for each (out_idx, in_idx) pair.
    for (out_idx, y_out) in y.iter_mut().enumerate() {
        let io: Vec<usize> = (0..n_sites)
            .map(|k| (out_idx / out_strides[k]) % d_out_dims[k])
            .collect();
        for (in_idx, &x_in) in x.iter().enumerate() {
            let ii: Vec<usize> = (0..n_sites)
                .map(|k| (in_idx / in_strides[k]) % d_in_dims[k])
                .collect();
            // Trace bond chain.
            let mut bond = vec![1.0f32]; // chi=1 left boundary
            for k in 0..n_sites {
                let site = &mpo.sites[k];
                let chi_r = site.chi_right;
                let mut next = vec![0.0f32; chi_r];
                for (il, &bond_val) in bond.iter().enumerate() {
                    for (ir, next_val) in next.iter_mut().enumerate() {
                        *next_val += bond_val * site.get(il, io[k], ii[k], ir);
                    }
                }
                bond = next;
            }
            *y_out += bond[0] * x_in;
        }
    }

    y
}

// ── utilities ─────────────────────────────────────────────────────────────────

/// Frobenius reconstruction error: ‖W_mpo − W‖_F / ‖W‖_F.
pub fn mpo_reconstruction_error(mpo: &Mpo, original: &[f32], n_out: usize, n_in: usize) -> f32 {
    let mut diff_sq = 0.0f32;
    let mut norm_sq = 0.0f32;

    let mut e_j = vec![0.0f32; n_in];
    for j in 0..n_in {
        e_j[j] = 1.0;
        let col = mpo_apply(mpo, &e_j);
        e_j[j] = 0.0;
        for i in 0..n_out {
            let d = col[i] - original[i * n_in + j];
            diff_sq += d * d;
            norm_sq += original[i * n_in + j] * original[i * n_in + j];
        }
    }

    if norm_sq < 1e-10 {
        return 0.0;
    }
    (diff_sq / norm_sq).sqrt()
}

/// Compression ratio: mpo parameter count / (n_out * n_in).
pub fn mpo_compression_ratio(mpo: &Mpo, n_out: usize, n_in: usize) -> f32 {
    let mpo_params: usize = mpo
        .sites
        .iter()
        .map(|s| s.chi_left * s.d_out * s.d_in * s.chi_right)
        .sum();
    mpo_params as f32 / (n_out * n_in) as f32
}

/// Compress a dense weight matrix into an MPO with automatic dimension handling.
///
/// Unlike [`compress_matrix_to_mpo`], this function accepts *any* `n_out × n_in`
/// matrix. It factors each dimension into a product of small integers, pads to
/// equal site counts with trivial (dim-1) sites, zero-pads the matrix to the
/// nearest compatible shape, compresses, then trims the output.
///
/// `chi_max` controls the maximum bond dimension (compression quality).
pub fn compress_matrix_to_mpo_auto(
    weights: &[f32],
    n_out: usize,
    n_in: usize,
    chi_max: usize,
) -> Mpo {
    assert_eq!(weights.len(), n_out * n_in);
    assert!(n_out >= 1 && n_in >= 1);

    let mut out_factors = factor_dimension(n_out);
    let mut in_factors = factor_dimension(n_in);

    let n_sites = out_factors.len().max(in_factors.len());
    pad_factors(&mut out_factors, n_sites);
    pad_factors(&mut in_factors, n_sites);

    let padded_out: usize = out_factors.iter().product();
    let padded_in: usize = in_factors.iter().product();

    let padded_weights = if padded_out == n_out && padded_in == n_in {
        weights.to_vec()
    } else {
        let mut padded = vec![0.0f32; padded_out * padded_in];
        for r in 0..n_out {
            let src = &weights[r * n_in..(r + 1) * n_in];
            let dst = &mut padded[r * padded_in..(r + 1) * padded_in];
            dst[..n_in].copy_from_slice(src);
        }
        padded
    };

    let mpo = compress_variable_mpo(
        &padded_weights,
        padded_out,
        padded_in,
        &out_factors,
        &in_factors,
        chi_max,
    );

    trim_mpo(&mpo, n_out, n_in, padded_out, padded_in)
}

/// Internal: compress with variable per-site dimensions.
fn compress_variable_mpo(
    weights: &[f32],
    n_out: usize,
    n_in: usize,
    d_out_dims: &[usize],
    d_in_dims: &[usize],
    chi_max: usize,
) -> Mpo {
    let n_sites = d_out_dims.len();
    assert_eq!(d_in_dims.len(), n_sites);
    assert_eq!(weights.len(), n_out * n_in);

    let phys_dims: Vec<usize> = (0..n_sites).map(|k| d_out_dims[k] * d_in_dims[k]).collect();

    let do_strides: Vec<usize> = {
        let mut s = vec![0usize; n_sites];
        let mut stride = 1;
        for k in (0..n_sites).rev() {
            s[k] = stride;
            stride *= d_out_dims[k];
        }
        s
    };
    let di_strides: Vec<usize> = {
        let mut s = vec![0usize; n_sites];
        let mut stride = 1;
        for k in (0..n_sites).rev() {
            s[k] = stride;
            stride *= d_in_dims[k];
        }
        s
    };

    let pair_strides: Vec<usize> = {
        let mut s = vec![0usize; n_sites];
        let mut stride = 1;
        for k in (0..n_sites).rev() {
            s[k] = stride;
            stride *= phys_dims[k];
        }
        s
    };

    let mut w_t = vec![0.0f32; n_out * n_in];
    for row in 0..n_out {
        for col in 0..n_in {
            let mut interleaved = 0usize;
            for k in 0..n_sites {
                let io_k = (row / do_strides[k]) % d_out_dims[k];
                let ii_k = (col / di_strides[k]) % d_in_dims[k];
                interleaved += (io_k * d_in_dims[k] + ii_k) * pair_strides[k];
            }
            w_t[interleaved] = weights[row * n_in + col];
        }
    }

    let mut current = w_t;
    let mut chi_left = 1usize;
    let mut sites = Vec::with_capacity(n_sites);

    for k in 0..n_sites {
        let phys_k = phys_dims[k];
        let d_out_k = d_out_dims[k];
        let d_in_k = d_in_dims[k];
        let total_right: usize = phys_dims[k..].iter().product();
        let n_svd = total_right / phys_k;
        let m_svd = chi_left * phys_k;

        let mut unfolded = vec![0.0f32; m_svd * n_svd];
        for il in 0..chi_left {
            for ip in 0..phys_k {
                let row = il * phys_k + ip;
                for col in 0..n_svd {
                    unfolded[row * n_svd + col] = current[il * total_right + ip * n_svd + col];
                }
            }
        }

        let (u, sigma, vt, k_svd) = if m_svd >= n_svd {
            let (u, s, vt) = svd_thin(&unfolded, m_svd, n_svd);
            let k = s.len();
            (u, s, vt, k)
        } else {
            let ut = transpose(&unfolded, m_svd, n_svd);
            let (v, s, ut2) = svd_thin(&ut, n_svd, m_svd);
            let k = s.len();
            let u2 = transpose(&ut2, m_svd, m_svd);
            let vt2 = transpose(&v, n_svd, m_svd);
            (u2, s, vt2, k)
        };

        const SVD_EPS: f32 = 1e-9;
        let chi_new = sigma
            .iter()
            .filter(|&&s| s > SVD_EPS)
            .count()
            .min(chi_max)
            .max(1);
        let chi_right = if k == n_sites - 1 { 1 } else { chi_new };

        let mut data = vec![0.0f32; chi_left * d_out_k * d_in_k * chi_right];
        for il in 0..chi_left {
            for io in 0..d_out_k {
                for ii in 0..d_in_k {
                    let u_row = il * phys_k + io * d_in_k + ii;
                    for ir in 0..chi_right {
                        let u_val = if ir < k_svd {
                            u[u_row * k_svd + ir]
                        } else {
                            0.0
                        };
                        let s_val = if ir < k_svd { sigma[ir] } else { 0.0 };
                        data[il * d_out_k * d_in_k * chi_right
                            + io * d_in_k * chi_right
                            + ii * chi_right
                            + ir] = u_val * s_val;
                    }
                }
            }
        }
        sites.push(MpoSite {
            chi_left,
            d_out: d_out_k,
            d_in: d_in_k,
            chi_right,
            data,
        });

        if k < n_sites - 1 {
            current = (0..chi_right * n_svd)
                .map(|idx| {
                    let ir = idx / n_svd;
                    let col = idx % n_svd;
                    if ir < k_svd {
                        vt[ir * n_svd + col]
                    } else {
                        0.0
                    }
                })
                .collect();
        }
        chi_left = chi_right;
    }

    Mpo { sites }
}

/// Trim an MPO that was built from a zero-padded matrix.
///
/// Sites with `d_in == 1` or `d_out == 1` (trivial padding sites) are removed.
/// The remaining sites are re-chained with correct bond dimensions.
fn trim_mpo(mpo: &Mpo, orig_out: usize, orig_in: usize, pad_out: usize, pad_in: usize) -> Mpo {
    if orig_out == pad_out && orig_in == pad_in {
        return mpo.clone();
    }

    let mut kept: Vec<MpoSite> = mpo
        .sites
        .iter()
        .filter(|s| s.d_out > 1 || s.d_in > 1)
        .cloned()
        .collect();

    if kept.is_empty() {
        kept = vec![MpoSite {
            chi_left: 1,
            d_out: orig_out,
            d_in: orig_in,
            chi_right: 1,
            data: vec![0.0; orig_out * orig_in],
        }];
    }

    kept[0].chi_left = 1;
    let last = kept.last_mut().unwrap();
    if last.chi_right != 1 {
        last.chi_right = 1;
        last.data.truncate(last.chi_left * last.d_out * last.d_in);
    }

    Mpo { sites: kept }
}

/// Convert an MPO to a chain of `GraphNode4D` nodes compatible with the MPS graph infrastructure.
///
/// Each site is stored as a graph node with its tensor data. Sites are linked by
/// `TemporalEdge`s whose `weight` encodes the bond dimension. This enables MPO
/// sites to participate in the same spatial-temporal graph queries as MPS states.
pub fn mpo_to_graph_nodes(mpo: &Mpo) -> Vec<GraphNode4D> {
    mpo.sites
        .iter()
        .enumerate()
        .map(|(i, site)| {
            let mut props = GraphProperties::new();
            props.insert(
                "shape".into(),
                serde_json::Value::Array(vec![
                    serde_json::Value::from(site.chi_left as u64),
                    serde_json::Value::from(site.d_out as u64),
                    serde_json::Value::from(site.d_in as u64),
                    serde_json::Value::from(site.chi_right as u64),
                ]),
            );
            props.insert(
                "mpo_kind".into(),
                serde_json::Value::String("mpo_site".into()),
            );
            let data_json = serde_json::Value::Array(
                site.data
                    .iter()
                    .map(|&x| serde_json::Value::from(x as f64))
                    .collect(),
            );
            props.insert("data".into(), data_json);

            let successors = if i + 1 < mpo.sites.len() {
                vec![TemporalEdge {
                    dst: (i + 1) as u64,
                    weight: site.chi_right as f32,
                    begin_ts: 0,
                    end_ts: 1,
                }]
            } else {
                vec![]
            };

            GraphNode4D {
                id: i as u64,
                x: i as f32,
                y: 0.0,
                z: 0.0,
                begin_ts: 0,
                end_ts: 1,
                properties: props,
                successors,
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn identity_matrix(n: usize) -> Vec<f32> {
        let mut m = vec![0.0f32; n * n];
        for i in 0..n {
            m[i * n + i] = 1.0;
        }
        m
    }

    fn outer_product(u: &[f32], v: &[f32]) -> Vec<f32> {
        let mut m = vec![0.0f32; u.len() * v.len()];
        for (i, &ui) in u.iter().enumerate() {
            for (j, &vj) in v.iter().enumerate() {
                m[i * v.len() + j] = ui * vj;
            }
        }
        m
    }

    fn norm(v: &[f32]) -> f32 {
        v.iter().map(|x| x * x).sum::<f32>().sqrt()
    }

    #[test]
    fn test_svd_thin_orthogonality() {
        let a = [
            1.0f32, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 0.0, 0.0,
            0.0, 0.0, 0.0, 0.0, 4.0, 0.0, 0.0,
        ];
        let (u, _sigma, vt) = svd_thin(&a, 4, 6);

        // U^T U ≈ I  (u is 4×4, checking inner products of columns)
        for i in 0..4 {
            for j in 0..4 {
                let dot: f32 = (0..4).map(|k| u[k * 4 + i] * u[k * 4 + j]).sum();
                let expected = if i == j { 1.0 } else { 0.0 };
                assert!(
                    (dot - expected).abs() < 1e-3,
                    "U^T U [{i},{j}] = {dot:.6}, expected {expected}"
                );
            }
        }
        // Vt Vt^T ≈ I
        for i in 0..4 {
            for j in 0..4 {
                let dot: f32 = (0..6).map(|k| vt[i * 6 + k] * vt[j * 6 + k]).sum();
                let expected = if i == j { 1.0 } else { 0.0 };
                assert!(
                    (dot - expected).abs() < 1e-3,
                    "Vt Vt^T [{i},{j}] = {dot:.6}, expected {expected}"
                );
            }
        }
    }

    #[test]
    fn test_compress_identity_4x4_exact() {
        let w = identity_matrix(4);
        let mpo = compress_matrix_to_mpo(&w, 4, 4, 2, 4);
        let err = mpo_reconstruction_error(&mpo, &w, 4, 4);
        assert!(
            err < 1e-3,
            "identity matrix MPO reconstruction error {err:.6} should be < 1e-3"
        );
    }

    #[test]
    fn test_compress_rank1_chi1_exact() {
        let u = [1.0f32, 2.0, 2.0, 4.0];
        let v = [1.0f32, 1.0, 1.0, 1.0];
        let w = outer_product(&u, &v);
        let mpo = compress_matrix_to_mpo(&w, 4, 4, 2, 1);
        let err = mpo_reconstruction_error(&mpo, &w, 4, 4);
        assert!(
            err < 1e-3,
            "rank-1 MPO reconstruction error {err:.6} should be < 1e-3 with chi_max=1"
        );
    }

    #[test]
    fn test_mpo_apply_identity() {
        let w = identity_matrix(4);
        let mpo = compress_matrix_to_mpo(&w, 4, 4, 2, 4);
        let x = [1.0f32, 0.0, 0.0, 0.0];
        let y = mpo_apply(&mpo, &x);
        assert_eq!(y.len(), 4);
        assert!(
            (y[0] - 1.0).abs() < 1e-3 && y[1].abs() < 1e-3,
            "identity MPO on e_0 should return e_0, got {y:?}"
        );
    }

    #[test]
    fn test_mpo_compression_ratio() {
        let w: Vec<f32> = (0..256).map(|i| i as f32 * 0.01).collect();
        let mpo = compress_matrix_to_mpo(&w, 16, 16, 2, 1);
        let ratio = mpo_compression_ratio(&mpo, 16, 16);
        assert!(
            ratio < 0.5,
            "chi_max=1 on 16×16 should give ratio < 0.5, got {ratio:.4}"
        );
    }

    #[test]
    fn test_reconstruction_error_decreases_with_chi() {
        let w: Vec<f32> = (0..256).map(|i| ((i * 7 + 3) % 17) as f32 - 8.0).collect();
        let mpo1 = compress_matrix_to_mpo(&w, 16, 16, 2, 1);
        let mpo4 = compress_matrix_to_mpo(&w, 16, 16, 2, 4);
        let err1 = mpo_reconstruction_error(&mpo1, &w, 16, 16);
        let err4 = mpo_reconstruction_error(&mpo4, &w, 16, 16);
        assert!(
            err4 < err1,
            "chi=4 error {err4:.4} should be less than chi=1 error {err1:.4}"
        );
    }

    #[test]
    fn test_factor_dimension_powers_of_4() {
        let factors = factor_dimension(16);
        assert_eq!(factors, vec![4, 4], "16 = 4×2");
        let factors64 = factor_dimension(64);
        assert_eq!(factors64, vec![4, 4, 4], "64 = 4³");
    }

    #[test]
    fn test_factor_dimension_non_power() {
        let factors = factor_dimension(6);
        assert!(
            factors.iter().product::<usize>() == 6,
            "factors must multiply to 6, got {factors:?}"
        );
        assert!(factors.iter().all(|&d| d >= 2), "all factors >= 2");
    }

    #[test]
    fn test_factor_dimension_prime() {
        let factors = factor_dimension(7);
        assert_eq!(factors, vec![7], "prime number factors to itself");
    }

    #[test]
    fn test_compress_auto_6x6() {
        let w: Vec<f32> = (0..36).map(|i| i as f32 * 0.1 - 1.8).collect();
        let mpo = compress_matrix_to_mpo_auto(&w, 6, 6, 4);
        assert_eq!(mpo.sites.len(), 2);
        let err = mpo_reconstruction_error(&mpo, &w, 6, 6);
        assert!(
            err < 0.3,
            "6×6 auto MPO error {err:.4} should be < 0.3 with chi=4"
        );
    }

    #[test]
    fn test_compress_auto_12x8() {
        let w: Vec<f32> = (0..96)
            .map(|i| (((i * 7 + 3) % 13) as f32 - 6.0) * 0.5)
            .collect();
        let mpo = compress_matrix_to_mpo_auto(&w, 12, 8, 8);
        assert!(!mpo.sites.is_empty());
        let err = mpo_reconstruction_error(&mpo, &w, 12, 8);
        assert!(
            err < 0.3,
            "12×8 auto MPO error {err:.4} should be < 0.3 with chi=8"
        );
    }

    #[test]
    fn test_compress_auto_identity_non_power() {
        let w = identity_matrix(6);
        let mpo = compress_matrix_to_mpo_auto(&w, 6, 6, 4);
        let err = mpo_reconstruction_error(&mpo, &w, 6, 6);
        assert!(err < 0.1, "6×6 identity MPO error {err:.4} should be < 0.1");
    }

    #[test]
    fn test_mpo_to_graph_nodes_chain() {
        let w = identity_matrix(4);
        let mpo = compress_matrix_to_mpo(&w, 4, 4, 2, 4);
        let nodes = mpo_to_graph_nodes(&mpo);
        assert_eq!(nodes.len(), 2, "2-site MPO → 2 graph nodes");
        assert_eq!(nodes[0].successors.len(), 1, "site 0 links to site 1");
        assert_eq!(nodes[0].successors[0].dst, 1);
        assert_eq!(
            nodes[0].successors[0].weight as usize,
            mpo.sites[0].chi_right
        );
    }

    #[test]
    fn test_auto_mpo_apply_matches_dense() {
        let w: Vec<f32> = (0..36).map(|i| ((i * 7 + 3) % 11) as f32 - 5.0).collect();
        let mpo = compress_matrix_to_mpo_auto(&w, 6, 6, 8);
        let x: Vec<f32> = vec![1.0, -1.0, 0.5, -0.5, 2.0, -2.0];
        let y_mpo = mpo_apply(&mpo, &x);
        let y_dense: Vec<f32> = (0..6)
            .map(|i| (0..6).map(|j| w[i * 6 + j] * x[j]).sum())
            .collect();
        let rel_err = {
            let diff: f32 = y_mpo
                .iter()
                .zip(y_dense.iter())
                .map(|(a, b)| (a - b).powi(2))
                .sum();
            let norm: f32 = y_dense.iter().map(|v| v.powi(2)).sum();
            (diff / (norm + 1e-8)).sqrt()
        };
        assert!(
            rel_err < 0.3,
            "auto MPO apply rel_err {rel_err:.4} should be < 0.3"
        );
    }

    #[test]
    fn test_mpo_apply_random_close() {
        let w: Vec<f32> = (0..256).map(|i| ((i * 7 + 3) % 17) as f32 - 8.0).collect();

        // 1. With chi_max=16 (exact reconstruction because it is full rank)
        let mpo_exact = compress_matrix_to_mpo(&w, 16, 16, 2, 16);
        let x: Vec<f32> = (0..16).map(|i| (i as f32 + 1.0) / 16.0).collect();
        let y_dense: Vec<f32> = (0..16)
            .map(|i| (0..16).map(|j| w[i * 16 + j] * x[j]).sum())
            .collect();
        let y_mpo_exact = mpo_apply(&mpo_exact, &x);
        let diff_exact: Vec<f32> = y_dense
            .iter()
            .zip(y_mpo_exact.iter())
            .map(|(a, b)| a - b)
            .collect();
        let rel_err_exact = norm(&diff_exact) / (norm(&y_dense) + 1e-8);
        assert!(
            rel_err_exact < 1e-3,
            "full-rank MPO reconstruction error {rel_err_exact:.6} should be < 1e-3"
        );

        // 2. With chi_max=4 (truncated reconstruction, error is bounded by theoretical rank limits)
        let mpo_trunc = compress_matrix_to_mpo(&w, 16, 16, 2, 4);
        let y_mpo_trunc = mpo_apply(&mpo_trunc, &x);
        let diff_trunc: Vec<f32> = y_dense
            .iter()
            .zip(y_mpo_trunc.iter())
            .map(|(a, b)| a - b)
            .collect();
        let rel_err_trunc = norm(&diff_trunc) / (norm(&y_dense) + 1e-8);
        assert!(
            rel_err_trunc < 0.70,
            "truncated MPO reconstruction error {rel_err_trunc:.4} should be < 0.70"
        );
    }
}