//! ODE Solvers for Liquid Neural Networks
//!
//! Implements Euler and Runge-Kutta 4th order solvers for continuous-time
//! neural dynamics with fixed-point arithmetic for ultra-low latency.

use super::activation::{self};
use super::{FixedPoint, MarketRegime, Result, PRECISION};
use serde::{Deserialize, Serialize};

/// ODE solver types for different accuracy/speed tradeoffs
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum SolverType {
    Euler,    // Fast, first-order accuracy
    RK4,      // Slower, fourth-order accuracy
    Adaptive, // Dynamic solver selection based on market regime
}

/// Function type for ODE right-hand side - now accepts closures
pub type ODEFunction<'a> = dyn Fn(FixedPoint, FixedPoint) -> FixedPoint + 'a;

/// Trait for ODE solvers
pub trait ODESolver: std::fmt::Debug + Clone {
    fn step<'a>(
        &self,
        f: &'a ODEFunction<'a>,
        x: FixedPoint,
        t: FixedPoint,
        dt: FixedPoint,
    ) -> Result<FixedPoint>;
    fn solver_type(&self) -> SolverType;
}

/// Euler method solver - fast but less accurate
#[derive(Debug, Clone)]
pub struct EulerSolver;

impl ODESolver for EulerSolver {
    fn step<'a>(
        &self,
        f: &'a ODEFunction<'a>,
        x: FixedPoint,
        t: FixedPoint,
        dt: FixedPoint,
    ) -> Result<FixedPoint> {
        let dx_dt = f(x, t);
        let step = (dt * dx_dt)?;
        x + step
    }

    fn solver_type(&self) -> SolverType {
        SolverType::Euler
    }
}

/// Runge-Kutta 4th order solver - more accurate but slower
#[derive(Debug, Clone)]
pub struct RK4Solver;

impl ODESolver for RK4Solver {
    fn step<'a>(
        &self,
        f: &'a ODEFunction<'a>,
        x: FixedPoint,
        t: FixedPoint,
        dt: FixedPoint,
    ) -> Result<FixedPoint> {
        let k1 = f(x, t);

        let half_dt = FixedPoint(dt.0 / 2);
        let k1_step = (half_dt * k1)?;
        let x_k1 = (x + k1_step)?;
        let t_half = (t + half_dt)?;
        let k2 = f(x_k1, t_half);

        let k2_step = (half_dt * k2)?;
        let x_k2 = (x + k2_step)?;
        let k3 = f(x_k2, t_half);

        let k3_step = (dt * k3)?;
        let x_k3 = (x + k3_step)?;
        let t_full = (t + dt)?;
        let k4 = f(x_k3, t_full);

        // Combine: x + dt/6 * (k1 + 2*k2 + 2*k3 + k4)
        let two_k2 = FixedPoint(k2.0 * 2);
        let two_k3 = FixedPoint(k3.0 * 2);
        let sum_partial = (k1 + two_k2)?;
        let sum_partial2 = (sum_partial + two_k3)?;
        let sum = (sum_partial2 + k4)?;
        let sixth_dt = FixedPoint(dt.0 / 6);
        let increment = (sixth_dt * sum)?;

        x + increment
    }

    fn solver_type(&self) -> SolverType {
        SolverType::RK4
    }
}

/// Adaptive solver that chooses method based on market conditions
#[derive(Debug, Clone)]
pub struct AdaptiveSolver {
    euler: EulerSolver,
    rk4: RK4Solver,
    current_regime: MarketRegime,
}

impl Default for AdaptiveSolver {
    fn default() -> Self {
        Self::new()
    }
}

impl AdaptiveSolver {
    pub const fn new() -> Self {
        Self {
            euler: EulerSolver,
            rk4: RK4Solver,
            current_regime: MarketRegime::Normal,
        }
    }

    pub const fn update_regime(&mut self, regime: MarketRegime) {
        self.current_regime = regime;
    }

    const fn use_high_accuracy(&self) -> bool {
        matches!(
            self.current_regime,
            MarketRegime::Bull | MarketRegime::Bear | MarketRegime::Crisis
        )
    }
}

impl ODESolver for AdaptiveSolver {
    fn step<'a>(
        &self,
        f: &'a ODEFunction<'a>,
        x: FixedPoint,
        t: FixedPoint,
        dt: FixedPoint,
    ) -> Result<FixedPoint> {
        if self.use_high_accuracy() {
            self.rk4.step(f, x, t, dt)
        } else {
            self.euler.step(f, x, t, dt)
        }
    }

    fn solver_type(&self) -> SolverType {
        SolverType::Adaptive
    }
}

/// Volatility-aware time constant adaptation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VolatilityAwareTimeConstants {
    base_tau: FixedPoint,
    min_tau: FixedPoint,
    max_tau: FixedPoint,
    current_tau: FixedPoint,
    volatility_factor: FixedPoint,
    adaptation_rate: FixedPoint,
}

impl VolatilityAwareTimeConstants {
    pub const fn new(base_tau: FixedPoint, min_tau: FixedPoint, max_tau: FixedPoint) -> Self {
        Self {
            base_tau,
            min_tau,
            max_tau,
            current_tau: base_tau,
            volatility_factor: FixedPoint::one(),
            adaptation_rate: FixedPoint(PRECISION / 10), // 0.1
        }
    }

    pub fn update_volatility(&mut self, volatility: FixedPoint) -> Result<()> {
        // Map volatility to time constant adjustment
        // High volatility -> lower time constants (faster adaptation)
        let volatility_normalized = if volatility.0 > 10 * PRECISION {
            FixedPoint(10 * PRECISION) // Cap at 10.0
        } else {
            volatility
        };

        // Inverse relationship: tau = base_tau / (1 + volatility_factor)
        let one = FixedPoint::one();
        let vol_factor = (one + volatility_normalized)?;
        let new_tau = (self.base_tau / vol_factor)?;

        // Clamp to valid range
        self.current_tau = if new_tau.0 < self.min_tau.0 {
            self.min_tau
        } else if new_tau.0 > self.max_tau.0 {
            self.max_tau
        } else {
            new_tau
        };

        self.volatility_factor = volatility_normalized;
        Ok(())
    }

    pub const fn current_tau(&self) -> FixedPoint {
        self.current_tau
    }
}

/// Liquid neural dynamics functions
#[derive(Debug)]
pub struct LiquidDynamics;

impl LiquidDynamics {
    /// LTC cell dynamics: dx/dt = (1/tau) * (-x + f(W*input + b))
    pub fn ltc_dynamics(
        weight: FixedPoint,
        bias: FixedPoint,
        tau: FixedPoint,
        activation_fn: &dyn Fn(FixedPoint) -> Result<FixedPoint>,
    ) -> impl Fn(FixedPoint, FixedPoint) -> FixedPoint + use<'_> {
        let w = weight;
        let b = bias;
        let t = tau;

        move |x: FixedPoint, _time: FixedPoint| -> FixedPoint {
            // Compute input: w*x + b (simplified for single neuron)
            let input = match w * x {
                Ok(wx) => match wx + b {
                    Ok(input) => input,
                    Err(_) => return FixedPoint::zero(), // Handle overflow
                },
                Err(_) => return FixedPoint::zero(),
            };

            // Apply activation function
            let Ok(activated) = activation_fn(input) else {
                return FixedPoint::zero();
            };

            // Compute dynamics: (1/tau) * (-x + activated)
            let neg_x = FixedPoint(-x.0);
            let Ok(diff) = neg_x + activated else {
                return FixedPoint::zero();
            };

            match diff / t {
                Ok(result) => result,
                Err(_) => FixedPoint::zero(),
            }
        }
    }

    /// `CfC` dynamics with closed-form solution approximation
    pub fn cfc_dynamics<'a>(
        input_weights: &'a [FixedPoint],
        recurrent_weights: &'a [FixedPoint],
        bias: FixedPoint,
    ) -> impl Fn(FixedPoint, FixedPoint) -> FixedPoint + 'a {
        move |x: FixedPoint, _time: FixedPoint| -> FixedPoint {
            // Simplified CfC dynamics for single cell
            // In practice, this would involve matrix operations
            let input_contrib = input_weights.first().copied().unwrap_or(FixedPoint::zero());
            let recurrent_contrib = recurrent_weights.first()
                .copied()
                .unwrap_or(FixedPoint::zero());

            let Some(total_input) = ((input_contrib * x).ok())
                .zip((recurrent_contrib * x).ok())
                .and_then(|(a, b)| (a + b).ok())
                .and_then(|sum| (sum + bias).ok())
            else {
                return FixedPoint::zero();
            };

            // Apply simple nonlinearity (tanh approximation)
            match activation::tanh(total_input) {
                Ok(result) => result,
                Err(_) => FixedPoint::zero(),
            }
        }
    }
}

/// Enum for different ODE solver types
#[derive(Debug, Clone)]
pub enum SolverEnum {
    Euler(EulerSolver),
    RK4(RK4Solver),
    Adaptive(AdaptiveSolver),
}

impl SolverEnum {
    pub fn step<'a>(
        &self,
        f: &'a ODEFunction<'a>,
        x: FixedPoint,
        t: FixedPoint,
        dt: FixedPoint,
    ) -> Result<FixedPoint> {
        match self {
            SolverEnum::Euler(solver) => solver.step(f, x, t, dt),
            SolverEnum::RK4(solver) => solver.step(f, x, t, dt),
            SolverEnum::Adaptive(solver) => solver.step(f, x, t, dt),
        }
    }

    pub fn solver_type(&self) -> SolverType {
        match self {
            SolverEnum::Euler(solver) => solver.solver_type(),
            SolverEnum::RK4(solver) => solver.solver_type(),
            SolverEnum::Adaptive(solver) => solver.solver_type(),
        }
    }
}

/// Factory for creating solvers
#[derive(Debug)]
pub struct SolverFactory;

impl SolverFactory {
    pub const fn create_solver(solver_type: SolverType) -> SolverEnum {
        match solver_type {
            SolverType::Euler => SolverEnum::Euler(EulerSolver),
            SolverType::RK4 => SolverEnum::RK4(RK4Solver),
            SolverType::Adaptive => SolverEnum::Adaptive(AdaptiveSolver::new()),
        }
    }
}

#[cfg(test)]
#[allow(clippy::unnecessary_wraps)]
mod tests {
    use super::*;
    use crate::liquid::PRECISION;
    use common::trading::MarketRegime;


    #[test]
    fn test_euler_solver() -> Result<()> {
        let solver = EulerSolver;
        let dt = FixedPoint(PRECISION / 100); // 0.01

        // Simple linear ODE: dx/dt = -x (exponential decay)
        let linear_decay = |x: FixedPoint, _t: FixedPoint| -> FixedPoint { FixedPoint(-x.0) };

        let x0 = FixedPoint(PRECISION); // 1.0
        let t0 = FixedPoint(0);

        let x1 = solver.step(&linear_decay, x0, t0, dt)?;

        // Expected: x1 = x0 + dt * (-x0) = 1.0 - 0.01 = 0.99
        let expected = FixedPoint((0.99 * PRECISION as f64) as i64);
        assert!((x1.0 - expected.0).abs() < PRECISION / 100); // Within 1% tolerance
        Ok(())
    }

    #[test]
    fn test_rk4_solver() -> Result<()> {
        let solver = RK4Solver;
        let dt = FixedPoint(PRECISION / 100); // 0.01

        // Linear ODE: dx/dt = -x
        let linear_decay = |x: FixedPoint, _t: FixedPoint| -> FixedPoint { FixedPoint(-x.0) };

        let x0 = FixedPoint(PRECISION); // 1.0
        let t0 = FixedPoint(0);

        let x1 = solver.step(&linear_decay, x0, t0, dt)?;

        // RK4 should be more accurate than Euler for this problem
        // Analytical solution: x(t) = exp(-t), so x(0.01) ≈ 0.9900498
        let expected = FixedPoint((0.990049 * PRECISION as f64) as i64);
        assert!((x1.0 - expected.0).abs() < PRECISION / 1000); // Higher accuracy expected
        Ok(())
    }

    #[test]
    fn test_volatility_aware_time_constants() -> Result<()> {
        let base_tau = FixedPoint(PRECISION / 10); // 0.1
        let min_tau = FixedPoint(PRECISION / 100); // 0.01
        let max_tau = FixedPoint(PRECISION); // 1.0

        let mut vol_aware = VolatilityAwareTimeConstants::new(base_tau, min_tau, max_tau);

        // High volatility should decrease time constant
        let high_volatility = FixedPoint(5 * PRECISION); // 5.0
        vol_aware.update_volatility(high_volatility)?;

        let adapted_tau = vol_aware.current_tau();
        assert!(adapted_tau.0 <= base_tau.0); // Should be smaller or equal
        assert!(adapted_tau.0 >= min_tau.0); // Should respect minimum
        Ok(())
    }

    #[test]
    fn test_ltc_dynamics() -> Result<()> {
        let weight = FixedPoint(PRECISION / 2); // 0.5
        let bias = FixedPoint(PRECISION / 10); // 0.1
        let tau = FixedPoint(PRECISION / 10); // 0.1

        let dynamics = LiquidDynamics::ltc_dynamics(weight, bias, tau, &activation::sigmoid);

        let x = FixedPoint(PRECISION / 2); // 0.5
        let t = FixedPoint(0);

        let dx_dt = dynamics(x, t);

        // Should be finite and reasonable
        assert!(dx_dt.0.abs() < 10 * PRECISION);
        Ok(())
    }

    #[test]
    fn test_adaptive_solver() -> Result<()> {
        let mut solver = AdaptiveSolver::new();

        // Test with normal regime (should use Euler)
        solver.update_regime(MarketRegime::Normal);
        assert_eq!(solver.solver_type(), SolverType::Adaptive);

        // Test with crisis regime (should use RK4)
        solver.update_regime(MarketRegime::Crisis);
        assert!(solver.use_high_accuracy());
        Ok(())
    }
}
