//! Tests for AD types inside nalgebra matrices.

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

use approx::{assert_abs_diff_eq, assert_relative_eq, assert_ulps_eq};
use echidna::{Dual64, DualVec32, DualVec64, Reverse64};
use nalgebra::{Matrix3, Vector3};
use num_traits::Float;
use simba::scalar::SubsetOf;

// ── Dual<f64> in nalgebra ──

#[test]
fn dual_vector3_dot_product() {
    let a = Vector3::new(
        Dual64::variable(1.0),
        Dual64::constant(2.0),
        Dual64::constant(3.0),
    );
    let b = Vector3::new(
        Dual64::constant(4.0),
        Dual64::constant(5.0),
        Dual64::constant(6.0),
    );
    let dot = a.dot(&b);
    // dot = 1*4 + 2*5 + 3*6 = 32
    assert_relative_eq!(dot.re, 32.0, max_relative = 1e-12);
    // d(dot)/d(a[0]) = b[0] = 4
    assert_relative_eq!(dot.eps, 4.0, max_relative = 1e-12);
}

#[test]
fn dual_vector3_norm() {
    // v = [x, 2, 3], norm = sqrt(x² + 4 + 9)
    // d(norm)/dx = x / sqrt(x² + 13)
    let v = Vector3::new(
        Dual64::variable(3.0),
        Dual64::constant(2.0),
        Dual64::constant(3.0),
    );
    let n = v.norm();
    let expected_norm = (9.0 + 4.0 + 9.0_f64).sqrt();
    assert_relative_eq!(n.re, expected_norm, max_relative = 1e-12);
    let expected_deriv = 3.0 / expected_norm;
    assert_relative_eq!(n.eps, expected_deriv, max_relative = 1e-10);
}

#[test]
fn dual_matrix_vector_product() {
    // M * v where v[0] is the variable
    let m = Matrix3::new(
        Dual64::constant(1.0),
        Dual64::constant(2.0),
        Dual64::constant(3.0),
        Dual64::constant(4.0),
        Dual64::constant(5.0),
        Dual64::constant(6.0),
        Dual64::constant(7.0),
        Dual64::constant(8.0),
        Dual64::constant(9.0),
    );
    let v = Vector3::new(
        Dual64::variable(1.0),
        Dual64::constant(0.0),
        Dual64::constant(0.0),
    );
    let result = m * v;
    // result = [1, 4, 7]
    assert_relative_eq!(result[0].re, 1.0, max_relative = 1e-12);
    assert_relative_eq!(result[1].re, 4.0, max_relative = 1e-12);
    assert_relative_eq!(result[2].re, 7.0, max_relative = 1e-12);
    // d/dx: derivatives are first column of M
    assert_relative_eq!(result[0].eps, 1.0, max_relative = 1e-12);
    assert_relative_eq!(result[1].eps, 4.0, max_relative = 1e-12);
    assert_relative_eq!(result[2].eps, 7.0, max_relative = 1e-12);
}

// ── DualVec<f64, N> in nalgebra ──

#[test]
fn dual_vec_vector3_dot_product() {
    let a = Vector3::new(
        DualVec64::<3>::new(1.0, [1.0, 0.0, 0.0]),
        DualVec64::<3>::new(2.0, [0.0, 1.0, 0.0]),
        DualVec64::<3>::new(3.0, [0.0, 0.0, 1.0]),
    );
    let b = Vector3::new(
        DualVec64::<3>::constant(4.0),
        DualVec64::<3>::constant(5.0),
        DualVec64::<3>::constant(6.0),
    );
    let dot = a.dot(&b);
    // dot = 1*4 + 2*5 + 3*6 = 32
    assert_relative_eq!(dot.re, 32.0, max_relative = 1e-12);
    // d(dot)/d(a[0]) = b[0] = 4
    // d(dot)/d(a[1]) = b[1] = 5
    // d(dot)/d(a[2]) = b[2] = 6
    assert_relative_eq!(dot.eps[0], 4.0, max_relative = 1e-12);
    assert_relative_eq!(dot.eps[1], 5.0, max_relative = 1e-12);
    assert_relative_eq!(dot.eps[2], 6.0, max_relative = 1e-12);
}

#[test]
fn dual_vec_vector3_norm() {
    // v = [x, y, z], norm = sqrt(x² + y² + z²)
    // d(norm)/dv_i = v_i / norm
    let v = Vector3::new(
        DualVec64::<3>::new(3.0, [1.0, 0.0, 0.0]),
        DualVec64::<3>::new(2.0, [0.0, 1.0, 0.0]),
        DualVec64::<3>::new(3.0, [0.0, 0.0, 1.0]),
    );
    let n = v.norm();
    let expected_norm = (9.0 + 4.0 + 9.0_f64).sqrt();
    assert_relative_eq!(n.re, expected_norm, max_relative = 1e-12);
    let expected_deriv = v.map(|x| x.re / expected_norm);
    assert_relative_eq!(n.eps[0], expected_deriv[0], max_relative = 1e-10);
    assert_relative_eq!(n.eps[1], expected_deriv[1], max_relative = 1e-10);
    assert_relative_eq!(n.eps[2], expected_deriv[2], max_relative = 1e-10);
}

#[test]
fn dual_vec_matrix_vector_product() {
    // M * v where v[0] is the variable
    let m = Matrix3::new(
        DualVec64::<3>::constant(1.0),
        DualVec64::<3>::constant(2.0),
        DualVec64::<3>::constant(3.0),
        DualVec64::<3>::constant(4.0),
        DualVec64::<3>::constant(5.0),
        DualVec64::<3>::constant(6.0),
        DualVec64::<3>::constant(7.0),
        DualVec64::<3>::constant(8.0),
        DualVec64::<3>::constant(9.0),
    );
    let v = Vector3::new(
        DualVec64::<3>::new(1.0, [1.0, 0.0, 0.0]),
        DualVec64::<3>::new(3.0, [0.0, 1.0, 0.0]),
        DualVec64::<3>::new(5.0, [0.0, 0.0, 1.0]),
    );
    let result = m * v;
    // result = [22, 49, 76]
    assert_relative_eq!(result[0].re, 22.0, max_relative = 1e-12);
    assert_relative_eq!(result[1].re, 49.0, max_relative = 1e-12);
    assert_relative_eq!(result[2].re, 76.0, max_relative = 1e-12);
    // dv_i/dx_j = M_ij
    for i in 0..3 {
        for j in 0..3 {
            assert_relative_eq!(result[i].eps[j], m[(i, j)].re, max_relative = 1e-12);
        }
    }
}

#[test]
fn dual_vec_eq() {
    let a = DualVec64::<3>::new(1.0, [1.0, 0.0, 0.0]);
    let b = DualVec64::<3>::new(1.0, [0.0, 1.0, 0.0]);
    // Verify that AbsDiffEq, RelativeEq, and UlpsEq compare only the real part
    assert_abs_diff_eq!(a, b, epsilon = DualVec64::<3>::epsilon());
    assert_relative_eq!(a, b, epsilon = DualVec64::<3>::epsilon());
    assert_ulps_eq!(a, b, max_ulps = 1);
}

// ── Reverse<f64> in nalgebra ──

#[test]
fn reverse_vector3_dot_product() {
    let g = echidna::grad(
        |x| {
            let a = Vector3::new(x[0], x[1], x[2]);
            let b = Vector3::new(
                Reverse64::constant(4.0),
                Reverse64::constant(5.0),
                Reverse64::constant(6.0),
            );
            a.dot(&b)
        },
        &[1.0, 2.0, 3.0],
    );
    // d(a·b)/d(a) = b
    assert_relative_eq!(g[0], 4.0, max_relative = 1e-12);
    assert_relative_eq!(g[1], 5.0, max_relative = 1e-12);
    assert_relative_eq!(g[2], 6.0, max_relative = 1e-12);
}

#[test]
fn reverse_vector3_norm_squared() {
    // f(x) = ||x||² = x₀² + x₁² + x₂²
    // ∇f = [2x₀, 2x₁, 2x₂]
    let g = echidna::grad(
        |x| {
            let v = Vector3::new(x[0], x[1], x[2]);
            v.dot(&v)
        },
        &[1.0, 2.0, 3.0],
    );
    assert_relative_eq!(g[0], 2.0, max_relative = 1e-12);
    assert_relative_eq!(g[1], 4.0, max_relative = 1e-12);
    assert_relative_eq!(g[2], 6.0, max_relative = 1e-12);
}

#[test]
fn grad_of_nalgebra_function() {
    // f(x) = sin(||x||)
    let g = echidna::grad(
        |x| {
            let v = Vector3::new(x[0], x[1], x[2]);
            let n: Reverse64 = v.norm();
            n.sin()
        },
        &[1.0, 2.0, 3.0],
    );
    // ∂f/∂xᵢ = cos(||x||) * xᵢ / ||x||
    let norm = (1.0 + 4.0 + 9.0_f64).sqrt();
    for (i, &xi) in [1.0, 2.0, 3.0].iter().enumerate() {
        let expected = norm.cos() * xi / norm;
        assert_relative_eq!(g[i], expected, max_relative = 1e-10);
    }
}

// ── Matrix inverse (exercises ComplexField / RealField deeply) ──

#[test]
fn dual_matrix3_try_inverse() {
    // A 3×3 matrix of constants — test that try_inverse compiles and runs.
    let m = Matrix3::new(
        Dual64::constant(2.0),
        Dual64::constant(1.0),
        Dual64::constant(0.0),
        Dual64::constant(1.0),
        Dual64::constant(3.0),
        Dual64::constant(1.0),
        Dual64::constant(0.0),
        Dual64::constant(1.0),
        Dual64::constant(2.0),
    );
    let inv = m.try_inverse().expect("matrix should be invertible");
    let identity = m * inv;
    for i in 0..3 {
        for j in 0..3 {
            let expected = if i == j { 1.0 } else { 0.0 };
            assert_relative_eq!(identity[(i, j)].re, expected, max_relative = 1e-10);
        }
    }
}

#[test]
fn dual_vec_matrix3_try_inverse() {
    // A 3×3 matrix of dual vectors
    let m = Matrix3::new(
        DualVec64::<3>::new(2.0, [1.0, 0.0, 0.0]),
        DualVec64::<3>::new(1.0, [0.0, 0.0, 0.0]),
        DualVec64::<3>::new(0.0, [0.0, 0.0, 0.0]),
        DualVec64::<3>::new(1.0, [0.0, 0.0, 0.0]),
        DualVec64::<3>::new(3.0, [0.0, 1.0, 0.0]),
        DualVec64::<3>::new(1.0, [0.0, 0.0, 0.0]),
        DualVec64::<3>::new(0.0, [0.0, 0.0, 0.0]),
        DualVec64::<3>::new(1.0, [0.0, 0.0, 0.0]),
        DualVec64::<3>::new(2.0, [0.0, 0.0, 1.0]),
    );
    let inv = m.try_inverse().expect("matrix should be invertible");
    // Verify that matrix times inv(matrix) is the identity with all entries with zero differential part
    let identity = m * inv;
    for i in 0..3 {
        for j in 0..3 {
            let actual = identity[(i, j)];
            let expected = if i == j { 1.0 } else { 0.0 };
            assert_relative_eq!(actual.re, expected, max_relative = 1e-10);
            for k in 0..3 {
                assert_relative_eq!(actual.eps[k], 0.0, max_relative = 1e-10);
            }
        }
    }
    // Verify finite-diff agreement.
    let h = 1e-7;
    let f = |a: f64| {
        let m = Matrix3::new(a, 1.0, 0.0, 1.0, 3.0, 1.0, 0.0, 1.0, 2.0);
        let inv = m.try_inverse().unwrap();
        inv[(0, 0)] + inv[(1, 1)] + inv[(2, 2)]
    };
    let fd = (f(2.0 + h) - f(2.0 - h)) / (2.0 * h);
    assert_relative_eq!(
        (inv[(0, 0)] + inv[(1, 1)] + inv[(2, 2)]).eps[0],
        fd,
        max_relative = 1e-4
    );
}

#[test]
fn reverse_matrix3_try_inverse() {
    // Same test for Reverse — validates the full ComplexField/RealField chain.
    let g = echidna::grad(
        |x| {
            let m = Matrix3::new(
                x[0],
                Reverse64::constant(1.0),
                Reverse64::constant(0.0),
                Reverse64::constant(1.0),
                Reverse64::constant(3.0),
                Reverse64::constant(1.0),
                Reverse64::constant(0.0),
                Reverse64::constant(1.0),
                Reverse64::constant(2.0),
            );
            let inv = m.try_inverse().expect("invertible");
            // Return trace of inverse (scalar output)
            inv[(0, 0)] + inv[(1, 1)] + inv[(2, 2)]
        },
        &[2.0],
    );
    // Verify finite-diff agreement.
    let h = 1e-7;
    let f = |a: f64| {
        let m = Matrix3::new(a, 1.0, 0.0, 1.0, 3.0, 1.0, 0.0, 1.0, 2.0);
        let inv = m.try_inverse().unwrap();
        inv[(0, 0)] + inv[(1, 1)] + inv[(2, 2)]
    };
    let fd = (f(2.0 + h) - f(2.0 - h)) / (2.0 * h);
    assert_relative_eq!(g[0], fd, max_relative = 1e-4);
}

// ── DualVec<F, N> SubsetOf conversions ──
//
// The `SubsetOf` impls are required for the `ComplexField` / `RealField` bounds,
// but nalgebra's matrix ops don't exercise `is_in_subset` / `from_superset` at
// runtime — so these direct tests guard against a silent regression in the
// conversion logic (e.g. `is_in_subset` no longer rejecting live tangents).

#[test]
fn dual_vec_subsetof_identity() {
    // DualVec<F, N> ⊂ DualVec<F, N>: the identity embedding round-trips exactly
    // and is always in-subset — even with non-zero tangents.
    let dv = DualVec64::<3>::new(2.5, [1.0, -2.0, 3.0]);
    let up: DualVec64<3> = dv.to_superset();
    assert_eq!(up.re, dv.re);
    assert_eq!(up.eps, dv.eps);
    let down = <DualVec64<3> as SubsetOf<DualVec64<3>>>::from_superset_unchecked(&dv);
    assert_eq!(down.re, dv.re);
    assert_eq!(down.eps, dv.eps);
    assert!(<DualVec64<3> as SubsetOf<DualVec64<3>>>::is_in_subset(&dv));
}

#[test]
fn dual_vec_subsetof_f64_embedding() {
    // f64 ⊂ DualVec<f64, N>: a scalar embeds as a constant dual vector.
    let dv: DualVec64<3> = 3.5_f64.to_superset();
    assert_eq!(dv.re, 3.5);
    assert!(
        dv.eps.iter().all(|&e| e == 0.0),
        "an embedded scalar must have zero tangents"
    );
    assert_eq!(
        <f64 as SubsetOf<DualVec64<3>>>::from_superset_unchecked(&dv),
        3.5
    );

    // is_in_subset accepts a constant dual but rejects one carrying a live
    // tangent; from_superset (the checked projection) follows suit.
    let with_tangent = DualVec64::<3>::new(3.5, [1.0, 0.0, 0.0]);
    assert!(<f64 as SubsetOf<DualVec64<3>>>::is_in_subset(&dv));
    assert!(!<f64 as SubsetOf<DualVec64<3>>>::is_in_subset(
        &with_tangent
    ));
    assert_eq!(
        <f64 as SubsetOf<DualVec64<3>>>::from_superset(&dv),
        Some(3.5)
    );
    assert_eq!(
        <f64 as SubsetOf<DualVec64<3>>>::from_superset(&with_tangent),
        None
    );
}

#[test]
fn dual_vec_subsetof_f32_embedding() {
    // f32 ⊂ DualVec<f32, N>.
    let dv: DualVec32<3> = 1.25_f32.to_superset();
    assert_eq!(dv.re, 1.25_f32);
    assert!(dv.eps.iter().all(|&e| e == 0.0));
    assert_eq!(
        <f32 as SubsetOf<DualVec32<3>>>::from_superset_unchecked(&dv),
        1.25_f32
    );

    let with_tangent = DualVec32::<3>::new(1.25, [1.0, 0.0, 0.0]);
    assert!(<f32 as SubsetOf<DualVec32<3>>>::is_in_subset(&dv));
    assert!(!<f32 as SubsetOf<DualVec32<3>>>::is_in_subset(
        &with_tangent
    ));
}

#[test]
fn dual_vec_subsetof_cross_precision() {
    // f32 ⊂ DualVec<f64, N> — widening is always lossless.
    let widened: DualVec64<3> = 0.1_f32.to_superset();
    assert_eq!(widened.re, 0.1_f32 as f64);
    assert!(widened.eps.iter().all(|&e| e == 0.0));
    assert_eq!(
        <f32 as SubsetOf<DualVec64<3>>>::from_superset_unchecked(&widened),
        0.1_f32
    );
    assert!(<f32 as SubsetOf<DualVec64<3>>>::is_in_subset(&widened));
    let dv64_with_tangent = DualVec64::<3>::new(0.5, [1.0, 0.0, 0.0]);
    assert!(!<f32 as SubsetOf<DualVec64<3>>>::is_in_subset(
        &dv64_with_tangent
    ));

    // f64 ⊂ DualVec<f32, N> — narrowing is lossy (required by the SupersetOf<f64>
    // bound on ComplexField). 0.1 is not f32-representable, so projecting back to
    // f64 does not recover the original value.
    let narrowed: DualVec32<3> = 0.1_f64.to_superset();
    assert_eq!(narrowed.re, 0.1_f64 as f32);
    assert!(narrowed.eps.iter().all(|&e| e == 0.0));
    let back = <f64 as SubsetOf<DualVec32<3>>>::from_superset_unchecked(&narrowed);
    assert_ne!(back, 0.1_f64);
    assert_eq!(back, 0.1_f64 as f32 as f64);
    assert!(<f64 as SubsetOf<DualVec32<3>>>::is_in_subset(&narrowed));
    let dv32_with_tangent = DualVec32::<3>::new(0.5, [2.0, 0.0, 0.0]);
    assert!(!<f64 as SubsetOf<DualVec32<3>>>::is_in_subset(
        &dv32_with_tangent
    ));
}
