#![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
/// Gradient: [2x, 2y]
/// Hessian: [[2, 0], [0, 2]]
/// Laplacian: 4
/// Diagonal: [2, 2]
fn sum_of_squares<T: Scalar>(x: &[T]) -> T {
    x[0] * x[0] + x[1] * x[1]
}

/// f(x,y) = x*y
/// Gradient: [y, x]
/// Hessian: [[0, 1], [1, 0]]
/// Laplacian: 0
/// Diagonal: [0, 0]
fn product<T: Scalar>(x: &[T]) -> T {
    x[0] * x[1]
}

/// f(x,y,z) = x^2*y + y^3
/// At (1, 2, 3):
/// f = 1*2 + 8 = 10
/// Gradient: [2xy, x^2+3y^2, 0] = [4, 13, 0]
/// Hessian: [[2y, 2x, 0], [2x, 6y, 0], [0, 0, 0]]
///        = [[4, 2, 0], [2, 12, 0], [0, 0, 0]]
/// Laplacian: 4 + 12 + 0 = 16
/// Diagonal: [4, 12, 0]
fn cubic_mix<T: Scalar>(x: &[T]) -> T {
    x[0] * x[0] * x[1] + x[1] * x[1] * x[1]
}

/// f(x) = x^3 (1D)
/// f''(x) = 6x, so at x=2: f''=12
fn cube_1d<T: Scalar>(x: &[T]) -> T {
    x[0] * x[0] * x[0]
}

/// f(x,y) = x + y (linear, all second derivatives zero)
fn linear_fn<T: Scalar>(x: &[T]) -> T {
    x[0] + x[1]
}

// ══════════════════════════════════════════════
//  1. Known Hessians via Rademacher vectors
// ══════════════════════════════════════════════

#[test]
fn laplacian_sum_of_squares() {
    let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
    // Rademacher vectors: entries +/-1, E[vv^T] = I
    let v0: Vec<f64> = vec![1.0, 1.0];
    let v1: Vec<f64> = vec![1.0, -1.0];
    let dirs: Vec<&[f64]> = vec![&v0, &v1];

    let (value, lap) = echidna::stde::laplacian(&tape, &[1.0, 2.0], &dirs);
    assert_relative_eq!(value, 5.0, epsilon = 1e-10);
    assert_relative_eq!(lap, 4.0, epsilon = 1e-10);
}

#[test]
fn laplacian_product() {
    let tape = record_fn(product, &[3.0, 4.0]);
    // Rademacher: v^T [[0,1],[1,0]] v = 2*v0*v1
    // For [1,1]: 2. For [1,-1]: -2. Average of 2*c2 = average(2, -2) = 0.
    let v0: Vec<f64> = vec![1.0, 1.0];
    let v1: Vec<f64> = vec![1.0, -1.0];
    let dirs: Vec<&[f64]> = vec![&v0, &v1];

    let (value, lap) = echidna::stde::laplacian(&tape, &[3.0, 4.0], &dirs);
    assert_relative_eq!(value, 12.0, epsilon = 1e-10);
    assert_relative_eq!(lap, 0.0, epsilon = 1e-10);
}

#[test]
fn laplacian_cubic_mix() {
    // H = [[4, 2, 0], [2, 12, 0], [0, 0, 0]], tr(H) = 16
    // Use all 8 Rademacher vectors for exact result
    let tape = record_fn(cubic_mix, &[1.0, 2.0, 3.0]);

    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 dirs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();

    let (value, lap) = echidna::stde::laplacian(&tape, &[1.0, 2.0, 3.0], &dirs);
    assert_relative_eq!(value, 10.0, epsilon = 1e-10);
    assert_relative_eq!(lap, 16.0, epsilon = 1e-10);
}

// ══════════════════════════════════════════════
//  2. Hessian diagonal
// ══════════════════════════════════════════════

#[test]
fn hessian_diagonal_sum_of_squares() {
    let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
    let (value, diag) = echidna::stde::hessian_diagonal(&tape, &[1.0, 2.0]);
    assert_relative_eq!(value, 5.0, epsilon = 1e-10);
    assert_eq!(diag.len(), 2);
    assert_relative_eq!(diag[0], 2.0, epsilon = 1e-10);
    assert_relative_eq!(diag[1], 2.0, epsilon = 1e-10);
}

#[test]
fn hessian_diagonal_product() {
    let tape = record_fn(product, &[3.0, 4.0]);
    let (_, diag) = echidna::stde::hessian_diagonal(&tape, &[3.0, 4.0]);
    assert_relative_eq!(diag[0], 0.0, epsilon = 1e-10);
    assert_relative_eq!(diag[1], 0.0, epsilon = 1e-10);
}

#[test]
fn hessian_diagonal_cubic_mix() {
    let tape = record_fn(cubic_mix, &[1.0, 2.0, 3.0]);
    let (_, diag) = echidna::stde::hessian_diagonal(&tape, &[1.0, 2.0, 3.0]);
    assert_eq!(diag.len(), 3);
    assert_relative_eq!(diag[0], 4.0, epsilon = 1e-10); // 2y = 4
    assert_relative_eq!(diag[1], 12.0, epsilon = 1e-10); // 6y = 12
    assert_relative_eq!(diag[2], 0.0, epsilon = 1e-10);
}

// ══════════════════════════════════════════════
//  3. Coordinate-basis Laplacian via hessian_diagonal sum
// ══════════════════════════════════════════════

#[test]
fn coordinate_basis_laplacian_via_diagonal_sum() {
    // The exact Laplacian is the sum of the Hessian diagonal.
    let tape = record_fn(cubic_mix, &[1.0, 2.0, 3.0]);
    let (_, diag) = echidna::stde::hessian_diagonal(&tape, &[1.0, 2.0, 3.0]);
    let laplacian: f64 = diag.iter().sum();
    assert_relative_eq!(laplacian, 16.0, epsilon = 1e-10);
}

// ══════════════════════════════════════════════
//  4. Cross-validation with hessian()
// ══════════════════════════════════════════════

#[test]
fn cross_validate_with_hessian_sum_of_squares() {
    let x = [1.0, 2.0];
    let (val_h, _grad, hess) = echidna::hessian(sum_of_squares, &x);

    let trace: f64 = (0..x.len()).map(|i| hess[i][i]).sum();
    let diag_from_hess: Vec<f64> = (0..x.len()).map(|i| hess[i][i]).collect();

    // Compare Laplacian via Rademacher
    let tape = record_fn(sum_of_squares, &x);
    let v0: Vec<f64> = vec![1.0, 1.0];
    let v1: Vec<f64> = vec![1.0, -1.0];
    let v2: Vec<f64> = vec![-1.0, 1.0];
    let v3: Vec<f64> = vec![-1.0, -1.0];
    let dirs: Vec<&[f64]> = vec![&v0, &v1, &v2, &v3];
    let (val_s, lap) = echidna::stde::laplacian(&tape, &x, &dirs);

    // Compare diagonal
    let (_, diag) = echidna::stde::hessian_diagonal(&tape, &x);

    assert_relative_eq!(val_h, val_s, epsilon = 1e-10);
    assert_relative_eq!(trace, lap, epsilon = 1e-10);
    for j in 0..x.len() {
        assert_relative_eq!(diag_from_hess[j], diag[j], epsilon = 1e-10);
    }
}

#[test]
fn cross_validate_with_hessian_cubic_mix() {
    let x = [1.0, 2.0, 3.0];
    let (val_h, _grad, hess) = echidna::hessian(cubic_mix, &x);

    let trace: f64 = (0..x.len()).map(|i| hess[i][i]).sum();
    let diag_from_hess: Vec<f64> = (0..x.len()).map(|i| hess[i][i]).collect();

    // Compare Laplacian via all 8 Rademacher vectors (exact)
    let tape = record_fn(cubic_mix, &x);
    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 dirs: Vec<&[f64]> = vecs.iter().map(|v| v.as_slice()).collect();
    let (val_s, lap) = echidna::stde::laplacian(&tape, &x, &dirs);

    let (_, diag) = echidna::stde::hessian_diagonal(&tape, &x);

    assert_relative_eq!(val_h, val_s, epsilon = 1e-10);
    assert_relative_eq!(trace, lap, epsilon = 1e-10);
    for j in 0..x.len() {
        assert_relative_eq!(diag_from_hess[j], diag[j], epsilon = 1e-10);
    }
}

// ══════════════════════════════════════════════
//  5. Cross-validation with hvp() / grad()
// ══════════════════════════════════════════════

#[test]
fn directional_derivative_matches_gradient() {
    // c1 for basis vector e_j should equal partial_j f
    let x = [1.0, 2.0, 3.0];
    let tape = record_fn(cubic_mix, &x);

    let grad = echidna::grad(cubic_mix, &x);

    let e0: Vec<f64> = vec![1.0, 0.0, 0.0];
    let e1: Vec<f64> = vec![0.0, 1.0, 0.0];
    let e2: Vec<f64> = vec![0.0, 0.0, 1.0];
    let dirs: Vec<&[f64]> = vec![&e0, &e1, &e2];

    let (_, first_order, _) = echidna::stde::directional_derivatives(&tape, &x, &dirs);

    for j in 0..x.len() {
        assert_relative_eq!(first_order[j], grad[j], epsilon = 1e-10);
    }
}

#[test]
fn directional_derivative_arbitrary_direction() {
    // For arbitrary v, c1 should equal grad . v
    let x = [1.0, 2.0, 3.0];
    let tape = record_fn(cubic_mix, &x);

    let grad = echidna::grad(cubic_mix, &x);
    let v: Vec<f64> = vec![0.5, -1.0, 2.0];
    let expected_c1: f64 = grad.iter().zip(v.iter()).map(|(g, vi)| g * vi).sum();

    let (_, c1, _) = echidna::stde::taylor_jet_2nd(&tape, &x, &v);
    assert_relative_eq!(c1, expected_c1, epsilon = 1e-10);
}

// ══════════════════════════════════════════════
//  7. Edge cases
// ══════════════════════════════════════════════

#[test]
fn n_equals_1() {
    let tape = record_fn(cube_1d, &[2.0]);
    let (value, diag) = echidna::stde::hessian_diagonal(&tape, &[2.0]);
    assert_relative_eq!(value, 8.0, epsilon = 1e-10);
    assert_eq!(diag.len(), 1);
    assert_relative_eq!(diag[0], 12.0, epsilon = 1e-10); // f''(2) = 6*2 = 12
}

#[test]
fn linear_function_all_zeros() {
    // f(x,y) = x + y — all second derivatives zero
    let tape = record_fn(linear_fn, &[1.0, 2.0]);
    let v0: Vec<f64> = vec![1.0, 1.0];
    let v1: Vec<f64> = vec![1.0, -1.0];
    let dirs: Vec<&[f64]> = vec![&v0, &v1];

    let (value, lap) = echidna::stde::laplacian(&tape, &[1.0, 2.0], &dirs);
    assert_relative_eq!(value, 3.0, epsilon = 1e-10);
    assert_relative_eq!(lap, 0.0, epsilon = 1e-10);

    let (_, diag) = echidna::stde::hessian_diagonal(&tape, &[1.0, 2.0]);
    assert_relative_eq!(diag[0], 0.0, epsilon = 1e-10);
    assert_relative_eq!(diag[1], 0.0, epsilon = 1e-10);
}

#[test]
fn single_direction() {
    let tape = record_fn(sum_of_squares, &[1.0, 2.0]);
    let v: Vec<f64> = vec![1.0, 0.0];
    let dirs: Vec<&[f64]> = vec![&v];

    let (_, first_order, second_order) =
        echidna::stde::directional_derivatives(&tape, &[1.0, 2.0], &dirs);
    assert_eq!(first_order.len(), 1);
    assert_eq!(second_order.len(), 1);
    // c1 = grad . e0 = 2*1 = 2
    assert_relative_eq!(first_order[0], 2.0, epsilon = 1e-10);
    // c2 = e0^T H e0 / 2 = 2/2 = 1
    assert_relative_eq!(second_order[0], 1.0, epsilon = 1e-10);
}

// ══════════════════════════════════════════════
//  9. With_buf reuse produces consistent results
// ══════════════════════════════════════════════

#[test]
fn buf_reuse_consistency() {
    let tape = record_fn(cubic_mix, &[1.0, 2.0, 3.0]);
    let x = [1.0, 2.0, 3.0];
    let v = [0.5, -1.0, 2.0];

    let (c0a, c1a, c2a) = echidna::stde::taylor_jet_2nd(&tape, &x, &v);

    let mut buf = Vec::new();
    let (c0b, c1b, c2b) = echidna::stde::taylor_jet_2nd_with_buf(&tape, &x, &v, &mut buf);
    // Reuse same buffer
    let (c0c, c1c, c2c) = echidna::stde::taylor_jet_2nd_with_buf(&tape, &x, &v, &mut buf);

    assert_relative_eq!(c0a, c0b, epsilon = 1e-14);
    assert_relative_eq!(c1a, c1b, epsilon = 1e-14);
    assert_relative_eq!(c2a, c2b, epsilon = 1e-14);

    assert_relative_eq!(c0b, c0c, epsilon = 1e-14);
    assert_relative_eq!(c1b, c1c, epsilon = 1e-14);
    assert_relative_eq!(c2b, c2c, epsilon = 1e-14);
}
