//! 2D Incompressible Navier-Stokes residual computation
//!
//! This module implements the physics-informed loss terms for the 2D
//! incompressible Navier-Stokes equations:
//!
//! **Momentum (x)**: ρ(∂u/∂t + u∂u/∂x + v∂u/∂y) = -∂p/∂x + μ(∂²u/∂x² + ∂²u/∂y²)
//! **Momentum (y)**: ρ(∂v/∂t + u∂v/∂x + v∂v/∂y) = -∂p/∂y + μ(∂²v/∂x² + ∂²v/∂y²)
//! **Continuity**: ∂u/∂x + ∂v/∂y = 0

use rtx_hemodynamics_shared::physics::FluidProperties;

/// Navier-Stokes residual computer for 2D incompressible flow
///
/// Computes the physics residuals for the momentum and continuity equations.
#[derive(Debug, Clone)]
pub struct NavierStokesResidual {
    /// Fluid density (kg/m³)
    rho: f64,
    /// Dynamic viscosity (Pa.s)
    mu: f64,
    /// Whether flow is steady-state (no time derivatives)
    steady_state: bool,
}

impl NavierStokesResidual {
    /// Creates a new Navier-Stokes residual computer
    #[must_use]
    pub fn new(fluid: &FluidProperties, steady_state: bool) -> Self {
        Self {
            rho: fluid.density(),
            mu: fluid.dynamic_viscosity(),
            steady_state,
        }
    }

    /// Creates residual computer for blood with steady-state assumption
    #[must_use]
    pub fn blood_steady() -> Self {
        Self::new(&FluidProperties::blood(), true)
    }

    /// Creates residual computer for blood with transient flow
    #[must_use]
    pub fn blood_transient() -> Self {
        Self::new(&FluidProperties::blood(), false)
    }

    /// Returns the fluid density
    #[must_use]
    pub const fn density(&self) -> f64 {
        self.rho
    }

    /// Returns the dynamic viscosity
    #[must_use]
    pub const fn viscosity(&self) -> f64 {
        self.mu
    }

    /// Returns whether using steady-state assumption
    #[must_use]
    pub const fn is_steady_state(&self) -> bool {
        self.steady_state
    }

    /// Computes the x-momentum residual at a single point
    ///
    /// `R_x` = ρ(∂u/∂t + u∂u/∂x + v∂u/∂y) + ∂p/∂x - μ(∂²u/∂x² + ∂²u/∂y²)
    ///
    /// For steady-state: `R_x` = ρ(u∂u/∂x + v∂u/∂y) + ∂p/∂x - μ(∂²u/∂x² + ∂²u/∂y²)
    ///
    /// # Arguments
    ///
    /// * `u` - x-velocity
    /// * `v` - y-velocity
    /// * `du_dt` - time derivative of u (ignored if steady-state)
    /// * `du_dx` - spatial derivative ∂u/∂x
    /// * `du_dy` - spatial derivative ∂u/∂y
    /// * `d2u_dx2` - second derivative ∂²u/∂x²
    /// * `d2u_dy2` - second derivative ∂²u/∂y²
    /// * `dp_dx` - pressure gradient ∂p/∂x
    #[must_use]
    #[allow(clippy::too_many_arguments)]
    pub fn x_momentum_residual(
        &self,
        u: f64,
        v: f64,
        du_dt: f64,
        du_dx: f64,
        du_dy: f64,
        d2u_dx2: f64,
        d2u_dy2: f64,
        dp_dx: f64,
    ) -> f64 {
        let time_term = if self.steady_state { 0.0 } else { du_dt };
        let convection = u * du_dx + v * du_dy;
        let diffusion = d2u_dx2 + d2u_dy2;

        self.rho * (time_term + convection) + dp_dx - self.mu * diffusion
    }

    /// Computes the y-momentum residual at a single point
    ///
    /// `R_y` = ρ(∂v/∂t + u∂v/∂x + v∂v/∂y) + ∂p/∂y - μ(∂²v/∂x² + ∂²v/∂y²)
    #[must_use]
    #[allow(clippy::too_many_arguments)]
    pub fn y_momentum_residual(
        &self,
        u: f64,
        v: f64,
        dv_dt: f64,
        dv_dx: f64,
        dv_dy: f64,
        d2v_dx2: f64,
        d2v_dy2: f64,
        dp_dy: f64,
    ) -> f64 {
        let time_term = if self.steady_state { 0.0 } else { dv_dt };
        let convection = u * dv_dx + v * dv_dy;
        let diffusion = d2v_dx2 + d2v_dy2;

        self.rho * (time_term + convection) + dp_dy - self.mu * diffusion
    }

    /// Computes the continuity residual at a single point
    ///
    /// `R_cont` = ∂u/∂x + ∂v/∂y
    ///
    /// For incompressible flow, this should be zero (divergence-free).
    #[must_use]
    pub fn continuity_residual(&self, du_dx: f64, dv_dy: f64) -> f64 {
        du_dx + dv_dy
    }

    /// Computes all three residuals at a single point
    ///
    /// Returns (`R_x`, `R_y`, `R_cont`)
    #[must_use]
    #[allow(clippy::too_many_arguments)]
    pub fn all_residuals(
        &self,
        u: f64,
        v: f64,
        du_dt: f64,
        dv_dt: f64,
        du_dx: f64,
        du_dy: f64,
        dv_dx: f64,
        dv_dy: f64,
        d2u_dx2: f64,
        d2u_dy2: f64,
        d2v_dx2: f64,
        d2v_dy2: f64,
        dp_dx: f64,
        dp_dy: f64,
    ) -> (f64, f64, f64) {
        let r_x = self.x_momentum_residual(u, v, du_dt, du_dx, du_dy, d2u_dx2, d2u_dy2, dp_dx);
        let r_y = self.y_momentum_residual(u, v, dv_dt, dv_dx, dv_dy, d2v_dx2, d2v_dy2, dp_dy);
        let r_cont = self.continuity_residual(du_dx, dv_dy);

        (r_x, r_y, r_cont)
    }

    /// Computes the total physics loss (MSE of all residuals)
    ///
    /// `L_physics` = `mean(R_x²` + `R_y²` + `R_cont²`)
    #[must_use]
    pub fn physics_loss(&self, residuals: &[(f64, f64, f64)]) -> f64 {
        if residuals.is_empty() {
            return 0.0;
        }

        let sum: f64 = residuals
            .iter()
            .map(|(rx, ry, rc)| rx.powi(2) + ry.powi(2) + rc.powi(2))
            .sum();

        sum / residuals.len() as f64
    }

    /// Computes individual loss components for monitoring
    ///
    /// Returns (`momentum_x_loss`, `momentum_y_loss`, `continuity_loss`)
    #[must_use]
    pub fn loss_components(&self, residuals: &[(f64, f64, f64)]) -> (f64, f64, f64) {
        if residuals.is_empty() {
            return (0.0, 0.0, 0.0);
        }

        let n = residuals.len() as f64;
        let (sum_x, sum_y, sum_c) = residuals.iter().fold((0.0, 0.0, 0.0), |acc, (rx, ry, rc)| {
            (acc.0 + rx.powi(2), acc.1 + ry.powi(2), acc.2 + rc.powi(2))
        });

        (sum_x / n, sum_y / n, sum_c / n)
    }
}

/// Analytical solution for 2D planar Poiseuille flow (steady channel flow)
///
/// This represents flow between two parallel plates (2D channel flow),
/// which exactly satisfies the 2D Cartesian Navier-Stokes equations.
///
/// For a channel of half-height h:
/// - Velocity profile: u(y) = (1/2μ)(-dp/dx)(h² - y²)
/// - Maximum velocity at centerline (y=0)
/// - Zero velocity at walls (y=±h)
#[derive(Debug, Clone)]
pub struct PoiseuilleFlow {
    /// Channel half-height (m) - distance from centerline to wall
    half_height: f64,
    /// Pressure gradient (Pa/m)
    dp_dx: f64,
    /// Dynamic viscosity (Pa.s)
    mu: f64,
}

impl PoiseuilleFlow {
    /// Creates a new 2D planar Poiseuille flow solution
    ///
    /// # Arguments
    ///
    /// * `half_height` - Channel half-height (distance from centerline to wall) in meters
    /// * `pressure_drop` - Total pressure drop over channel length
    /// * `length` - Channel length in meters
    /// * `mu` - Dynamic viscosity in Pa.s
    #[must_use]
    pub fn new(half_height: f64, pressure_drop: f64, length: f64, mu: f64) -> Self {
        Self {
            half_height,
            dp_dx: -pressure_drop / length, // Negative because pressure decreases
            mu,
        }
    }

    /// Creates Poiseuille flow with blood properties
    #[must_use]
    pub fn blood(half_height: f64, pressure_drop: f64, length: f64) -> Self {
        Self::new(
            half_height,
            pressure_drop,
            length,
            FluidProperties::blood().dynamic_viscosity(),
        )
    }

    /// Returns the channel half-height (also called radius for compatibility)
    #[must_use]
    pub const fn radius(&self) -> f64 {
        self.half_height
    }

    /// Returns the pressure gradient
    #[must_use]
    pub const fn pressure_gradient(&self) -> f64 {
        self.dp_dx
    }

    /// Computes the analytical velocity at position y from centerline
    ///
    /// u(y) = (1/2μ)(-dp/dx)(h² - y²)
    ///
    /// The velocity profile is parabolic with maximum at centerline.
    #[must_use]
    pub fn velocity(&self, y: f64) -> f64 {
        if y.abs() > self.half_height {
            return 0.0; // Outside channel
        }
        (-self.dp_dx / (2.0 * self.mu)) * (self.half_height.powi(2) - y.powi(2))
    }

    /// Returns the maximum (centerline) velocity
    #[must_use]
    pub fn max_velocity(&self) -> f64 {
        self.velocity(0.0)
    }

    /// Returns the average velocity
    ///
    /// `u_avg` = (2/3) * `u_max` for 2D channel flow
    #[must_use]
    pub fn average_velocity(&self) -> f64 {
        self.max_velocity() * 2.0 / 3.0
    }

    /// Computes the volumetric flow rate per unit depth
    ///
    /// Q = (2h³ / 3μ) * (-dp/dx)
    #[must_use]
    pub fn flow_rate(&self) -> f64 {
        2.0 * self.half_height.powi(3) * (-self.dp_dx) / (3.0 * self.mu)
    }

    /// Computes the wall shear stress
    ///
    /// `τ_w` = μ|du/dy|_wall = h * (-dp/dx)
    #[must_use]
    pub fn wall_shear_stress(&self) -> f64 {
        self.half_height * (-self.dp_dx)
    }

    /// Computes the velocity gradient at position y
    ///
    /// du/dy = (1/μ)(dp/dx) * y
    #[must_use]
    pub fn velocity_gradient(&self, y: f64) -> f64 {
        (self.dp_dx / self.mu) * y
    }

    /// Computes the pressure at axial position x (relative to inlet)
    ///
    /// p(x) = `p_inlet` + dp/dx * x
    #[must_use]
    pub fn pressure(&self, x: f64, p_inlet: f64) -> f64 {
        p_inlet + self.dp_dx * x
    }

    /// Validates that Navier-Stokes residuals are zero for this analytical solution
    ///
    /// For 2D planar Poiseuille flow:
    /// - u = u(y), v = 0
    /// - All time derivatives are zero
    /// - du/dx = 0 (fully developed)
    /// - dv/dx = dv/dy = 0
    /// - d²u/dx² = 0
    /// - d²u/dy² = (dp/dx) / μ
    #[must_use]
    pub fn validate_residuals(&self, ns: &NavierStokesResidual, y: f64) -> (f64, f64, f64) {
        let u = self.velocity(y);
        let v = 0.0;

        // Time derivatives (steady-state)
        let du_dt = 0.0;
        let dv_dt = 0.0;

        // Spatial derivatives
        let du_dx = 0.0; // Fully developed
        let du_dy = self.velocity_gradient(y);
        let dv_dx = 0.0;
        let dv_dy = 0.0;

        // Second derivatives
        let d2u_dx2 = 0.0;
        // d²u/dy² = (dp/dx) / μ for 2D planar flow
        let d2u_dy2 = self.dp_dx / self.mu;
        let d2v_dx2 = 0.0;
        let d2v_dy2 = 0.0;

        // Pressure gradients
        let dp_dx = self.dp_dx;
        let dp_dy = 0.0;

        ns.all_residuals(
            u, v, du_dt, dv_dt, du_dx, du_dy, dv_dx, dv_dy, d2u_dx2, d2u_dy2, d2v_dx2, d2v_dy2,
            dp_dx, dp_dy,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_ns_residual_creation() {
        let ns = NavierStokesResidual::blood_steady();
        assert!((ns.density() - 1060.0).abs() < 1.0);
        assert!((ns.viscosity() - 0.0035).abs() < 0.001);
        assert!(ns.is_steady_state());
    }

    #[test]
    fn test_continuity_residual() {
        let ns = NavierStokesResidual::blood_steady();

        // Divergence-free flow should have zero residual
        let residual = ns.continuity_residual(0.5, -0.5);
        assert!(residual.abs() < f64::EPSILON);

        // Non-divergence-free should have non-zero residual
        let residual = ns.continuity_residual(0.5, 0.3);
        assert!((residual - 0.8).abs() < f64::EPSILON);
    }

    #[test]
    fn test_poiseuille_velocity_profile() {
        // Create Poiseuille flow with known parameters
        let radius = 0.005; // 5mm pipe
        let pressure_drop = 100.0; // 100 Pa drop
        let length = 0.1; // 10cm length
        let mu = 0.0035; // Blood viscosity

        let flow = PoiseuilleFlow::new(radius, pressure_drop, length, mu);

        // Maximum velocity at centerline
        let u_max = flow.max_velocity();
        assert!(u_max > 0.0);

        // Zero velocity at wall
        let u_wall = flow.velocity(radius);
        assert!(u_wall.abs() < 1e-10);

        // Velocity at r=0 should be maximum
        let u_center = flow.velocity(0.0);
        assert!((u_center - u_max).abs() < f64::EPSILON);

        // Parabolic profile: u(R/2) = 3/4 * u_max
        let u_half = flow.velocity(radius / 2.0);
        assert!((u_half - 0.75 * u_max).abs() < 1e-10);
    }

    #[test]
    fn test_poiseuille_wall_shear_stress() {
        let half_height = 0.005;
        let pressure_drop = 100.0;
        let length = 0.1;
        let mu = 0.0035;

        let flow = PoiseuilleFlow::new(half_height, pressure_drop, length, mu);
        let wss = flow.wall_shear_stress();

        // For 2D planar flow: τ_w = h * |dp/dx|
        let expected = half_height * (pressure_drop / length);
        assert!((wss - expected).abs() < 1e-10);
    }

    #[test]
    fn test_poiseuille_validates_ns() {
        // The Poiseuille solution should satisfy Navier-Stokes exactly
        let radius = 0.005;
        let pressure_drop = 100.0;
        let length = 0.1;
        let mu = 0.0035;

        let flow = PoiseuilleFlow::new(radius, pressure_drop, length, mu);
        let ns = NavierStokesResidual::new(&FluidProperties::new(1060.0, mu).unwrap(), true);

        // Test at multiple radial positions
        for i in 0..10 {
            let r = radius * (i as f64) / 10.0;
            let (rx, ry, rc) = flow.validate_residuals(&ns, r);

            // All residuals should be near zero (within numerical precision)
            assert!(
                rx.abs() < 1e-6,
                "X-momentum residual too large at r={}: {}",
                r,
                rx
            );
            assert!(
                ry.abs() < 1e-10,
                "Y-momentum residual too large at r={}: {}",
                r,
                ry
            );
            assert!(
                rc.abs() < 1e-10,
                "Continuity residual too large at r={}: {}",
                r,
                rc
            );
        }
    }

    #[test]
    fn test_physics_loss_computation() {
        let ns = NavierStokesResidual::blood_steady();

        // Perfect solution should have zero loss
        let perfect_residuals = vec![(0.0, 0.0, 0.0), (0.0, 0.0, 0.0)];
        let loss = ns.physics_loss(&perfect_residuals);
        assert!(loss.abs() < f64::EPSILON);

        // Non-zero residuals should give positive loss
        let residuals = vec![(1.0, 0.0, 0.0), (0.0, 1.0, 0.0)];
        let loss = ns.physics_loss(&residuals);
        assert!((loss - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_loss_components() {
        let ns = NavierStokesResidual::blood_steady();

        let residuals = vec![(1.0, 2.0, 3.0), (1.0, 2.0, 3.0)];

        let (lx, ly, lc) = ns.loss_components(&residuals);
        assert!((lx - 1.0).abs() < f64::EPSILON);
        assert!((ly - 4.0).abs() < f64::EPSILON);
        assert!((lc - 9.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_poiseuille_flow_rate() {
        let half_height = 0.005;
        let pressure_drop = 100.0;
        let length = 0.1;
        let mu = 0.0035;

        let flow = PoiseuilleFlow::new(half_height, pressure_drop, length, mu);

        // 2D planar flow rate per unit depth: Q = (2h³ / 3μ) * |dp/dx|
        let dp_dx = pressure_drop / length;
        let expected_q = 2.0 * half_height.powi(3) * dp_dx / (3.0 * mu);
        let actual_q = flow.flow_rate();

        assert!((actual_q - expected_q).abs() < 1e-15);
    }

    #[test]
    fn test_steady_vs_transient() {
        let fluid = FluidProperties::blood();
        let steady = NavierStokesResidual::new(&fluid, true);
        let transient = NavierStokesResidual::new(&fluid, false);

        // With non-zero du/dt, steady should ignore it
        let r_steady = steady.x_momentum_residual(
            1.0, 0.0,  // u, v
            10.0, // du_dt (should be ignored)
            0.0, 0.0, // du_dx, du_dy
            0.0, 0.0, // d2u_dx2, d2u_dy2
            0.0, // dp_dx
        );

        let r_transient = transient.x_momentum_residual(
            1.0, 0.0,  // u, v
            10.0, // du_dt (should be included)
            0.0, 0.0, // du_dx, du_dy
            0.0, 0.0, // d2u_dx2, d2u_dy2
            0.0, // dp_dx
        );

        // Steady should have zero (no time term)
        assert!(r_steady.abs() < f64::EPSILON);

        // Transient should have ρ * du/dt = 1060 * 10 = 10600
        assert!((r_transient - 1060.0 * 10.0).abs() < 0.1);
    }
}