#![cfg(feature = "stde")]

use approx::assert_relative_eq;
use echidna::{BReverse, BytecodeTape, Scalar};

fn record_fn(f: impl FnOnce(&[BReverse<f64>]) -> BReverse<f64>, x: &[f64]) -> BytecodeTape<f64> {
    let (tape, _) = echidna::record(f, x);
    tape
}

// ══════════════════════════════════════════════
//  Test functions
// ══════════════════════════════════════════════

/// f(x,y) = x^2 + y^2
fn sum_of_squares<T: Scalar>(x: &[T]) -> T {
    x[0] * x[0] + x[1] * x[1]
}

/// f(x,y,z) = x^2*y + y^3
fn cubic_mix<T: Scalar>(x: &[T]) -> T {
    x[0] * x[0] * x[1] + x[1] * x[1] * x[1]
}

/// f(x,y) = x^4 + y^4
#[cfg(feature = "diffop")]
fn quartic<T: Scalar>(x: &[T]) -> T {
    let x0 = x[0];
    let y0 = x[1];
    x0 * x0 * x0 * x0 + y0 * y0 * y0 * y0
}

// ══════════════════════════════════════════════
//  20. Dense STDE for positive-definite operators
// ══════════════════════════════════════════════

#[test]
fn dense_stde_identity_is_laplacian() {
    // L=I, dense_stde_2nd should equal laplacian (same z_vectors as directions)
    let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
    let x = [1.0, 2.0];

    // Identity Cholesky factor
    let row0: Vec<f64> = vec![1.0, 0.0];
    let row1: Vec<f64> = vec![0.0, 1.0];
    let cholesky: Vec<&[f64]> = vec![&row0, &row1];

    // Use Rademacher vectors as z
    let z0: Vec<f64> = vec![1.0, 1.0];
    let z1: Vec<f64> = vec![1.0, -1.0];
    let z2: Vec<f64> = vec![-1.0, 1.0];
    let z3: Vec<f64> = vec![-1.0, -1.0];
    let z_vecs: Vec<&[f64]> = vec![&z0, &z1, &z2, &z3];

    let dense_result = echidna::stde::dense_stde_2nd(&tape, &x, &cholesky, &z_vecs);

    // With L=I, v=z, so dense_stde_2nd is the same as laplacian
    let (_, lap) = echidna::stde::laplacian(&tape, &x, &z_vecs);

    assert_relative_eq!(dense_result.estimate, lap, epsilon = 1e-10);
}

#[test]
fn dense_stde_diagonal_scaling() {
    // L=diag(a), C=diag(a²): tr(C·H) = Σ a_j² ∂²u/∂x_j²
    // f(x,y) = x²+y², H=diag(2,2), a=(2,3)
    // tr(C·H) = 4*2 + 9*2 = 26
    let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
    let x = [1.0, 2.0];

    let row0: Vec<f64> = vec![2.0, 0.0];
    let row1: Vec<f64> = vec![0.0, 3.0];
    let cholesky: Vec<&[f64]> = vec![&row0, &row1];

    // All 4 Rademacher z-vectors for exact result
    let z0: Vec<f64> = vec![1.0, 1.0];
    let z1: Vec<f64> = vec![1.0, -1.0];
    let z2: Vec<f64> = vec![-1.0, 1.0];
    let z3: Vec<f64> = vec![-1.0, -1.0];
    let z_vecs: Vec<&[f64]> = vec![&z0, &z1, &z2, &z3];

    let result = echidna::stde::dense_stde_2nd(&tape, &x, &cholesky, &z_vecs);
    assert_relative_eq!(result.estimate, 26.0, epsilon = 1e-10);
}

#[test]
fn dense_stde_off_diagonal() {
    // L with off-diagonal entries, verify against exact tr(C·H) from full Hessian
    // f(x,y,z) = x²y + y³ at (1,2,3)
    // H = [[4, 2, 0], [2, 12, 0], [0, 0, 0]]
    // L = [[1, 0, 0], [0.5, 1, 0], [0, 0, 1]] (lower triangular)
    // C = L·L^T = [[1, 0.5, 0], [0.5, 1.25, 0], [0, 0, 1]]
    // tr(C·H) = C[0][0]*H[0][0] + C[0][1]*H[1][0] + C[1][0]*H[0][1] + C[1][1]*H[1][1]
    //         + C[0][2]*H[2][0] + ... (all zero)
    //         = 1*4 + 0.5*2 + 0.5*2 + 1.25*12 + 0 + 0 + 0 + 0 + 1*0
    //         = 4 + 1 + 1 + 15 = 21
    let tape = record_fn(cubic_mix, &[1.0, 2.0, 3.0]);
    let x = [1.0, 2.0, 3.0];

    let row0: Vec<f64> = vec![1.0, 0.0, 0.0];
    let row1: Vec<f64> = vec![0.5, 1.0, 0.0];
    let row2: Vec<f64> = vec![0.0, 0.0, 1.0];
    let cholesky: Vec<&[f64]> = vec![&row0, &row1, &row2];

    // All 8 Rademacher z-vectors for exact result
    let signs: [f64; 2] = [1.0, -1.0];
    let mut vecs = Vec::new();
    for &s0 in &signs {
        for &s1 in &signs {
            for &s2 in &signs {
                vecs.push(vec![s0, s1, s2]);
            }
        }
    }
    let z_vecs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();

    let result = echidna::stde::dense_stde_2nd(&tape, &x, &cholesky, &z_vecs);
    assert_relative_eq!(result.estimate, 21.0, epsilon = 1e-8);
}

#[test]
fn dense_stde_matches_parabolic() {
    // L=σ, dense_stde_2nd matches 2*parabolic_diffusion (½ tr(σσ^T H) = ½ * dense_stde_2nd)
    // σ = diag(2, 3) as column vectors
    // parabolic_diffusion computes ½ tr(σσ^T H)
    // dense_stde_2nd computes tr(σσ^T H) = tr(C·H)
    let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
    let x = [1.0, 2.0];

    // σ columns (for parabolic_diffusion)
    let c0: Vec<f64> = vec![2.0, 0.0];
    let c1: Vec<f64> = vec![0.0, 3.0];
    let cols: Vec<&[f64]> = vec![&c0, &c1];
    let (_, diffusion) = echidna::stde::parabolic_diffusion(&tape, &x, &cols);

    // Same σ as Cholesky rows (for dense_stde_2nd)
    let row0: Vec<f64> = vec![2.0, 0.0];
    let row1: Vec<f64> = vec![0.0, 3.0];
    let cholesky: Vec<&[f64]> = vec![&row0, &row1];

    let z0: Vec<f64> = vec![1.0, 1.0];
    let z1: Vec<f64> = vec![1.0, -1.0];
    let z2: Vec<f64> = vec![-1.0, 1.0];
    let z3: Vec<f64> = vec![-1.0, -1.0];
    let z_vecs: Vec<&[f64]> = vec![&z0, &z1, &z2, &z3];

    let dense_result = echidna::stde::dense_stde_2nd(&tape, &x, &cholesky, &z_vecs);

    // parabolic_diffusion = ½ tr(C·H), dense_stde_2nd = tr(C·H)
    assert_relative_eq!(dense_result.estimate, 2.0 * diffusion, epsilon = 1e-10);
}

#[test]
fn dense_stde_stats_populated() {
    let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
    let x = [1.0, 2.0];

    let row0: Vec<f64> = vec![1.0, 0.0];
    let row1: Vec<f64> = vec![0.0, 1.0];
    let cholesky: Vec<&[f64]> = vec![&row0, &row1];

    let z0: Vec<f64> = vec![1.0, 1.0];
    let z1: Vec<f64> = vec![1.0, -1.0];
    let z_vecs: Vec<&[f64]> = vec![&z0, &z1];

    let result = echidna::stde::dense_stde_2nd(&tape, &x, &cholesky, &z_vecs);
    assert_eq!(result.num_samples, 2);
    assert_relative_eq!(result.value, 5.0, epsilon = 1e-10);
    // With diagonal H and identity Cholesky, variance is zero
    assert_relative_eq!(result.sample_variance, 0.0, epsilon = 1e-10);
}

// ══════════════════════════════════════════════
//  21. Sparse STDE (requires stde + diffop)
// ══════════════════════════════════════════════

#[cfg(feature = "diffop")]
mod sparse_stde_tests {
    use super::*;
    use echidna::diffop::DiffOp;

    #[test]
    fn stde_sparse_full_sample_matches_exact() {
        // Full sample (all entries) should match DiffOp::eval exactly
        // Use Laplacian on sum_of_squares: exact = 4
        let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
        let x = [1.0, 2.0];

        let op = DiffOp::<f64>::laplacian(2);
        let (_, exact) = op.eval(&tape, &x);

        let dist = op.sparse_distribution();
        let all_indices: Vec<usize> = (0..dist.len()).collect();
        let result = echidna::stde::stde_sparse(&tape, &x, &dist, &all_indices);

        assert_relative_eq!(result.estimate, exact, epsilon = 1e-6);
    }

    #[test]
    fn stde_sparse_laplacian_convergence() {
        // 1000 deterministic samples: mean should be within tolerance of exact
        let tape = record_fn(cubic_mix, &[1.0, 2.0, 3.0]);
        let x = [1.0, 2.0, 3.0];

        let op = DiffOp::<f64>::laplacian(3);
        let (_, exact) = op.eval(&tape, &x); // 16.0

        let dist = op.sparse_distribution();

        // Generate deterministic "random" indices via simple hash
        let num_samples = 1000;
        let indices: Vec<usize> = (0..num_samples)
            .map(|i| {
                let u = ((i as u64 * 2654435761u64) % 1000) as f64 / 1000.0;
                dist.sample_index(u)
            })
            .collect();

        let result = echidna::stde::stde_sparse(&tape, &x, &dist, &indices);

        // Mean should be close to exact (within 3 standard errors)
        let error = (result.estimate - exact).abs();
        let bound = 3.0 * result.standard_error;
        assert!(
            error < bound || error < 1.0,
            "stde_sparse estimate {} too far from exact {}: error = {}, 3*SE = {}",
            result.estimate,
            exact,
            error,
            bound,
        );
    }

    #[test]
    fn stde_sparse_diagonal_4th() {
        // Diagonal 4th-order operator on quartic: Σ_j ∂⁴/∂x_j⁴ = 24 + 24 = 48
        // Uses diagonal(n, 4) for uniform-weight stochastic test (biharmonic has
        // non-uniform weights which require importance sampling, not full enumeration).
        let tape = record_fn(quartic, &[2.0, 3.0]);
        let x = [2.0, 3.0];

        let op = DiffOp::<f64>::diagonal(2, 4);
        let (_, exact) = op.eval(&tape, &x);
        assert_relative_eq!(exact, 48.0, epsilon = 1e-4);

        let dist = op.sparse_distribution();
        let all_indices: Vec<usize> = (0..dist.len()).collect();
        let result = echidna::stde::stde_sparse(&tape, &x, &dist, &all_indices);

        assert_relative_eq!(result.estimate, 48.0, epsilon = 1e-4);
    }

    /// f(x,y) = exp(x) + sin(y): pure second partials genuinely differ
    /// (∂²/∂x² = exp(x) vs ∂²/∂y² = -sin(y)), so a mixed-coefficient
    /// operator cannot be estimated without importance sampling.
    fn exp_plus_sin<T: Scalar>(x: &[T]) -> T {
        x[0].exp() + x[1].sin()
    }

    #[test]
    fn stde_sparse_mixed_second_order() {
        // Mixed-coefficient operator L = ∂²/∂x² + 2∂²/∂y² on exp(x)+sin(y)
        // at (0.5, 1.0). Crucially the two pure second partials DIFFER
        // (∂²/∂x² = exp(0.5) ≈ 1.649, ∂²/∂y² = -sin(1) ≈ -0.841), so this is
        // a genuinely non-vacuous fixture: unlike a symmetric operand where
        // the partials coincide, the |C_α|-weighted true value cannot be
        // recovered by a uniform average over the sparse entries. `stde_sparse`
        // is unbiased only when indices are drawn ∝|C_α| (importance sampling).
        let x = [0.5, 1.0];
        let tape = record_fn(exp_plus_sin, &x);

        // True value: L(exp(x)+sin(y)) = exp(x) + 2·(-sin(y)).
        let expected = 0.5_f64.exp() - 2.0 * 1.0_f64.sin();

        let op = DiffOp::from_orders(
            2,
            &[
                (1.0, &[2, 0]), // ∂²/∂x²
                (2.0, &[0, 2]), // 2∂²/∂y²
            ],
        );
        let (_, exact) = op.eval(&tape, &x);
        assert_relative_eq!(exact, expected, epsilon = 1e-6);

        let dist = op.sparse_distribution();

        // Non-vacuity guard: a uniform enumeration of the sparse entries is
        // biased here (unequal |C_α| = 1 vs 2 AND unequal partials), landing
        // far from the truth. This pins the ∝|C_α| precondition — if the
        // fixture ever regressed to equal partials the bias would vanish and
        // this assertion would fail, exposing the vacuity.
        let all_indices: Vec<usize> = (0..dist.len()).collect();
        let uniform = echidna::stde::stde_sparse(&tape, &x, &dist, &all_indices);
        assert!(
            (uniform.estimate - exact).abs() > 1.0,
            "fixture must be non-vacuous: uniform enumeration {} should be far \
             from exact {} (importance sampling ∝|C_α| is required)",
            uniform.estimate,
            exact,
        );

        // Correct usage: draw indices ∝|C_α| via `sample_index`. The Monte
        // Carlo mean converges to the true operator value. Deterministic
        // quasi-random `u` sequence keeps the test reproducible across runs.
        let num_samples = 50_000;
        let indices: Vec<usize> = (0..num_samples)
            .map(|i| {
                let u = ((i as u64).wrapping_mul(2654435761) % 100_000) as f64 / 100_000.0;
                dist.sample_index(u)
            })
            .collect();
        let result = echidna::stde::stde_sparse(&tape, &x, &dist, &indices);

        let error = (result.estimate - exact).abs();
        let bound = 3.0 * result.standard_error;
        assert!(
            error < bound,
            "importance-sampled estimate {} too far from exact {}: error = {}, 3*SE = {}",
            result.estimate,
            exact,
            error,
            bound,
        );
    }
}

// ══════════════════════════════════════════════
//  22. Indefinite Dense STDE (requires stde + nalgebra)
// ══════════════════════════════════════════════

#[cfg(feature = "nalgebra")]
mod indefinite_stde_tests {
    use super::*;

    /// PD matrix: verify result matches dense_stde_2nd with same z-vectors.
    #[test]
    fn indefinite_stde_matches_positive_definite() {
        // C = [[4, 1], [1, 3]] (positive definite)
        // Cholesky: L = [[2, 0], [0.5, sqrt(2.75)]]
        let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
        let x = [1.0, 2.0];

        let c = nalgebra::DMatrix::from_row_slice(2, 2, &[4.0, 1.0, 1.0, 3.0]);

        // Cholesky factor for comparison
        let l00 = 2.0;
        let l10 = 0.5;
        let l11 = (3.0 - 0.25_f64).sqrt(); // sqrt(2.75)
        let row0 = vec![l00, 0.0];
        let row1 = vec![l10, l11];
        let cholesky: Vec<&[f64]> = vec![&row0, &row1];

        // Use Rademacher-like z-vectors
        let z0 = vec![1.0, 1.0];
        let z1 = vec![1.0, -1.0];
        let z2 = vec![-1.0, 1.0];
        let z3 = vec![-1.0, -1.0];
        let z_vecs: Vec<&[f64]> = vec![&z0, &z1, &z2, &z3];

        let chol_result = echidna::stde::dense_stde_2nd(&tape, &x, &cholesky, &z_vecs);
        let indef_result = echidna::stde::dense_stde_2nd_indefinite(&tape, &x, &c, &z_vecs, 1e-12);

        // H = diag(2, 2), tr(C·H) = 4*2 + 3*2 = 14
        assert_relative_eq!(chol_result.estimate, 14.0, epsilon = 1e-8);
        assert_relative_eq!(indef_result.estimate, 14.0, epsilon = 1e-8);
    }

    /// Diagonal indefinite C = diag(2, -3): verify tr(C·H) against analytical 2·H₀₀ - 3·H₁₁.
    #[test]
    fn indefinite_stde_diagonal_indefinite() {
        // f(x,y) = x² + y², H = diag(2, 2)
        // C = diag(2, -3), tr(C·H) = 2·2 + (-3)·2 = 4 - 6 = -2
        let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
        let x = [1.0, 2.0];

        let c = nalgebra::DMatrix::from_row_slice(2, 2, &[2.0, 0.0, 0.0, -3.0]);

        let z0 = vec![1.0, 1.0];
        let z1 = vec![1.0, -1.0];
        let z2 = vec![-1.0, 1.0];
        let z3 = vec![-1.0, -1.0];
        let z_vecs: Vec<&[f64]> = vec![&z0, &z1, &z2, &z3];

        let result = echidna::stde::dense_stde_2nd_indefinite(&tape, &x, &c, &z_vecs, 1e-12);
        assert_relative_eq!(result.estimate, -2.0, epsilon = 1e-8);
    }

    /// Full symmetric indefinite C, verify against tr(C·H) computed from dense Hessian.
    #[test]
    fn indefinite_stde_full_indefinite() {
        // f(x,y,z) = x²y + y³, H at (1,2,3):
        // H = [[2y, 2x, 0], [2x, 6y, 0], [0, 0, 0]] = [[4, 2, 0], [2, 12, 0], [0, 0, 0]]
        let tape = record_fn(cubic_mix, &[1.0, 2.0, 3.0]);
        let x = [1.0, 2.0, 3.0];

        // C = [[1, 0, -1], [0, -2, 0], [-1, 0, 3]] — indefinite
        let c = nalgebra::DMatrix::from_row_slice(
            3,
            3,
            &[1.0, 0.0, -1.0, 0.0, -2.0, 0.0, -1.0, 0.0, 3.0],
        );

        // tr(C·H) = Σ_{ij} C_{ij} H_{ij}
        // = 1·4 + 0·2 + (-1)·0 + 0·2 + (-2)·12 + 0·0 + (-1)·0 + 0·0 + 3·0
        // = 4 - 24 = -20
        let expected = -20.0;

        // Use many z-vectors for convergence (Rademacher-like from all sign combos)
        let signs: Vec<Vec<f64>> = vec![
            vec![1.0, 1.0, 1.0],
            vec![1.0, 1.0, -1.0],
            vec![1.0, -1.0, 1.0],
            vec![1.0, -1.0, -1.0],
            vec![-1.0, 1.0, 1.0],
            vec![-1.0, 1.0, -1.0],
            vec![-1.0, -1.0, 1.0],
            vec![-1.0, -1.0, -1.0],
        ];
        let z_vecs: Vec<&[f64]> = signs.iter().map(|v| v.as_slice()).collect();

        let result = echidna::stde::dense_stde_2nd_indefinite(&tape, &x, &c, &z_vecs, 1e-12);
        assert_relative_eq!(result.estimate, expected, epsilon = 1e-6);
    }

    /// All-negative eigenvalues: C = diag(-2, -3), H = diag(2, 2).
    /// tr(C·H) = -2·2 + (-3)·2 = -10.
    #[test]
    fn indefinite_stde_all_negative() {
        let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
        let x = [1.0, 2.0];

        let c = nalgebra::DMatrix::from_row_slice(2, 2, &[-2.0, 0.0, 0.0, -3.0]);

        let z0 = vec![1.0, 1.0];
        let z1 = vec![1.0, -1.0];
        let z2 = vec![-1.0, 1.0];
        let z3 = vec![-1.0, -1.0];
        let z_vecs: Vec<&[f64]> = vec![&z0, &z1, &z2, &z3];

        let result = echidna::stde::dense_stde_2nd_indefinite(&tape, &x, &c, &z_vecs, 1e-12);
        assert_relative_eq!(result.estimate, -10.0, epsilon = 1e-8);
    }

    /// C = 0: result should be 0.
    #[test]
    fn indefinite_stde_zero_matrix() {
        let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
        let x = [1.0, 2.0];

        let c = nalgebra::DMatrix::zeros(2, 2);

        let z0 = vec![1.0, 1.0];
        let z1 = vec![1.0, -1.0];
        let z_vecs: Vec<&[f64]> = vec![&z0, &z1];

        let result = echidna::stde::dense_stde_2nd_indefinite(&tape, &x, &c, &z_vecs, 1e-12);
        assert_relative_eq!(result.estimate, 0.0, epsilon = 1e-10);
    }

    /// C with eigenvalue ~1e-15: verify epsilon clamping prevents sign-flip.
    #[test]
    fn indefinite_stde_near_zero_eigenvalue() {
        // C = [[1, 0], [0, 1e-15]] — the tiny eigenvalue should be clamped to zero
        let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
        let x = [1.0, 2.0];

        let c = nalgebra::DMatrix::from_row_slice(2, 2, &[1.0, 0.0, 0.0, 1e-15]);

        let z0 = vec![1.0, 1.0];
        let z1 = vec![1.0, -1.0];
        let z2 = vec![-1.0, 1.0];
        let z3 = vec![-1.0, -1.0];
        let z_vecs: Vec<&[f64]> = vec![&z0, &z1, &z2, &z3];

        // With eps_factor=1e-12, threshold = 1e-12 * 1.0 = 1e-12.
        // The eigenvalue 1e-15 < 1e-12, so it's clamped to zero.
        // Result should be tr(diag(1,0) · diag(2,2)) = 1·2 + 0·2 = 2
        let result = echidna::stde::dense_stde_2nd_indefinite(&tape, &x, &c, &z_vecs, 1e-12);
        assert_relative_eq!(result.estimate, 2.0, epsilon = 1e-8);
    }
}
