//! Physics-informed losses for CFD neural operators.

use aeroflow_shared::{FlowConditions, FlowField2D};

/// Navier-Stokes residual calculator for physics-informed losses.
#[derive(Debug)]
pub struct NavierStokesLoss {
    /// Kinematic viscosity (m²/s).
    nu: f32,
    /// Density (kg/m³).
    rho: f32,
}

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

impl NavierStokesLoss {
    /// Create a new Navier-Stokes loss calculator.
    pub fn new() -> Self {
        Self {
            nu: 1.5e-5, // Air at sea level
            rho: 1.225,
        }
    }

    /// Create with specific fluid properties.
    pub fn with_properties(nu: f32, rho: f32) -> Self {
        Self { nu, rho }
    }

    /// Evaluate physics residuals on the flow field.
    pub fn evaluate(&self, flow_field: &FlowField2D, conditions: &FlowConditions) -> f32 {
        let continuity = self.continuity_residual(flow_field);
        let momentum_x = self.momentum_x_residual(flow_field, conditions);
        let momentum_y = self.momentum_y_residual(flow_field, conditions);

        (continuity.powi(2) + momentum_x.powi(2) + momentum_y.powi(2)).sqrt()
    }

    /// Compute all residual components.
    pub fn compute_residuals(
        &self,
        flow_field: &FlowField2D,
        conditions: &FlowConditions,
    ) -> PhysicsResiduals {
        PhysicsResiduals {
            continuity: self.continuity_residual(flow_field),
            momentum_x: self.momentum_x_residual(flow_field, conditions),
            momentum_y: self.momentum_y_residual(flow_field, conditions),
            boundary: self.boundary_residual(flow_field, conditions),
        }
    }

    /// Continuity equation residual: ∂u/∂x + ∂v/∂y = 0
    fn continuity_residual(&self, flow_field: &FlowField2D) -> f32 {
        let nx = flow_field.velocity_x.len();
        let ny = flow_field.velocity_x[0].len();

        if nx < 3 || ny < 3 {
            return 0.0;
        }

        let dx = (flow_field.bounds.x_max - flow_field.bounds.x_min) / (nx - 1) as f32;
        let dy = (flow_field.bounds.y_max - flow_field.bounds.y_min) / (ny - 1) as f32;

        let mut total_residual = 0.0;
        let mut count = 0;

        for i in 1..nx - 1 {
            for j in 1..ny - 1 {
                let du_dx = (flow_field.velocity_x[i + 1][j] - flow_field.velocity_x[i - 1][j])
                    / (2.0 * dx);
                let dv_dy = (flow_field.velocity_y[i][j + 1] - flow_field.velocity_y[i][j - 1])
                    / (2.0 * dy);

                total_residual += (du_dx + dv_dy).abs();
                count += 1;
            }
        }

        if count > 0 {
            total_residual / count as f32
        } else {
            0.0
        }
    }

    /// X-momentum equation residual.
    /// u·∂u/∂x + v·∂u/∂y = -1/ρ·∂p/∂x + ν·∇²u
    fn momentum_x_residual(&self, flow_field: &FlowField2D, _conditions: &FlowConditions) -> f32 {
        let nx = flow_field.velocity_x.len();
        let ny = flow_field.velocity_x[0].len();

        if nx < 3 || ny < 3 {
            return 0.0;
        }

        let dx = (flow_field.bounds.x_max - flow_field.bounds.x_min) / (nx - 1) as f32;
        let dy = (flow_field.bounds.y_max - flow_field.bounds.y_min) / (ny - 1) as f32;

        let mut total_residual = 0.0;
        let mut count = 0;

        for i in 2..nx - 2 {
            for j in 2..ny - 2 {
                let u = flow_field.velocity_x[i][j];
                let v = flow_field.velocity_y[i][j];

                // First derivatives
                let du_dx = (flow_field.velocity_x[i + 1][j] - flow_field.velocity_x[i - 1][j])
                    / (2.0 * dx);
                let du_dy = (flow_field.velocity_x[i][j + 1] - flow_field.velocity_x[i][j - 1])
                    / (2.0 * dy);

                // Second derivatives (Laplacian)
                let d2u_dx2 = (flow_field.velocity_x[i + 1][j] - 2.0 * flow_field.velocity_x[i][j]
                    + flow_field.velocity_x[i - 1][j])
                    / (dx * dx);
                let d2u_dy2 = (flow_field.velocity_x[i][j + 1] - 2.0 * flow_field.velocity_x[i][j]
                    + flow_field.velocity_x[i][j - 1])
                    / (dy * dy);

                // Pressure gradient
                let dp_dx =
                    (flow_field.pressure[i + 1][j] - flow_field.pressure[i - 1][j]) / (2.0 * dx);

                // Momentum residual
                let convection = u * du_dx + v * du_dy;
                let pressure_term = dp_dx / self.rho;
                let diffusion = self.nu * (d2u_dx2 + d2u_dy2);

                let residual = convection + pressure_term - diffusion;
                total_residual += residual.abs();
                count += 1;
            }
        }

        if count > 0 {
            total_residual / count as f32
        } else {
            0.0
        }
    }

    /// Y-momentum equation residual.
    fn momentum_y_residual(&self, flow_field: &FlowField2D, _conditions: &FlowConditions) -> f32 {
        let nx = flow_field.velocity_y.len();
        let ny = flow_field.velocity_y[0].len();

        if nx < 3 || ny < 3 {
            return 0.0;
        }

        let dx = (flow_field.bounds.x_max - flow_field.bounds.x_min) / (nx - 1) as f32;
        let dy = (flow_field.bounds.y_max - flow_field.bounds.y_min) / (ny - 1) as f32;

        let mut total_residual = 0.0;
        let mut count = 0;

        for i in 2..nx - 2 {
            for j in 2..ny - 2 {
                let u = flow_field.velocity_x[i][j];
                let v = flow_field.velocity_y[i][j];

                // First derivatives
                let dv_dx = (flow_field.velocity_y[i + 1][j] - flow_field.velocity_y[i - 1][j])
                    / (2.0 * dx);
                let dv_dy = (flow_field.velocity_y[i][j + 1] - flow_field.velocity_y[i][j - 1])
                    / (2.0 * dy);

                // Second derivatives
                let d2v_dx2 = (flow_field.velocity_y[i + 1][j] - 2.0 * flow_field.velocity_y[i][j]
                    + flow_field.velocity_y[i - 1][j])
                    / (dx * dx);
                let d2v_dy2 = (flow_field.velocity_y[i][j + 1] - 2.0 * flow_field.velocity_y[i][j]
                    + flow_field.velocity_y[i][j - 1])
                    / (dy * dy);

                // Pressure gradient
                let dp_dy =
                    (flow_field.pressure[i][j + 1] - flow_field.pressure[i][j - 1]) / (2.0 * dy);

                // Momentum residual
                let convection = u * dv_dx + v * dv_dy;
                let pressure_term = dp_dy / self.rho;
                let diffusion = self.nu * (d2v_dx2 + d2v_dy2);

                let residual = convection + pressure_term - diffusion;
                total_residual += residual.abs();
                count += 1;
            }
        }

        if count > 0 {
            total_residual / count as f32
        } else {
            0.0
        }
    }

    /// Boundary condition residual.
    fn boundary_residual(&self, flow_field: &FlowField2D, conditions: &FlowConditions) -> f32 {
        let nx = flow_field.velocity_x.len();
        let ny = flow_field.velocity_x[0].len();

        let mut residual = 0.0;
        let mut count = 0;

        // Inlet boundary (left edge)
        for j in 0..ny {
            let u_inlet = conditions.velocity * conditions.angle_of_attack.to_radians().cos();
            let v_inlet = conditions.velocity * conditions.angle_of_attack.to_radians().sin();

            residual += (flow_field.velocity_x[0][j] - u_inlet).abs();
            residual += (flow_field.velocity_y[0][j] - v_inlet).abs();
            count += 2;
        }

        // Farfield boundaries (top and bottom)
        for i in 0..nx {
            // Top
            residual += (flow_field.velocity_x[i][ny - 1] - conditions.velocity).abs();
            // Bottom
            residual += (flow_field.velocity_x[i][0] - conditions.velocity).abs();
            count += 2;
        }

        if count > 0 {
            residual / count as f32
        } else {
            0.0
        }
    }
}

/// Physics residuals for monitoring.
#[derive(Debug, Clone, Copy, Default)]
pub struct PhysicsResiduals {
    /// Continuity equation residual.
    pub continuity: f32,
    /// X-momentum residual.
    pub momentum_x: f32,
    /// Y-momentum residual.
    pub momentum_y: f32,
    /// Boundary condition residual.
    pub boundary: f32,
}

impl PhysicsResiduals {
    /// Total L2 norm of residuals.
    pub fn total(&self) -> f32 {
        (self.continuity.powi(2)
            + self.momentum_x.powi(2)
            + self.momentum_y.powi(2)
            + self.boundary.powi(2))
        .sqrt()
    }
}

/// Euler equations residual (inviscid).
#[derive(Debug)]
pub struct EulerLoss {
    /// Ratio of specific heats.
    #[allow(dead_code)]
    gamma: f32,
}

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

impl EulerLoss {
    /// Create a new Euler loss calculator.
    pub fn new() -> Self {
        Self { gamma: 1.4 }
    }

    /// Evaluate Euler equation residuals.
    pub fn evaluate(&self, _flow_field: &FlowField2D, conditions: &FlowConditions) -> f32 {
        // For compressible flow, would compute:
        // ∂ρ/∂t + ∇·(ρu) = 0 (mass)
        // ∂(ρu)/∂t + ∇·(ρuu + pI) = 0 (momentum)
        // ∂E/∂t + ∇·((E + p)u) = 0 (energy)

        // Simplified: just check for low Mach incompressible behavior
        if conditions.mach < 0.3 {
            0.0
        } else {
            0.1 * (conditions.mach - 0.3)
        }
    }
}

/// Wall function for near-wall treatment.
#[derive(Debug)]
pub struct WallFunction {
    /// Von Karman constant.
    kappa: f32,
    /// Integration constant.
    b: f32,
}

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

impl WallFunction {
    /// Create a new wall function.
    pub fn new() -> Self {
        Self {
            kappa: 0.41,
            b: 5.0,
        }
    }

    /// Law of the wall: u+ = (1/κ)ln(y+) + B
    pub fn u_plus(&self, y_plus: f32) -> f32 {
        if y_plus < 11.6 {
            // Viscous sublayer
            y_plus
        } else {
            // Log layer
            (1.0 / self.kappa) * y_plus.ln() + self.b
        }
    }

    /// Compute wall shear stress.
    pub fn wall_shear(&self, u_tau: f32, rho: f32) -> f32 {
        rho * u_tau.powi(2)
    }
}

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

    fn create_test_flow_field() -> FlowField2D {
        let nx = 10;
        let ny = 8;

        FlowField2D {
            pressure: vec![vec![101325.0; ny]; nx],
            velocity_x: vec![vec![100.0; ny]; nx],
            velocity_y: vec![vec![0.0; ny]; nx],
            velocity_magnitude: vec![vec![100.0; ny]; nx],
            cp: vec![vec![0.0; ny]; nx],
            vorticity: vec![vec![0.0; ny]; nx],
            tke: None,
            grid_x: vec![vec![0.0; ny]; nx],
            grid_y: vec![vec![0.0; ny]; nx],
            bounds: DomainBounds::default(),
        }
    }

    #[test]
    fn test_ns_loss_creation() {
        let loss = NavierStokesLoss::new();
        assert!((loss.nu - 1.5e-5).abs() < 1e-7);
    }

    #[test]
    fn test_continuity_residual() {
        let loss = NavierStokesLoss::new();
        let flow_field = create_test_flow_field();
        let residual = loss.continuity_residual(&flow_field);
        // Uniform flow should have zero continuity residual
        assert!(residual.abs() < 1e-3);
    }

    #[test]
    fn test_evaluate() {
        let loss = NavierStokesLoss::new();
        let flow_field = create_test_flow_field();
        let conditions = FlowConditions::default();
        let total = loss.evaluate(&flow_field, &conditions);
        assert!(total.is_finite());
    }

    #[test]
    fn test_physics_residuals() {
        let loss = NavierStokesLoss::new();
        let flow_field = create_test_flow_field();
        let conditions = FlowConditions::default();
        let residuals = loss.compute_residuals(&flow_field, &conditions);

        assert!(residuals.continuity.is_finite());
        assert!(residuals.momentum_x.is_finite());
        assert!(residuals.momentum_y.is_finite());
        assert!(residuals.total().is_finite());
    }

    #[test]
    fn test_euler_loss() {
        let loss = EulerLoss::new();
        let flow_field = create_test_flow_field();

        let low_mach = FlowConditions {
            mach: 0.2,
            ..Default::default()
        };
        assert_eq!(loss.evaluate(&flow_field, &low_mach), 0.0);

        let high_mach = FlowConditions {
            mach: 0.5,
            ..Default::default()
        };
        assert!(loss.evaluate(&flow_field, &high_mach) > 0.0);
    }

    #[test]
    fn test_wall_function() {
        let wf = WallFunction::new();

        // Viscous sublayer: u+ = y+
        assert!((wf.u_plus(5.0) - 5.0).abs() < 0.01);

        // Log layer: u+ = (1/κ)ln(y+) + B
        let y_plus: f32 = 100.0;
        let expected = (1.0 / 0.41) * y_plus.ln() + 5.0;
        assert!((wf.u_plus(y_plus) - expected).abs() < 0.01);
    }
}