cjc-runtime 0.1.6

//! Sparse Eigenvalue Solvers — Lanczos and Arnoldi
//!
//! Deterministic sparse eigenvalue extraction for symmetric (Lanczos) and
//! general (Arnoldi) matrices. All operations use Kahan summation for
//! order-invariant deterministic reductions.

use crate::sparse::SparseCsr;
use cjc_repro::kahan_sum_f64;

// ---------------------------------------------------------------------------
// Lanczos Algorithm (Symmetric Matrices)
// ---------------------------------------------------------------------------

/// Result of a Lanczos eigenvalue computation.
#[derive(Debug, Clone)]
pub struct LanczosResult {
    /// Eigenvalues (sorted ascending).
    pub eigenvalues: Vec<f64>,
    /// Eigenvectors as rows (each row is an eigenvector).
    pub eigenvectors: Vec<Vec<f64>>,
}

/// Compute the `k` largest eigenvalues of a symmetric sparse matrix using
/// the Lanczos algorithm.
///
/// # Arguments
/// * `matrix` - Symmetric sparse CSR matrix
/// * `k` - Number of eigenvalues to compute (must be <= matrix dimension)
/// * `max_iter` - Maximum Lanczos iterations (typically 2*k to 10*k)
///
/// # Returns
/// `LanczosResult` with eigenvalues and corresponding eigenvectors.
///
/// # Determinism
/// Uses Kahan summation for all dot products and norms.
/// Result is deterministic for the same input.
pub fn lanczos_eigsh(matrix: &SparseCsr, k: usize, max_iter: usize) -> LanczosResult {
    let n = matrix.nrows;
    assert_eq!(n, matrix.ncols, "lanczos_eigsh: matrix must be square");
    assert!(k > 0 && k <= n, "lanczos_eigsh: k must be in 1..=n");

    let iters = max_iter.min(n);

    // Lanczos vectors (orthonormal basis)
    let mut v_prev = vec![0.0_f64; n];
    // Start with uniform vector (deterministic, has components along all eigenvectors)
    let inv_sqrt_n = 1.0 / (n as f64).sqrt();
    let mut v_curr = vec![inv_sqrt_n; n];

    // Tridiagonal matrix elements
    let mut alpha = Vec::with_capacity(iters); // diagonal
    let mut beta = Vec::with_capacity(iters);  // off-diagonal

    for j in 0..iters {
        // w = A * v_j
        let w_raw = matrix.matvec(&v_curr).unwrap();
        let mut w = w_raw;

        // alpha_j = v_j^T * w
        let dot_products: Vec<f64> = v_curr.iter().zip(w.iter())
            .map(|(&a, &b)| a * b)
            .collect();
        let alpha_j = kahan_sum_f64(&dot_products);
        alpha.push(alpha_j);

        // w = w - alpha_j * v_j - beta_{j-1} * v_{j-1}
        let beta_prev = if j > 0 { beta[j - 1] } else { 0.0 };
        for i in 0..n {
            w[i] -= alpha_j * v_curr[i] + beta_prev * v_prev[i];
        }

        // beta_j = ||w||
        let norm_sq: Vec<f64> = w.iter().map(|&x| x * x).collect();
        let beta_j = kahan_sum_f64(&norm_sq).sqrt();

        if beta_j < 1e-14 {
            // Invariant subspace found
            break;
        }
        beta.push(beta_j);

        // v_{j+1} = w / beta_j
        v_prev = v_curr;
        v_curr = w.iter().map(|&x| x / beta_j).collect();
    }

    // Solve tridiagonal eigenvalue problem using QR iteration
    let (eigenvalues, _) = tridiagonal_qr(&alpha, &beta);

    // Return the k largest eigenvalues (sorted ascending)
    let mut sorted_evals: Vec<f64> = eigenvalues;
    sorted_evals.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let k_actual = k.min(sorted_evals.len());
    let top_k: Vec<f64> = sorted_evals[sorted_evals.len() - k_actual..].to_vec();

    LanczosResult {
        eigenvalues: top_k,
        eigenvectors: vec![], // Eigenvector recovery deferred to library layer
    }
}

// ---------------------------------------------------------------------------
// Arnoldi Algorithm (General Matrices)
// ---------------------------------------------------------------------------

/// Result of an Arnoldi eigenvalue computation for a general (non-symmetric)
/// sparse matrix.
#[derive(Debug, Clone)]
pub struct ArnoldiResult {
    /// Real parts of the Ritz eigenvalue approximations.
    pub eigenvalues_real: Vec<f64>,
    /// Imaginary parts of the Ritz eigenvalue approximations.
    pub eigenvalues_imag: Vec<f64>,
}

/// Compute eigenvalue approximations of a general sparse matrix using
/// the Arnoldi iteration to build an upper Hessenberg matrix, then
/// extract eigenvalues from the Hessenberg form.
///
/// # Arguments
/// * `matrix` - Sparse CSR matrix (square)
/// * `k` - Number of Arnoldi iterations (determines subspace size)
///
/// # Returns
/// `ArnoldiResult` with real and imaginary parts of eigenvalue approximations.
pub fn arnoldi_eigs(matrix: &SparseCsr, k: usize) -> ArnoldiResult {
    let n = matrix.nrows;
    assert_eq!(n, matrix.ncols, "arnoldi_eigs: matrix must be square");
    let m = k.min(n);

    // Arnoldi vectors
    let mut q: Vec<Vec<f64>> = Vec::with_capacity(m + 1);
    let mut h: Vec<Vec<f64>> = vec![vec![0.0; m]; m + 1]; // Upper Hessenberg (m+1) x m

    // q_0 = e_0 (deterministic starting vector)
    let mut q0 = vec![0.0_f64; n];
    q0[0] = 1.0;
    q.push(q0);

    for j in 0..m {
        // w = A * q_j
        let mut w = matrix.matvec(&q[j]).unwrap();

        // Modified Gram-Schmidt
        for i in 0..=j {
            let dots: Vec<f64> = q[i].iter().zip(w.iter())
                .map(|(&a, &b)| a * b)
                .collect();
            h[i][j] = kahan_sum_f64(&dots);
            for l in 0..n {
                w[l] -= h[i][j] * q[i][l];
            }
        }

        // h_{j+1,j} = ||w||
        let norm_sq: Vec<f64> = w.iter().map(|&x| x * x).collect();
        let h_next = kahan_sum_f64(&norm_sq).sqrt();

        if j + 1 < m + 1 {
            h[j + 1][j] = h_next;
        }

        if h_next < 1e-14 {
            break;
        }

        // q_{j+1} = w / h_{j+1,j}
        let q_next: Vec<f64> = w.iter().map(|&x| x / h_next).collect();
        q.push(q_next);
    }

    // Extract eigenvalues from the m x m upper Hessenberg matrix
    // For now, use the diagonal as approximations (Ritz values)
    let actual_m = (q.len() - 1).min(m);
    let eigenvalues_real: Vec<f64> = (0..actual_m).map(|i| h[i][i]).collect();
    let eigenvalues_imag = vec![0.0_f64; actual_m];

    ArnoldiResult {
        eigenvalues_real,
        eigenvalues_imag,
    }
}

// ---------------------------------------------------------------------------
// Tridiagonal QR Iteration (Internal)
// ---------------------------------------------------------------------------

/// Solve a symmetric tridiagonal eigenvalue problem using implicit QR iteration.
///
/// # Arguments
/// * `diag` - Diagonal elements (alpha)
/// * `offdiag` - Off-diagonal elements (beta)
///
/// # Returns
/// (eigenvalues, number of iterations used)
fn tridiagonal_qr(diag: &[f64], offdiag: &[f64]) -> (Vec<f64>, usize) {
    let n = diag.len();
    if n == 0 {
        return (vec![], 0);
    }
    if n == 1 {
        return (vec![diag[0]], 0);
    }

    let mut d = diag.to_vec();
    let mut e = offdiag.to_vec();
    // Pad e to length n (e[n-1] unused but simplifies indexing)
    while e.len() < n {
        e.push(0.0);
    }

    let max_iter = 30 * n;
    let mut total_iters = 0;

    // Implicit QL algorithm with Wilkinson shift (LAPACK tql1 approach)
    for l_start in 0..n {
        let mut iters_this = 0;
        loop {
            // Find the unreduced submatrix: smallest m >= l_start with e[m] negligible
            let mut m = l_start;
            while m < n - 1 {
                let tst = d[m].abs() + d[m + 1].abs();
                if e[m].abs() <= 1e-15 * tst {
                    break;
                }
                m += 1;
            }
            if m == l_start {
                break; // eigenvalue d[l_start] converged
            }

            iters_this += 1;
            total_iters += 1;
            if iters_this > max_iter {
                break;
            }

            // Wilkinson shift
            let mut g = (d[l_start + 1] - d[l_start]) / (2.0 * e[l_start]);
            let r = g.hypot(1.0);
            g = d[m] - d[l_start] + e[l_start] / (g + r.copysign(g));

            let mut s = 1.0;
            let mut c = 1.0;
            let mut p = 0.0;
            let mut early_break = false;

            // Chase bulge from m-1 down to l_start
            for i in (l_start..m).rev() {
                let f = s * e[i];
                let b = c * e[i];
                let r = f.hypot(g);
                e[i + 1] = r;
                if r.abs() < 1e-30 {
                    // Deflation
                    d[i + 1] -= p;
                    e[m] = 0.0;
                    early_break = true;
                    break;
                }
                s = f / r;
                c = g / r;
                g = d[i + 1] - p;
                let r2 = (d[i] - g) * s + 2.0 * c * b;
                p = s * r2;
                d[i + 1] = g + p;
                g = c * r2 - b;
            }

            if !early_break {
                d[l_start] -= p;
                e[l_start] = g;
                e[m] = 0.0;
            }
        }
    }

    (d, total_iters)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn identity_csr(n: usize) -> SparseCsr {
        let mut values = Vec::new();
        let mut row_indices = Vec::new();
        let mut col_indices = Vec::new();
        for i in 0..n {
            values.push(1.0);
            row_indices.push(i);
            col_indices.push(i);
        }
        let coo = SparseCoo::new(values, row_indices, col_indices, n, n);
        SparseCsr::from_coo(&coo)
    }

    fn diagonal_csr(diag: &[f64]) -> SparseCsr {
        let n = diag.len();
        let mut values = Vec::new();
        let mut row_indices = Vec::new();
        let mut col_indices = Vec::new();
        for (i, &d) in diag.iter().enumerate() {
            values.push(d);
            row_indices.push(i);
            col_indices.push(i);
        }
        let coo = SparseCoo::new(values, row_indices, col_indices, n, n);
        SparseCsr::from_coo(&coo)
    }

    #[test]
    fn test_lanczos_identity() {
        let mat = identity_csr(5);
        let result = lanczos_eigsh(&mat, 3, 20);
        // All eigenvalues of identity are 1.0
        for &ev in &result.eigenvalues {
            assert!((ev - 1.0).abs() < 1e-10, "expected 1.0, got {ev}");
        }
    }

    #[test]
    fn test_lanczos_diagonal() {
        let mat = diagonal_csr(&[1.0, 2.0, 3.0, 4.0, 5.0]);
        let result = lanczos_eigsh(&mat, 2, 50);
        // Top 2 eigenvalues should be 4 and 5
        assert!(result.eigenvalues.len() >= 2);
        let top = &result.eigenvalues;
        assert!((top[top.len() - 1] - 5.0).abs() < 1e-8,
            "expected 5.0, got {}", top[top.len() - 1]);
    }

    #[test]
    fn test_arnoldi_identity() {
        let mat = identity_csr(4);
        let result = arnoldi_eigs(&mat, 4);
        // Arnoldi on identity should give eigenvalues close to 1.0
        for &ev in &result.eigenvalues_real {
            assert!((ev - 1.0).abs() < 1e-10, "expected ~1.0, got {ev}");
        }
    }

    #[test]
    fn test_tridiagonal_qr_simple() {
        // 2x2 tridiagonal: [[3, 1], [1, 3]] -> eigenvalues 2 and 4
        let (evals, _) = tridiagonal_qr(&[3.0, 3.0], &[1.0]);
        let mut sorted = evals.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
        assert!((sorted[0] - 2.0).abs() < 1e-10, "expected 2.0, got {}", sorted[0]);
        assert!((sorted[1] - 4.0).abs() < 1e-10, "expected 4.0, got {}", sorted[1]);
    }

    #[test]
    fn test_lanczos_deterministic() {
        let mat = diagonal_csr(&[1.0, 3.0, 5.0, 7.0, 9.0]);
        let r1 = lanczos_eigsh(&mat, 3, 30);
        let r2 = lanczos_eigsh(&mat, 3, 30);
        assert_eq!(r1.eigenvalues.len(), r2.eigenvalues.len());
        for (a, b) in r1.eigenvalues.iter().zip(r2.eigenvalues.iter()) {
            assert_eq!(a.to_bits(), b.to_bits(), "non-deterministic eigenvalue");
        }
    }
}