use nalgebra::SVector;

use super::super::ode::ODE;
use super::Integrator;

/// Bogacki-Shampine 3/2 integrator trait for tableau coefficients
pub trait BS3Integrator<'a> {
    const A: &'a [f64];
    const B: &'a [f64];
    const B_ERROR: &'a [f64];
    const C: &'a [f64];
}

/// Bogacki-Shampine 3(2) method
///
/// A 3rd order explicit Runge-Kutta method with an embedded 2nd order method for
/// error estimation. This method is efficient for moderate accuracy requirements
/// (tolerances around 1e-3 to 1e-6) and uses fewer stages than Dormand-Prince 4(5).
///
/// # Characteristics
/// - Order: 3(2) - 3rd order solution with 2nd order error estimate
/// - Stages: 4
/// - FSAL: Yes (First Same As Last - reuses last function evaluation)
/// - Adaptive: Yes
/// - Dense output: 3rd order Hermite interpolation
///
/// # When to use BS3
/// - Problems requiring moderate accuracy (rtol ~ 1e-3 to 1e-6)
/// - When function evaluations are expensive (fewer stages than DP5)
/// - Non-stiff problems
///
/// # Example
/// ```rust
/// use ordinary_diffeq::prelude::*;
/// use nalgebra::Vector1;
///
/// let params = ();
/// fn derivative(_t: f64, y: Vector1<f64>, _p: &()) -> Vector1<f64> {
///     Vector1::new(-y[0])
/// }
///
/// let y0 = Vector1::new(1.0);
/// let ode = ODE::new(&derivative, 0.0, 5.0, y0, ());
/// let bs3 = BS3::new().a_tol(1e-6).r_tol(1e-4);
/// let controller = PIController::default();
///
/// let mut problem = Problem::new(ode, bs3, controller);
/// let solution = problem.solve();
/// ```
///
/// # References
/// - Bogacki, P. and Shampine, L.F. (1989), "A 3(2) pair of Runge-Kutta formulas",
///   Applied Mathematics Letters, Vol. 2, No. 4, pp. 321-325
#[derive(Debug, Clone, Copy)]
pub struct BS3<const D: usize> {
    a_tol: SVector<f64, D>,
    r_tol: f64,
}

impl<const D: usize> BS3<D>
where
    BS3<D>: Integrator<D>,
{
    /// Create a new BS3 integrator with default tolerances
    ///
    /// Default: atol = 1e-8, rtol = 1e-8
    pub fn new() -> Self {
        Self {
            a_tol: SVector::<f64, D>::from_element(1e-8),
            r_tol: 1e-8,
        }
    }

    /// Set absolute tolerance (same value for all components)
    pub fn a_tol(mut self, a_tol: f64) -> Self {
        self.a_tol = SVector::<f64, D>::from_element(a_tol);
        self
    }

    /// Set absolute tolerance (different value per component)
    pub fn a_tol_full(mut self, a_tol: SVector<f64, D>) -> Self {
        self.a_tol = a_tol;
        self
    }

    /// Set relative tolerance
    pub fn r_tol(mut self, r_tol: f64) -> Self {
        self.r_tol = r_tol;
        self
    }
}

impl<'a, const D: usize> BS3Integrator<'a> for BS3<D> {
    // Butcher tableau for BS3
    // The A matrix is stored in lower-triangular form as a flat array
    // Row 1: []
    // Row 2: [1/2]
    // Row 3: [0, 3/4]
    // Row 4: [2/9, 1/3, 4/9]
    const A: &'a [f64] = &[
        1.0 / 2.0,           // a[1,0]
        0.0,                 // a[2,0]
        3.0 / 4.0,           // a[2,1]
        2.0 / 9.0,           // a[3,0]
        1.0 / 3.0,           // a[3,1]
        4.0 / 9.0,           // a[3,2]
    ];

    // Solution weights (3rd order)
    const B: &'a [f64] = &[
        2.0 / 9.0,           // b[0]
        1.0 / 3.0,           // b[1]
        4.0 / 9.0,           // b[2]
        0.0,                 // b[3] - FSAL property: this is zero
    ];

    // Error estimate weights (difference between 3rd and 2nd order)
    const B_ERROR: &'a [f64] = &[
        2.0 / 9.0 - 7.0 / 24.0,      // b[0] - b*[0]
        1.0 / 3.0 - 1.0 / 4.0,       // b[1] - b*[1]
        4.0 / 9.0 - 1.0 / 3.0,       // b[2] - b*[2]
        0.0 - 1.0 / 8.0,             // b[3] - b*[3]
    ];

    // Stage times
    const C: &'a [f64] = &[
        0.0,                 // c[0]
        1.0 / 2.0,           // c[1]
        3.0 / 4.0,           // c[2]
        1.0,                 // c[3]
    ];
}

impl<'a, const D: usize> Integrator<D> for BS3<D>
where
    BS3<D>: BS3Integrator<'a>,
{
    const ORDER: usize = 3;
    const STAGES: usize = 4;
    const ADAPTIVE: bool = true;
    const DENSE: bool = true;

    fn step<P>(
        &self,
        ode: &ODE<D, P>,
        h: f64,
    ) -> (SVector<f64, D>, Option<f64>, Option<Vec<SVector<f64, D>>>) {
        // Allocate storage for the 4 stages
        let mut k: Vec<SVector<f64, D>> = vec![SVector::<f64, D>::zeros(); Self::STAGES];

        // Stage 1: k1 = f(t, y)
        k[0] = (ode.f)(ode.t, ode.y, &ode.params);

        // Stage 2: k2 = f(t + c[1]*h, y + h*a[1,0]*k1)
        let y2 = ode.y + h * Self::A[0] * k[0];
        k[1] = (ode.f)(ode.t + Self::C[1] * h, y2, &ode.params);

        // Stage 3: k3 = f(t + c[2]*h, y + h*(a[2,0]*k1 + a[2,1]*k2))
        let y3 = ode.y + h * (Self::A[1] * k[0] + Self::A[2] * k[1]);
        k[2] = (ode.f)(ode.t + Self::C[2] * h, y3, &ode.params);

        // Stage 4: k4 = f(t + c[3]*h, y + h*(a[3,0]*k1 + a[3,1]*k2 + a[3,2]*k3))
        let y4 = ode.y + h * (Self::A[3] * k[0] + Self::A[4] * k[1] + Self::A[5] * k[2]);
        k[3] = (ode.f)(ode.t + Self::C[3] * h, y4, &ode.params);

        // Compute 3rd order solution
        let next_y = ode.y + h * (Self::B[0] * k[0] + Self::B[1] * k[1] + Self::B[2] * k[2] + Self::B[3] * k[3]);

        // Compute error estimate (difference between 3rd and 2nd order solutions)
        let err = h * (Self::B_ERROR[0] * k[0] + Self::B_ERROR[1] * k[1] + Self::B_ERROR[2] * k[2] + Self::B_ERROR[3] * k[3]);

        // Compute error norm scaled by tolerance
        let tol = self.a_tol + ode.y.abs() * self.r_tol;
        let error_norm = (err.component_div(&tol)).norm();

        // Store coefficients for dense output (cubic Hermite interpolation)
        // BS3 uses standard cubic Hermite interpolation with derivatives at endpoints
        // Store: y0, y1, f0=k[0], f1=k[3] (FSAL)
        let dense_coeffs = vec![
            ode.y,           // y0 at start of step
            next_y,          // y1 at end of step
            k[0],            // f(t0, y0) - derivative at start
            k[3],            // f(t1, y1) - derivative at end (FSAL)
        ];

        (next_y, Some(error_norm), Some(dense_coeffs))
    }

    fn interpolate(
        &self,
        t_start: f64,
        t_end: f64,
        dense: &[SVector<f64, D>],
        t: f64,
    ) -> SVector<f64, D> {
        // Compute interpolation parameter θ ∈ [0, 1]
        let theta = (t - t_start) / (t_end - t_start);
        let h = t_end - t_start;

        // Cubic Hermite interpolation using values and derivatives at endpoints
        // dense[0] = y0 (value at start)
        // dense[1] = y1 (value at end)
        // dense[2] = f0 (derivative at start)
        // dense[3] = f1 (derivative at end)
        //
        // Standard cubic Hermite formula:
        // y(θ) = (1 + 2θ)(1-θ)²*y0 + θ²(3-2θ)*y1 + θ(1-θ)²*h*f0 + θ²(θ-1)*h*f1
        //
        // Equivalently (Horner form):
        // y(θ) = y0 + θ*[h*f0 + θ*(-3*y0 - 2*h*f0 + 3*y1 - h*f1 + θ*(2*y0 + h*f0 - 2*y1 + h*f1))]

        let y0 = &dense[0];
        let y1 = &dense[1];
        let f0 = &dense[2];
        let f1 = &dense[3];

        let theta2 = theta * theta;
        let one_minus_theta = 1.0 - theta;
        let one_minus_theta2 = one_minus_theta * one_minus_theta;

        // Apply cubic Hermite interpolation formula
        (1.0 + 2.0 * theta) * one_minus_theta2 * y0
            + theta2 * (3.0 - 2.0 * theta) * y1
            + theta * one_minus_theta2 * h * f0
            + theta2 * (theta - 1.0) * h * f1
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use nalgebra::Vector1;

    #[test]
    fn test_bs3_creation() {
        let _bs3: BS3<1> = BS3::new();
        assert_eq!(BS3::<1>::ORDER, 3);
        assert_eq!(BS3::<1>::STAGES, 4);
        assert!(BS3::<1>::ADAPTIVE);
        assert!(BS3::<1>::DENSE);
    }

    #[test]
    fn test_bs3_step() {
        type Params = ();
        fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
            Vector1::new(y[0]) // y' = y, solution is e^t
        }

        let y0 = Vector1::new(1.0);
        let ode = ODE::new(&derivative, 0.0, 1.0, y0, ());

        let bs3 = BS3::new();
        let h = 0.001;  // Smaller step size for tighter tolerances

        let (y_next, err, dense) = bs3.step(&ode, h);

        // At t=0.001, exact solution is e^0.001 ≈ 1.0010005001667084
        let exact = (0.001_f64).exp();
        assert_relative_eq!(y_next[0], exact, max_relative = 1e-6);

        // Error should be reasonable for h=0.001
        assert!(err.is_some());
        // The error estimate is scaled by tolerance, so err < 1 means step is acceptable
        assert!(err.unwrap() < 1.0);

        // Dense output should be provided
        assert!(dense.is_some());
        assert_eq!(dense.unwrap().len(), 4);
    }

    #[test]
    fn test_bs3_interpolation() {
        type Params = ();
        fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
            Vector1::new(y[0])
        }

        let y0 = Vector1::new(1.0);
        let ode = ODE::new(&derivative, 0.0, 1.0, y0, ());

        let bs3 = BS3::new();
        let h = 0.001;  // Smaller step size

        let (_y_next, _err, dense) = bs3.step(&ode, h);
        let dense = dense.unwrap();

        // Interpolate at midpoint
        let t_mid = 0.0005;
        let y_mid = bs3.interpolate(0.0, 0.001, &dense, t_mid);

        // Should be close to e^0.0005
        let exact = (0.0005_f64).exp();
        // Cubic Hermite interpolation should be quite accurate
        assert_relative_eq!(y_mid[0], exact, max_relative = 1e-10);
    }

    #[test]
    fn test_bs3_accuracy() {
        // Test BS3 on a simple problem with known solution
        // y' = -y, y(0) = 1, solution is y(t) = e^(-t)
        type Params = ();
        fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
            Vector1::new(-y[0])
        }

        let y0 = Vector1::new(1.0);
        let bs3 = BS3::new().a_tol(1e-10).r_tol(1e-10);
        let h = 0.01;

        // Take 100 steps to reach t = 1.0
        let mut ode = ODE::new(&derivative, 0.0, 1.0, y0, ());
        for _ in 0..100 {
            let (y_new, _, _) = bs3.step(&ode, h);
            ode.y = y_new;
            ode.t += h;
        }

        // At t=1.0, exact solution is e^(-1) ≈ 0.36787944117
        let exact = (-1.0_f64).exp();
        assert_relative_eq!(ode.y[0], exact, max_relative = 1e-7);
    }

    #[test]
    fn test_bs3_convergence() {
        // Test that BS3 achieves 3rd order convergence
        // For a 3rd order method, halving h should reduce error by factor of ~2^3 = 8
        type Params = ();
        fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
            Vector1::new(y[0]) // y' = y, solution is e^t
        }

        let bs3 = BS3::new();
        let t_start = 0.0;
        let t_end = 1.0;
        let y0 = Vector1::new(1.0);

        // Test with different step sizes
        let step_sizes = [0.1, 0.05, 0.025];
        let mut errors = Vec::new();

        for &h in &step_sizes {
            let mut ode = ODE::new(&derivative, t_start, t_end, y0, ());

            // Take steps until we reach t_end
            while ode.t < t_end - 1e-10 {
                let (y_new, _, _) = bs3.step(&ode, h);
                ode.y = y_new;
                ode.t += h;
            }

            // Compute error at final time
            let exact = t_end.exp();
            let error = (ode.y[0] - exact).abs();
            errors.push(error);
        }

        // Check convergence rate between consecutive step sizes
        for i in 0..errors.len() - 1 {
            let ratio = errors[i] / errors[i + 1];
            // For order 3, we expect ratio ≈ 2^3 = 8 (since we halve the step size)
            // Allow some tolerance due to floating point arithmetic
            assert!(
                ratio > 6.0 && ratio < 10.0,
                "Expected convergence ratio ~8, got {:.2}",
                ratio
            );
        }

        // The error should decrease as step size decreases
        for i in 0..errors.len() - 1 {
            assert!(errors[i] > errors[i + 1]);
        }
    }

    #[test]
    fn test_bs3_fsal_property() {
        // Test that BS3 correctly implements the FSAL (First Same As Last) property
        // The last function evaluation of one step should equal the first of the next
        type Params = ();
        fn derivative(_t: f64, y: Vector1<f64>, _p: &Params) -> Vector1<f64> {
            Vector1::new(2.0 * y[0]) // y' = 2y
        }

        let y0 = Vector1::new(1.0);
        let bs3 = BS3::new();
        let h = 0.1;

        // First step
        let ode1 = ODE::new(&derivative, 0.0, 1.0, y0, ());
        let (y_new1, _, dense1) = bs3.step(&ode1, h);
        let dense1 = dense1.unwrap();

        // Extract f1 from first step (derivative at end of step)
        let f1_end = &dense1[3]; // f(t1, y1)

        // Second step starts where first ended
        let ode2 = ODE::new(&derivative, h, 1.0, y_new1, ());
        let (_, _, dense2) = bs3.step(&ode2, h);
        let dense2 = dense2.unwrap();

        // Extract f0 from second step (derivative at start of step)
        let f0_start = &dense2[2]; // f(t0, y0) of second step

        // These should be equal (FSAL property)
        assert_relative_eq!(f1_end[0], f0_start[0], max_relative = 1e-14);
    }
}