//! Boundary condition enforcement for hemodynamics PINN
//!
//! This module handles the enforcement of boundary conditions:
//! - Inlet: prescribed velocity profile
//! - Outlet: prescribed pressure
//! - Wall: no-slip condition (velocity = 0)

use crate::vessel::VesselSdf;
use rtx_hemodynamics_shared::geometry::Point2D;
use rtx_hemodynamics_shared::physics::{BoundaryCondition, BoundaryType};

/// Boundary condition enforcer for PINN training
#[derive(Debug)]
pub struct BoundaryEnforcer {
    /// Inlet boundary condition
    inlet: BoundaryCondition,
    /// Outlet boundary condition
    outlet: BoundaryCondition,
    /// Vessel SDF for wall detection
    vessel: VesselSdf,
    /// Tolerance for boundary detection
    tolerance: f64,
}

impl BoundaryEnforcer {
    /// Creates a new boundary enforcer
    #[must_use]
    pub fn new(inlet: BoundaryCondition, outlet: BoundaryCondition, vessel: VesselSdf) -> Self {
        Self {
            inlet,
            outlet,
            vessel,
            tolerance: 1e-6,
        }
    }

    /// Creates a boundary enforcer with Poiseuille inlet profile
    pub fn poiseuille_inlet(
        inlet_velocity: f64,
        outlet_pressure: f64,
        vessel: VesselSdf,
    ) -> Result<Self, String> {
        let inlet = BoundaryCondition::inlet_velocity(inlet_velocity).map_err(|e| e.to_string())?;
        let outlet =
            BoundaryCondition::outlet_pressure(outlet_pressure).map_err(|e| e.to_string())?;

        Ok(Self::new(inlet, outlet, vessel))
    }

    /// Returns the inlet boundary condition
    #[must_use]
    pub const fn inlet(&self) -> &BoundaryCondition {
        &self.inlet
    }

    /// Returns the outlet boundary condition
    #[must_use]
    pub const fn outlet(&self) -> &BoundaryCondition {
        &self.outlet
    }

    /// Returns the vessel SDF
    #[must_use]
    pub const fn vessel(&self) -> &VesselSdf {
        &self.vessel
    }

    /// Classifies a point as inlet, outlet, wall, or interior
    #[must_use]
    pub fn classify_point(&self, point: &Point2D) -> BoundaryLocation {
        let length = self.vessel.geometry().length();

        // Check if at inlet (x ≈ 0)
        if point.x.abs() < self.tolerance {
            return BoundaryLocation::Inlet;
        }

        // Check if at outlet (x ≈ length)
        if (point.x - length).abs() < self.tolerance {
            return BoundaryLocation::Outlet;
        }

        // Check if on wall
        if self.vessel.is_boundary(point, self.tolerance) {
            return BoundaryLocation::Wall;
        }

        // Check if outside
        if !self.vessel.is_inside(point) {
            return BoundaryLocation::Outside;
        }

        BoundaryLocation::Interior
    }

    /// Computes the target velocity at inlet using Poiseuille profile
    ///
    /// u(y) = `u_max` * (1 - (y/R)²)
    #[must_use]
    pub fn inlet_velocity_target(&self, y: f64, time: f64) -> (f64, f64) {
        let r = self.vessel.geometry().base_radius();

        // Get inlet velocity from boundary condition
        let u_max = match self.inlet.boundary_type() {
            BoundaryType::InletVelocity(v) => *v,
            BoundaryType::PulsatileInlet { .. } => self.inlet.evaluate(time),
            _ => 0.0,
        };

        // Parabolic profile: u(y) = u_max * (1 - (y/R)²)
        let normalized_y = y / r;
        let u = u_max * (1.0 - normalized_y.powi(2));

        (u, 0.0) // v = 0 at inlet for fully-developed flow
    }

    /// Computes the target pressure at outlet
    #[must_use]
    pub fn outlet_pressure_target(&self, _time: f64) -> f64 {
        match self.outlet.boundary_type() {
            BoundaryType::OutletPressure(p) => *p,
            _ => 0.0,
        }
    }

    /// Computes the boundary loss for a set of predicted values
    ///
    /// `L_bc` = `L_inlet` + `L_outlet` + `L_wall`
    #[must_use]
    pub fn boundary_loss(
        &self,
        inlet_points: &[Point2D],
        inlet_predictions: &[(f64, f64, f64)],
        _outlet_points: &[Point2D],
        outlet_predictions: &[(f64, f64, f64)],
        _wall_points: &[Point2D],
        wall_predictions: &[(f64, f64, f64)],
        time: f64,
    ) -> BoundaryLoss {
        let inlet_loss = self.inlet_loss(inlet_points, inlet_predictions, time);
        let outlet_loss = self.outlet_loss(outlet_predictions, time);
        let wall_loss = self.wall_loss(wall_predictions);

        BoundaryLoss {
            inlet: inlet_loss,
            outlet: outlet_loss,
            wall: wall_loss,
            total: inlet_loss + outlet_loss + wall_loss,
        }
    }

    /// Computes inlet loss (velocity MSE)
    fn inlet_loss(&self, points: &[Point2D], predictions: &[(f64, f64, f64)], time: f64) -> f64 {
        if points.is_empty() {
            return 0.0;
        }

        let mut loss = 0.0;
        for (point, (u_pred, v_pred, _)) in points.iter().zip(predictions.iter()) {
            let (u_target, v_target) = self.inlet_velocity_target(point.y, time);
            loss += (u_pred - u_target).powi(2) + (v_pred - v_target).powi(2);
        }

        loss / points.len() as f64
    }

    /// Computes outlet loss (pressure MSE)
    fn outlet_loss(&self, predictions: &[(f64, f64, f64)], time: f64) -> f64 {
        if predictions.is_empty() {
            return 0.0;
        }

        let p_target = self.outlet_pressure_target(time);
        let mut loss = 0.0;

        for (_, _, p_pred) in predictions {
            loss += (p_pred - p_target).powi(2);
        }

        loss / predictions.len() as f64
    }

    /// Computes wall loss (no-slip: u = v = 0)
    fn wall_loss(&self, predictions: &[(f64, f64, f64)]) -> f64 {
        if predictions.is_empty() {
            return 0.0;
        }

        let mut loss = 0.0;
        for (u, v, _) in predictions {
            loss += u.powi(2) + v.powi(2);
        }

        loss / predictions.len() as f64
    }

    /// Samples boundary points for training
    #[must_use]
    pub fn sample_boundary_points(
        &self,
        n_inlet: usize,
        n_outlet: usize,
        n_wall: usize,
    ) -> (Vec<Point2D>, Vec<Point2D>, Vec<Point2D>) {
        let inlet = self.vessel.sample_inlet(n_inlet);
        let outlet = self.vessel.sample_outlet(n_outlet);
        let wall = self.vessel.sample_boundary(n_wall);

        (inlet, outlet, wall)
    }
}

/// Classification of point location relative to boundaries
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BoundaryLocation {
    /// At inlet (x = 0)
    Inlet,
    /// At outlet (x = length)
    Outlet,
    /// On vessel wall
    Wall,
    /// Inside vessel (not on boundary)
    Interior,
    /// Outside vessel domain
    Outside,
}

/// Boundary loss components
#[derive(Debug, Clone, Copy)]
pub struct BoundaryLoss {
    /// Inlet velocity loss
    pub inlet: f64,
    /// Outlet pressure loss
    pub outlet: f64,
    /// Wall no-slip loss
    pub wall: f64,
    /// Total boundary loss
    pub total: f64,
}

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

    fn create_test_enforcer() -> BoundaryEnforcer {
        let vessel = VesselSdf::straight(0.1, 0.005).unwrap();
        BoundaryEnforcer::poiseuille_inlet(0.1, 0.0, vessel).unwrap()
    }

    #[test]
    fn test_boundary_enforcer_creation() {
        let enforcer = create_test_enforcer();
        assert!(enforcer.inlet().is_inlet());
        assert!(enforcer.outlet().is_outlet());
    }

    #[test]
    fn test_classify_point_inlet() {
        let enforcer = create_test_enforcer();
        let inlet_point = Point2D::new(0.0, 0.002);

        assert_eq!(
            enforcer.classify_point(&inlet_point),
            BoundaryLocation::Inlet
        );
    }

    #[test]
    fn test_classify_point_outlet() {
        let enforcer = create_test_enforcer();
        let outlet_point = Point2D::new(0.1, 0.002);

        assert_eq!(
            enforcer.classify_point(&outlet_point),
            BoundaryLocation::Outlet
        );
    }

    #[test]
    fn test_classify_point_interior() {
        let enforcer = create_test_enforcer();
        let interior_point = Point2D::new(0.05, 0.002);

        assert_eq!(
            enforcer.classify_point(&interior_point),
            BoundaryLocation::Interior
        );
    }

    #[test]
    fn test_inlet_velocity_profile() {
        let enforcer = create_test_enforcer();

        // Maximum at centerline
        let (u_center, v_center) = enforcer.inlet_velocity_target(0.0, 0.0);
        assert!((u_center - 0.1).abs() < f64::EPSILON);
        assert!(v_center.abs() < f64::EPSILON);

        // Zero at wall
        let (u_wall, _) = enforcer.inlet_velocity_target(0.005, 0.0);
        assert!(u_wall.abs() < f64::EPSILON);
    }

    #[test]
    fn test_wall_loss() {
        let enforcer = create_test_enforcer();

        // Perfect no-slip should have zero loss
        let perfect = vec![(0.0, 0.0, 100.0)];
        let loss = enforcer.wall_loss(&perfect);
        assert!(loss.abs() < f64::EPSILON);

        // Non-zero velocity should have positive loss
        let imperfect = vec![(0.01, 0.0, 100.0)];
        let loss = enforcer.wall_loss(&imperfect);
        assert!((loss - 0.0001).abs() < f64::EPSILON);
    }

    #[test]
    fn test_sample_boundary_points() {
        let enforcer = create_test_enforcer();

        let (inlet, outlet, wall) = enforcer.sample_boundary_points(10, 10, 20);

        assert_eq!(inlet.len(), 10);
        assert_eq!(outlet.len(), 10);
        assert_eq!(wall.len(), 20);

        // Inlet points should be at x=0
        for p in &inlet {
            assert!(p.x.abs() < f64::EPSILON);
        }

        // Outlet points should be at x=length
        for p in &outlet {
            assert!((p.x - 0.1).abs() < f64::EPSILON);
        }
    }

    #[test]
    fn test_boundary_loss_components() {
        let enforcer = create_test_enforcer();

        let inlet_points = vec![Point2D::new(0.0, 0.0)];
        let inlet_preds = vec![(0.1, 0.0, 50.0)]; // Perfect velocity

        let outlet_points = vec![Point2D::new(0.1, 0.0)];
        let outlet_preds = vec![(0.05, 0.0, 0.0)]; // Perfect pressure

        let wall_points = vec![Point2D::new(0.05, 0.005)];
        let wall_preds = vec![(0.0, 0.0, 50.0)]; // Perfect no-slip

        let loss = enforcer.boundary_loss(
            &inlet_points,
            &inlet_preds,
            &outlet_points,
            &outlet_preds,
            &wall_points,
            &wall_preds,
            0.0,
        );

        assert!(
            loss.inlet.abs() < f64::EPSILON,
            "inlet loss: {}",
            loss.inlet
        );
        assert!(
            loss.outlet.abs() < f64::EPSILON,
            "outlet loss: {}",
            loss.outlet
        );
        assert!(loss.wall.abs() < f64::EPSILON, "wall loss: {}", loss.wall);
    }
}