geographdb-core 0.5.4

use super::mpo::{compress_matrix_to_mpo, mpo_apply, svd_thin, Mpo};

/// MPO-compressed rotation matrix for KV cache (OSCAR-style).
/// OSCAR derives per-layer rotation V from empirical query covariance.
/// For large d, V is expensive to store (d²×4 bytes per layer).
/// This module experiments with MPO compression of V.
#[derive(Debug, Clone)]
pub struct RotationMpo {
    pub mpo: Mpo,
    pub d: usize,
}

fn mat_vec(a: &[f32], x: &[f32], d: usize) -> Vec<f32> {
    (0..d)
        .map(|i| (0..d).map(|j| a[i * d + j] * x[j]).sum())
        .collect()
}

fn mat_t_vec(a: &[f32], x: &[f32], d: usize) -> Vec<f32> {
    (0..d)
        .map(|i| (0..d).map(|j| a[j * d + i] * x[j]).sum())
        .collect()
}

/// Gram-Schmidt orthogonalization of rows of a (d×d row-major).
/// Handles rank-deficient input by substituting standard basis vectors for degenerate rows.
fn gram_schmidt(a: &[f32], d: usize) -> Vec<f32> {
    let mut out = vec![0.0f32; d * d];
    for i in 0..d {
        let mut v: Vec<f32> = (0..d).map(|j| a[i * d + j]).collect();
        let initial_norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        for k in 0..i {
            let dot: f32 = (0..d).map(|j| v[j] * out[k * d + j]).sum();
            for j in 0..d {
                v[j] -= dot * out[k * d + j];
            }
        }
        let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        // Relative threshold: f32 accumulation leaves residuals ~0.1% of original
        // for linearly dependent rows (rank-deficient SVD output).
        let is_degenerate = norm < 0.01 * initial_norm.max(1e-8);
        if !is_degenerate {
            for j in 0..d {
                out[i * d + j] = v[j] / norm;
            }
        } else {
            // Row is in span of previous rows; find a basis vector for the orthogonal complement
            'find: for e in 0..d {
                let mut v2 = vec![0.0f32; d];
                v2[e] = 1.0;
                for k in 0..i {
                    let dot: f32 = (0..d).map(|j| v2[j] * out[k * d + j]).sum();
                    for j in 0..d {
                        v2[j] -= dot * out[k * d + j];
                    }
                }
                let n2: f32 = v2.iter().map(|x| x * x).sum::<f32>().sqrt();
                if n2 > 0.5 {
                    for j in 0..d {
                        out[i * d + j] = v2[j] / n2;
                    }
                    break 'find;
                }
            }
        }
    }
    out
}

fn l2_error(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b)
        .map(|(x, y)| (x - y).powi(2))
        .sum::<f32>()
        .sqrt()
}

/// Empirical query covariance C = (1/N) Σ q_i q_i^T.
/// Returns d×d matrix, row-major.
pub fn empirical_query_covariance(queries: &[Vec<f32>], d: usize) -> Vec<f32> {
    let mut cov = vec![0.0f32; d * d];
    for q in queries {
        assert_eq!(q.len(), d);
        for i in 0..d {
            for j in 0..d {
                cov[i * d + j] += q[i] * q[j];
            }
        }
    }
    let n = queries.len() as f32;
    cov.iter_mut().for_each(|c| *c /= n);
    cov
}

/// Rotation matrix V from covariance via SVD (= eigendecomposition for PSD).
/// Returns d×d orthogonal matrix, row-major.
pub fn covariance_to_rotation(cov: &[f32], d: usize) -> Vec<f32> {
    let (u, _s, _vt) = svd_thin(cov, d, d);
    // Gram-Schmidt ensures exact orthogonality when cov is rank-deficient
    // (Jacobi SVD leaves null-space vectors unnormalized for zero eigenvalues)
    gram_schmidt(&u, d)
}

/// Compress rotation V (d×d) to MPO. d must be a perfect n_sites-th power.
pub fn compress_rotation(v: &[f32], d: usize, n_sites: usize, chi_max: usize) -> RotationMpo {
    RotationMpo {
        mpo: compress_matrix_to_mpo(v, d, d, n_sites, chi_max),
        d,
    }
}

/// Apply MPO-compressed rotation to a key vector.
pub fn apply_rotation(rot: &RotationMpo, key: &[f32]) -> Vec<f32> {
    mpo_apply(&rot.mpo, key)
}

/// Mean cosine similarity between exact V*k and MPO(k) over test vectors.
pub fn rotation_fidelity(v_exact: &[f32], rot: &RotationMpo, test_vecs: &[Vec<f32>]) -> f32 {
    let d = rot.d;
    let total: f32 = test_vecs
        .iter()
        .map(|k| {
            let exact = mat_vec(v_exact, k, d);
            let approx = apply_rotation(rot, k);
            let dot: f32 = exact.iter().zip(&approx).map(|(a, b)| a * b).sum();
            let ne: f32 = exact.iter().map(|x| x * x).sum::<f32>().sqrt();
            let na: f32 = approx.iter().map(|x| x * x).sum::<f32>().sqrt();
            if ne > 1e-9 && na > 1e-9 {
                dot / (ne * na)
            } else {
                1.0
            }
        })
        .sum();
    total / test_vecs.len() as f32
}

/// Ratio of MPO parameter count to dense d×d matrix count.
/// < 1.0 means the MPO is cheaper to store.
pub fn mpo_compression_ratio_rotation(rot: &RotationMpo) -> f32 {
    let original = (rot.d * rot.d) as f32;
    let compressed: usize = rot
        .mpo
        .sites
        .iter()
        .map(|s| s.chi_left * s.d_out * s.d_in * s.chi_right)
        .sum();
    compressed as f32 / original
}

/// INT2 quantization: 4 levels.
/// Returns (levels in 0..=3, scale, zero) where x_hat[i] = levels[i] * scale + zero.
pub fn quantize_int2(x: &[f32]) -> (Vec<u8>, f32, f32) {
    let min = x.iter().cloned().fold(f32::INFINITY, f32::min);
    let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let scale = if (max - min).abs() < 1e-9 {
        1.0
    } else {
        (max - min) / 3.0
    };
    let zero = min;
    let q = x
        .iter()
        .map(|&v| ((v - zero) / scale).round().clamp(0.0, 3.0) as u8)
        .collect();
    (q, scale, zero)
}

/// Inverse of quantize_int2.
pub fn dequantize_int2(q: &[u8], scale: f32, zero: f32) -> Vec<f32> {
    q.iter().map(|&v| v as f32 * scale + zero).collect()
}

/// L2 error after: exact V*key → INT2 → dequant → V^T (un-rotate).
pub fn rotation_quantization_error(v: &[f32], key: &[f32], d: usize) -> f32 {
    let rotated = mat_vec(v, key, d);
    let (q, scale, zero) = quantize_int2(&rotated);
    let dequantized = dequantize_int2(&q, scale, zero);
    let recovered = mat_t_vec(v, &dequantized, d);
    l2_error(key, &recovered)
}

/// L2 error after naive INT2 → dequant (no rotation).
pub fn identity_quantization_error(key: &[f32]) -> f32 {
    let (q, scale, zero) = quantize_int2(key);
    let dequantized = dequantize_int2(&q, scale, zero);
    l2_error(key, &dequantized)
}

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

    fn basis_vecs(d: usize) -> Vec<Vec<f32>> {
        (0..d)
            .map(|i| {
                let mut e = vec![0.0f32; d];
                e[i] = 1.0;
                e
            })
            .collect()
    }

    // 1. All-same query → rank-1 covariance
    #[test]
    fn test_empirical_cov_rank1() {
        let d = 4;
        let norm = (2.0f32).sqrt();
        let v_norm = vec![1.0 / norm, 0.0, 1.0 / norm, 0.0_f32];
        let queries: Vec<Vec<f32>> = (0..10).map(|_| v_norm.clone()).collect();
        let cov = empirical_query_covariance(&queries, d);
        for i in 0..d {
            for j in 0..d {
                let expected = v_norm[i] * v_norm[j];
                assert!(
                    (cov[i * d + j] - expected).abs() < 1e-5,
                    "cov[{i},{j}]={}, expected {expected}",
                    cov[i * d + j]
                );
            }
        }
    }

    // 2. Standard basis queries → covariance = I/d
    #[test]
    fn test_empirical_cov_basis() {
        let d = 4;
        let queries = basis_vecs(d);
        let cov = empirical_query_covariance(&queries, d);
        let inv_d = 1.0 / d as f32;
        for i in 0..d {
            for j in 0..d {
                let expected = if i == j { inv_d } else { 0.0 };
                assert!(
                    (cov[i * d + j] - expected).abs() < 1e-5,
                    "cov[{i},{j}]={}, expected {expected}",
                    cov[i * d + j]
                );
            }
        }
    }

    // 3. Rotation is orthogonal: V * V^T ≈ I
    #[test]
    fn test_rotation_orthogonal() {
        let d = 4;
        let queries = vec![
            vec![1.0f32, 1.0, 0.0, 0.0],
            vec![0.0f32, 1.0, 1.0, 0.0],
            vec![1.0f32, 0.0, 0.0, 1.0],
        ];
        let cov = empirical_query_covariance(&queries, d);
        let v = covariance_to_rotation(&cov, d);
        for i in 0..d {
            for j in 0..d {
                let dot: f32 = (0..d).map(|k| v[i * d + k] * v[j * d + k]).sum();
                let expected = if i == j { 1.0 } else { 0.0 };
                assert!(
                    (dot - expected).abs() < 1e-4,
                    "V*V^T[{i},{j}]={dot:.5}, expected {expected}"
                );
            }
        }
    }

    // 4. Rotation preserves L2 norm
    #[test]
    fn test_rotation_preserves_norm() {
        let d = 4;
        let queries = vec![vec![1.0f32, 2.0, 0.5, -1.0]];
        let cov = empirical_query_covariance(&queries, d);
        let v = covariance_to_rotation(&cov, d);
        let key = vec![3.0f32, -1.0, 2.0, 0.5];
        let rotated = mat_vec(&v, &key, d);
        let norm_k: f32 = key.iter().map(|x| x * x).sum::<f32>().sqrt();
        let norm_r: f32 = rotated.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            (norm_k - norm_r).abs() < 1e-4,
            "|k|={norm_k:.4}, |Vk|={norm_r:.4}"
        );
    }

    // 5. MPO at full chi reproduces exact rotation (fidelity ≈ 1)
    #[test]
    fn test_compress_full_chi_exact() {
        let d = 16;
        let n_sites = 2;
        let chi_max = 16;
        // Identity rotation — simplest nontrivial case
        let v: Vec<f32> = (0..d * d)
            .map(|idx| if idx / d == idx % d { 1.0 } else { 0.0 })
            .collect();
        let rot = compress_rotation(&v, d, n_sites, chi_max);
        let test_vecs = basis_vecs(d);
        let fidelity = rotation_fidelity(&v, &rot, &test_vecs);
        assert!(fidelity > 0.999, "fidelity={fidelity:.5}");
    }

    // 6. Fidelity is monotone in chi_max
    #[test]
    fn test_fidelity_increases_with_chi() {
        let d = 16;
        let n_sites = 2;
        // Non-trivial rotation from mixed queries
        let queries: Vec<Vec<f32>> = (0..8)
            .map(|i| {
                let mut q = vec![0.0f32; d];
                q[i] = 1.0;
                q[i + 8] = 0.5;
                q
            })
            .collect();
        let cov = empirical_query_covariance(&queries, d);
        let v = covariance_to_rotation(&cov, d);
        let test_vecs = basis_vecs(d);
        let fid1 = rotation_fidelity(&v, &compress_rotation(&v, d, n_sites, 1), &test_vecs);
        let fid4 = rotation_fidelity(&v, &compress_rotation(&v, d, n_sites, 4), &test_vecs);
        let fid16 = rotation_fidelity(&v, &compress_rotation(&v, d, n_sites, 16), &test_vecs);
        assert!(
            fid4 >= fid1 - 1e-4,
            "fid4={fid4:.4} should >= fid1={fid1:.4}"
        );
        assert!(
            fid16 >= fid4 - 1e-4,
            "fid16={fid16:.4} should >= fid4={fid4:.4}"
        );
        assert!(fid16 > 0.999, "fid16={fid16:.4} should be near 1.0");
    }

    // 7. chi_max=2 compresses d=16 matrix below original size
    #[test]
    fn test_compression_ratio_small_chi() {
        let d = 16;
        let n_sites = 2;
        let v: Vec<f32> = (0..d * d)
            .map(|idx| if idx / d == idx % d { 1.0 } else { 0.0 })
            .collect();
        let rot2 = compress_rotation(&v, d, n_sites, 2);
        let rot4 = compress_rotation(&v, d, n_sites, 4);
        let r2 = mpo_compression_ratio_rotation(&rot2);
        let r4 = mpo_compression_ratio_rotation(&rot4);
        assert!(r2 < r4, "chi=2 ratio={r2:.3} should < chi=4 ratio={r4:.3}");
        assert!(r2 < 1.0, "chi=2 should compress d=16: ratio={r2:.3}");
    }

    // 8. INT2 levels are always in 0..=3
    #[test]
    fn test_quantize_int2_levels_in_range() {
        let x = vec![-5.0f32, -3.0, 0.0, 1.0, 2.5, 7.0, 100.0, -100.0];
        let (q, _scale, _zero) = quantize_int2(&x);
        for &level in &q {
            assert!(level <= 3, "level {level} out of range");
        }
    }

    // 9. INT2 dequantize error bounded by scale/2
    #[test]
    fn test_quantize_int2_error_bounded() {
        let x = vec![-3.7f32, -1.2, 0.3, 1.8, 3.1, -0.5, 2.2, -2.9];
        let (q, scale, zero) = quantize_int2(&x);
        let deq = dequantize_int2(&q, scale, zero);
        for (orig, rec) in x.iter().zip(deq.iter()) {
            assert!(
                (orig - rec).abs() <= scale / 2.0 + 1e-4,
                "error {:.4} > scale/2 {:.4}",
                (orig - rec).abs(),
                scale / 2.0
            );
        }
    }

    // 10. At full chi, MPO rotation output matches exact rotation
    #[test]
    fn test_mpo_rotation_matches_exact() {
        let d = 16;
        let n_sites = 2;
        let queries: Vec<Vec<f32>> = (0..4)
            .map(|i| {
                let mut q = vec![0.0f32; d];
                q[i * 4] = 1.0;
                q[i * 4 + 1] = 0.7;
                q
            })
            .collect();
        let cov = empirical_query_covariance(&queries, d);
        let v = covariance_to_rotation(&cov, d);
        let rot = compress_rotation(&v, d, n_sites, 16);
        let key = vec![
            7.0, 3.0, 1.5, 1.0, 0.8, 0.6, 0.4, 0.2, 0.2, 0.4, 0.6, 0.8, 1.0, 1.5, 3.0, 7.0_f32,
        ];
        let exact = mat_vec(&v, &key, d);
        let approx = apply_rotation(&rot, &key);
        let diff = l2_error(&exact, &approx);
        assert!(diff < 0.01, "MPO vs exact rotation L2 diff: {diff:.5}");
    }
}