//! Optimized inference engine for hemodynamics PINN
//!
//! Provides efficient batched inference with optional GPU acceleration.

use crate::network::VesselPinn;
use crate::vessel::VesselSdf;
use crate::wss::WssComputer;
use rtx_hemodynamics_shared::fields::{FieldResponse, PressureField, VelocityField, WssField};
use rtx_hemodynamics_shared::geometry::Point2D;

/// Inference engine for field evaluation
#[derive(Debug)]
pub struct InferenceEngine {
    /// The trained PINN model
    model: VesselPinn,
    /// Vessel geometry
    vessel: VesselSdf,
    /// WSS computer
    wss_computer: WssComputer,
    /// Inference batch size
    batch_size: usize,
}

impl InferenceEngine {
    /// Creates a new inference engine
    #[must_use]
    pub fn new(model: VesselPinn, vessel: VesselSdf) -> Self {
        Self {
            model,
            vessel,
            wss_computer: WssComputer::blood(),
            batch_size: 1000,
        }
    }

    /// Sets the batch size for inference
    #[must_use]
    pub const fn with_batch_size(mut self, size: usize) -> Self {
        self.batch_size = size;
        self
    }

    /// Returns the model
    #[must_use]
    pub const fn model(&self) -> &VesselPinn {
        &self.model
    }

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

    /// Evaluates velocity and pressure at given points
    #[must_use]
    pub fn evaluate(&self, points: &[Point2D], time: f64) -> Vec<(f64, f64, f64)> {
        self.model.forward_batch(points, time)
    }

    /// Evaluates and returns full field response
    #[must_use]
    pub fn evaluate_fields(&self, points: &[Point2D], time: f64) -> FieldResponse {
        let predictions = self.evaluate(points, time);

        let u: Vec<f64> = predictions.iter().map(|(u, _, _)| *u).collect();
        let v: Vec<f64> = predictions.iter().map(|(_, v, _)| *v).collect();
        let p: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();

        let velocity = VelocityField::new(points.to_vec(), u, v).unwrap();
        let pressure = PressureField::new(points.to_vec(), p).unwrap();

        FieldResponse::new(velocity, pressure, None)
    }

    /// Evaluates fields including WSS at wall points
    #[must_use]
    pub fn evaluate_with_wss(
        &self,
        interior_points: &[Point2D],
        wall_points: &[Point2D],
        time: f64,
    ) -> FieldResponse {
        // Evaluate interior fields
        let interior_preds = self.evaluate(interior_points, time);

        let u: Vec<f64> = interior_preds.iter().map(|(u, _, _)| *u).collect();
        let v: Vec<f64> = interior_preds.iter().map(|(_, v, _)| *v).collect();
        let p: Vec<f64> = interior_preds.iter().map(|(_, _, p)| *p).collect();

        let velocity = VelocityField::new(interior_points.to_vec(), u, v).unwrap();
        let pressure = PressureField::new(interior_points.to_vec(), p).unwrap();

        // Compute WSS at wall points
        let wss_values = self.compute_wss(wall_points, time);
        let wss = WssField::new(wall_points.to_vec(), wss_values).unwrap();

        FieldResponse::new(velocity, pressure, Some(wss))
    }

    /// Computes WSS at wall points
    #[must_use]
    pub fn compute_wss(&self, wall_points: &[Point2D], time: f64) -> Vec<f64> {
        let velocity_fn = |p: &Point2D| {
            let (u, v, _) = self.model.forward(p.x, p.y, time);
            (u, v)
        };

        self.wss_computer
            .compute_wss_batch(wall_points, &self.vessel, velocity_fn)
    }

    /// Evaluates on a regular grid for visualization
    #[must_use]
    pub fn evaluate_grid(&self, nx: usize, ny: usize, time: f64) -> GridFieldData {
        let length = self.vessel.geometry().length();
        let (bbox_min, bbox_max) = self.vessel.geometry().bounding_box();

        let dx = length / (nx - 1) as f64;
        let dy = (bbox_max.y - bbox_min.y) / (ny - 1) as f64;

        let mut points = Vec::with_capacity(nx * ny);
        let mut mask = Vec::with_capacity(nx * ny);

        for j in 0..ny {
            for i in 0..nx {
                let x = i as f64 * dx;
                let y = bbox_min.y + j as f64 * dy;
                let p = Point2D::new(x, y);

                points.push(p);
                mask.push(self.vessel.is_inside(&p));
            }
        }

        let predictions = self.evaluate(&points, time);

        let u: Vec<f64> = predictions.iter().map(|(u, _, _)| *u).collect();
        let v: Vec<f64> = predictions.iter().map(|(_, v, _)| *v).collect();
        let p: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();

        GridFieldData {
            nx,
            ny,
            dx,
            dy,
            x_min: 0.0,
            y_min: bbox_min.y,
            u,
            v,
            p,
            mask,
        }
    }

    /// Samples streamlines for visualization
    #[must_use]
    pub fn compute_streamlines(
        &self,
        seed_points: &[Point2D],
        time: f64,
        max_steps: usize,
        step_size: f64,
    ) -> Vec<Vec<Point2D>> {
        seed_points
            .iter()
            .map(|seed| self.trace_streamline(*seed, time, max_steps, step_size))
            .collect()
    }

    /// Traces a single streamline from a seed point
    fn trace_streamline(
        &self,
        seed: Point2D,
        time: f64,
        max_steps: usize,
        step_size: f64,
    ) -> Vec<Point2D> {
        let mut streamline = vec![seed];
        let mut current = seed;

        for _ in 0..max_steps {
            // Get velocity at current point
            let (u, v, _) = self.model.forward(current.x, current.y, time);
            let vel_mag = (u * u + v * v).sqrt();

            if vel_mag < 1e-10 {
                break; // Stagnation point
            }

            // Normalize and step
            let next = Point2D::new(
                current.x + step_size * u / vel_mag,
                current.y + step_size * v / vel_mag,
            );

            // Check if still inside vessel
            if !self.vessel.is_inside(&next) {
                break;
            }

            streamline.push(next);
            current = next;
        }

        streamline
    }

    /// Computes flow statistics for the current solution
    #[must_use]
    pub fn compute_statistics(&self, time: f64) -> FlowStatistics {
        let interior = self.vessel.sample_interior(1000, 42);
        let wall = self.vessel.sample_boundary(200);

        let predictions = self.evaluate(&interior, time);
        let wss_values = self.compute_wss(&wall, time);

        let velocities: Vec<f64> = predictions
            .iter()
            .map(|(u, v, _)| (u * u + v * v).sqrt())
            .collect();
        let pressures: Vec<f64> = predictions.iter().map(|(_, _, p)| *p).collect();

        FlowStatistics {
            max_velocity: velocities.iter().copied().fold(0.0_f64, f64::max),
            mean_velocity: velocities.iter().sum::<f64>() / velocities.len() as f64,
            max_pressure: pressures.iter().copied().fold(f64::NEG_INFINITY, f64::max),
            min_pressure: pressures.iter().copied().fold(f64::INFINITY, f64::min),
            pressure_drop: pressures[0] - pressures[pressures.len() - 1],
            max_wss: WssComputer::max_wss(&wss_values),
            mean_wss: WssComputer::mean_wss(&wss_values),
        }
    }
}

/// Grid-based field data for visualization
#[derive(Debug, Clone)]
pub struct GridFieldData {
    /// Number of grid points in x
    pub nx: usize,
    /// Number of grid points in y
    pub ny: usize,
    /// Grid spacing in x
    pub dx: f64,
    /// Grid spacing in y
    pub dy: f64,
    /// Minimum x coordinate
    pub x_min: f64,
    /// Minimum y coordinate
    pub y_min: f64,
    /// X-velocity component (row-major: [j * nx + i])
    pub u: Vec<f64>,
    /// Y-velocity component
    pub v: Vec<f64>,
    /// Pressure
    pub p: Vec<f64>,
    /// Interior mask (true if inside vessel)
    pub mask: Vec<bool>,
}

impl GridFieldData {
    /// Returns the value at grid index (i, j)
    ///
    /// Returns None if outside grid or masked.
    #[must_use]
    pub fn get(&self, i: usize, j: usize) -> Option<(f64, f64, f64)> {
        if i >= self.nx || j >= self.ny {
            return None;
        }

        let idx = j * self.nx + i;
        if !self.mask[idx] {
            return None;
        }

        Some((self.u[idx], self.v[idx], self.p[idx]))
    }

    /// Converts grid index to physical coordinates
    #[must_use]
    pub fn grid_to_physical(&self, i: usize, j: usize) -> Point2D {
        Point2D::new(
            self.x_min + i as f64 * self.dx,
            self.y_min + j as f64 * self.dy,
        )
    }
}

/// Flow statistics summary
#[derive(Debug, Clone, Copy)]
pub struct FlowStatistics {
    /// Maximum velocity magnitude
    pub max_velocity: f64,
    /// Mean velocity magnitude
    pub mean_velocity: f64,
    /// Maximum pressure
    pub max_pressure: f64,
    /// Minimum pressure
    pub min_pressure: f64,
    /// Pressure drop across vessel
    pub pressure_drop: f64,
    /// Maximum wall shear stress
    pub max_wss: f64,
    /// Mean wall shear stress
    pub mean_wss: f64,
}

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

    fn create_test_engine() -> InferenceEngine {
        let config = PinnConfig::default().with_layers(2).with_hidden_dim(16);
        let model = VesselPinn::new(config);
        let vessel = VesselSdf::straight(0.1, 0.005).unwrap();

        InferenceEngine::new(model, vessel)
    }

    #[test]
    fn test_inference_engine_creation() {
        let engine = create_test_engine();
        assert!(engine.batch_size > 0);
    }

    #[test]
    fn test_evaluate_points() {
        let engine = create_test_engine();
        let points = vec![Point2D::new(0.01, 0.0), Point2D::new(0.05, 0.002)];

        let results = engine.evaluate(&points, 0.0);

        assert_eq!(results.len(), 2);
        for (u, v, p) in &results {
            assert!(u.is_finite());
            assert!(v.is_finite());
            assert!(p.is_finite());
        }
    }

    #[test]
    fn test_evaluate_fields() {
        let engine = create_test_engine();
        let points = engine.vessel.sample_interior(50, 42);

        let response = engine.evaluate_fields(&points, 0.0);

        assert_eq!(response.velocity().len(), 50);
        assert_eq!(response.pressure().len(), 50);
        assert!(!response.has_wss());
    }

    #[test]
    fn test_evaluate_with_wss() {
        let engine = create_test_engine();
        let interior = engine.vessel.sample_interior(50, 42);
        let wall = engine.vessel.sample_boundary(20);

        let response = engine.evaluate_with_wss(&interior, &wall, 0.0);

        assert_eq!(response.velocity().len(), 50);
        assert!(response.has_wss());
        assert_eq!(response.wss().as_ref().unwrap().len(), 20);
    }

    #[test]
    fn test_evaluate_grid() {
        let engine = create_test_engine();
        let grid = engine.evaluate_grid(10, 5, 0.0);

        assert_eq!(grid.nx, 10);
        assert_eq!(grid.ny, 5);
        assert_eq!(grid.u.len(), 50);
        assert_eq!(grid.mask.len(), 50);
    }

    #[test]
    fn test_streamline_tracing() {
        let engine = create_test_engine();
        let seeds = vec![Point2D::new(0.01, 0.0)];

        let streamlines = engine.compute_streamlines(&seeds, 0.0, 100, 0.001);

        assert_eq!(streamlines.len(), 1);
        assert!(!streamlines[0].is_empty());

        // First point should be seed
        let first = &streamlines[0][0];
        assert!((first.x - 0.01).abs() < f64::EPSILON);
    }

    #[test]
    fn test_compute_statistics() {
        let engine = create_test_engine();
        let stats = engine.compute_statistics(0.0);

        assert!(stats.max_velocity >= 0.0);
        assert!(stats.mean_velocity >= 0.0);
        assert!(stats.max_wss >= 0.0);
    }

    #[test]
    fn test_grid_field_data_access() {
        let grid = GridFieldData {
            nx: 3,
            ny: 2,
            dx: 0.05,
            dy: 0.005,
            x_min: 0.0,
            y_min: -0.005,
            u: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            v: vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
            p: vec![100.0, 90.0, 80.0, 100.0, 90.0, 80.0],
            mask: vec![true, true, true, true, true, true],
        };

        let (u, v, p) = grid.get(1, 0).unwrap();
        assert!((u - 2.0).abs() < f64::EPSILON);
        assert!((v - 0.2).abs() < f64::EPSILON);
        assert!((p - 90.0).abs() < f64::EPSILON);

        let point = grid.grid_to_physical(1, 1);
        assert!((point.x - 0.05).abs() < f64::EPSILON);
        assert!((point.y - 0.0).abs() < f64::EPSILON);
    }
}