//! Physics simulation for digital twins.
//!
//! This module implements physics-based simulations for medical digital twins,
//! including the Pennes bioheat equation for thermal therapy planning.

use serde::{Deserialize, Serialize};

use crate::error::{DigitalTwinError, Result};
use crate::geometry::OrganGeometry;

/// Parameters for bioheat simulation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BioheatParams {
    /// Blood temperature [°C]
    pub blood_temperature: f32,
    /// Blood density [kg/m³]
    pub blood_density: f32,
    /// Blood specific heat [J/(kg·K)]
    pub blood_specific_heat: f32,
    /// Time step for simulation [s]
    pub dt: f32,
    /// Maximum number of iterations
    pub max_iterations: usize,
    /// Convergence tolerance for steady-state
    pub tolerance: f32,
}

impl Default for BioheatParams {
    fn default() -> Self {
        Self {
            blood_temperature: 37.0,
            blood_density: 1050.0,
            blood_specific_heat: 3617.0,
            dt: 0.01,
            max_iterations: 10000,
            tolerance: 1e-4,
        }
    }
}

/// Boundary condition types.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BoundaryCondition {
    /// Fixed temperature (Dirichlet)
    Temperature(f32),
    /// Heat flux (Neumann) [W/m²]
    HeatFlux(f32),
    /// Convective (Robin) with heat transfer coefficient [W/(m²·K)] and ambient temp [°C]
    Convective {
        /// Heat transfer coefficient
        h: f32,
        /// Ambient temperature
        t_ambient: f32,
    },
    /// Adiabatic (no heat transfer)
    Adiabatic,
}

impl Default for BoundaryCondition {
    fn default() -> Self {
        Self::Temperature(37.0)
    }
}

/// Result from a physics simulation.
#[derive(Debug, Clone)]
pub struct SimulationResult {
    /// Temperature field [°C]
    pub temperature: Vec<f32>,
    /// Thermal damage (Arrhenius integral) [-]
    pub damage: Vec<f32>,
    /// Number of iterations used
    pub iterations: usize,
    /// Final residual (for convergence check)
    pub residual: f32,
    /// Simulation time [s]
    pub time: f32,
    /// Maximum temperature reached [°C]
    pub max_temperature: f32,
    /// Volume with significant damage (damage > 1) [mm³]
    pub damaged_volume: f32,
}

impl SimulationResult {
    /// Create a new simulation result.
    pub fn new(size: usize) -> Self {
        Self {
            temperature: vec![37.0; size],
            damage: vec![0.0; size],
            iterations: 0,
            residual: f32::INFINITY,
            time: 0.0,
            max_temperature: 37.0,
            damaged_volume: 0.0,
        }
    }
}

/// Bioheat model implementing the Pennes bioheat equation.
///
/// The Pennes bioheat equation describes heat transfer in living tissue:
///
/// ```text
/// ρc ∂T/∂t = ∇·(k∇T) + ρ_b c_b ω_b (T_b - T) + Q_m + Q_ext
/// ```
///
/// where:
/// - ρ, c = tissue density and specific heat
/// - k = thermal conductivity
/// - ρ_b, c_b = blood density and specific heat
/// - ω_b = blood perfusion rate
/// - T_b = arterial blood temperature
/// - Q_m = metabolic heat generation
/// - Q_ext = external heat source (e.g., ablation probe)
pub struct BioheatModel {
    /// Simulation parameters
    params: BioheatParams,
    /// External heat source field [W/m³]
    heat_source: Vec<f32>,
}

impl BioheatModel {
    /// Create a new bioheat model.
    pub fn new(params: BioheatParams) -> Self {
        Self {
            params,
            heat_source: Vec::new(),
        }
    }

    /// Set external heat source field.
    pub fn set_heat_source(&mut self, source: Vec<f32>) {
        self.heat_source = source;
    }

    /// Clear external heat source.
    pub fn clear_heat_source(&mut self) {
        self.heat_source.clear();
    }

    /// Solve steady-state bioheat equation.
    ///
    /// Uses Gauss-Seidel iteration to find the steady-state temperature
    /// distribution (∂T/∂t = 0).
    pub fn solve_steady_state(
        &self,
        geometry: &OrganGeometry,
        boundary: &BoundaryCondition,
    ) -> Result<SimulationResult> {
        let shape = geometry.shape();
        let [nx, ny, nz] = shape;
        let n = nx * ny * nz;

        if !self.heat_source.is_empty() && self.heat_source.len() != n {
            return Err(DigitalTwinError::ShapeMismatch {
                expected: shape,
                got: [self.heat_source.len(), 1, 1],
            });
        }

        let spacing = geometry.spacing();
        let dx = spacing[0] * 1e-3; // Convert mm to m
        let dy = spacing[1] * 1e-3;
        let dz = spacing[2] * 1e-3;

        // Get property fields
        let k = geometry.thermal_conductivity_field();
        let rho = geometry.density_field();
        let c = geometry.specific_heat_field();
        let omega = geometry.perfusion_field();

        // Get metabolic heat from tissue database
        let q_m: Vec<f32> = geometry
            .data()
            .iter()
            .map(|v| {
                geometry
                    .tissue_db()
                    .get_or_default(v.label.tissue_type())
                    .metabolic_heat
            })
            .collect();

        // Initialize temperature field
        let mut result = SimulationResult::new(n);

        // Initial temperature from geometry
        for (i, voxel) in geometry.data().iter().enumerate() {
            result.temperature[i] = voxel.temperature;
        }

        // Blood perfusion coefficient: ρ_b * c_b * ω_b
        let rho_b_c_b = self.params.blood_density * self.params.blood_specific_heat;
        let t_b = self.params.blood_temperature;

        // Gauss-Seidel iteration
        for iter in 0..self.params.max_iterations {
            let mut max_change = 0.0f32;

            for z in 1..nz - 1 {
                for y in 1..ny - 1 {
                    for x in 1..nx - 1 {
                        let idx = z * nx * ny + y * nx + x;

                        // Get tissue properties at this voxel
                        let k_c = k[idx];
                        let omega_c = omega[idx];

                        // Skip air voxels
                        if k_c < 1e-6 {
                            continue;
                        }

                        // Neighbor indices
                        let idx_xm = idx - 1;
                        let idx_xp = idx + 1;
                        let idx_ym = idx - nx;
                        let idx_yp = idx + nx;
                        let idx_zm = idx - nx * ny;
                        let idx_zp = idx + nx * ny;

                        // Neighbor temperatures
                        let t_xm = result.temperature[idx_xm];
                        let t_xp = result.temperature[idx_xp];
                        let t_ym = result.temperature[idx_ym];
                        let t_yp = result.temperature[idx_yp];
                        let t_zm = result.temperature[idx_zm];
                        let t_zp = result.temperature[idx_zp];

                        // Interface conductivities (harmonic mean)
                        let k_xm = 2.0 * k_c * k[idx_xm] / (k_c + k[idx_xm] + 1e-10);
                        let k_xp = 2.0 * k_c * k[idx_xp] / (k_c + k[idx_xp] + 1e-10);
                        let k_ym = 2.0 * k_c * k[idx_ym] / (k_c + k[idx_ym] + 1e-10);
                        let k_yp = 2.0 * k_c * k[idx_yp] / (k_c + k[idx_yp] + 1e-10);
                        let k_zm = 2.0 * k_c * k[idx_zm] / (k_c + k[idx_zm] + 1e-10);
                        let k_zp = 2.0 * k_c * k[idx_zp] / (k_c + k[idx_zp] + 1e-10);

                        // Diffusion coefficients
                        let ax = 1.0 / (dx * dx);
                        let ay = 1.0 / (dy * dy);
                        let az = 1.0 / (dz * dz);

                        // Perfusion term
                        let perf = rho_b_c_b * omega_c;

                        // External heat source
                        let q_ext = if self.heat_source.is_empty() {
                            0.0
                        } else {
                            self.heat_source[idx]
                        };

                        // Total source: metabolic + external + perfusion heating
                        let source = q_m[idx] + q_ext + perf * t_b;

                        // Coefficient matrix diagonal
                        let diag =
                            (k_xm + k_xp) * ax + (k_ym + k_yp) * ay + (k_zm + k_zp) * az + perf;

                        // Off-diagonal terms
                        let off_diag = k_xm * ax * t_xm
                            + k_xp * ax * t_xp
                            + k_ym * ay * t_ym
                            + k_yp * ay * t_yp
                            + k_zm * az * t_zm
                            + k_zp * az * t_zp;

                        // New temperature
                        let t_new = (off_diag + source) / (diag + 1e-10);

                        // Track maximum change
                        let change = (t_new - result.temperature[idx]).abs();
                        max_change = max_change.max(change);

                        result.temperature[idx] = t_new;
                    }
                }
            }

            // Apply boundary conditions
            self.apply_boundary(&mut result.temperature, shape, boundary);

            result.iterations = iter + 1;
            result.residual = max_change;

            // Check convergence
            if max_change < self.params.tolerance {
                break;
            }
        }

        // Compute statistics
        result.max_temperature = result
            .temperature
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max);

        // Compute thermal damage (simplified Arrhenius model)
        // Ω = A * exp(-Ea / RT) integrated over time
        // For steady-state, we use instantaneous damage indicator
        for (i, &temp) in result.temperature.iter().enumerate() {
            // CEM43 equivalent - normalized to 43°C
            if temp > 43.0 {
                // Simple exponential damage model
                result.damage[i] = (0.5f32).powf(43.0 - temp);
            }
        }

        // Calculate damaged volume (voxels with damage > 1)
        let voxel_volume = spacing[0] * spacing[1] * spacing[2]; // mm³
        result.damaged_volume =
            result.damage.iter().filter(|&&d| d > 1.0).count() as f32 * voxel_volume;

        Ok(result)
    }

    /// Solve transient bioheat equation.
    ///
    /// Uses explicit finite difference time stepping.
    pub fn solve_transient(
        &self,
        geometry: &OrganGeometry,
        boundary: &BoundaryCondition,
        duration: f32,
    ) -> Result<SimulationResult> {
        let shape = geometry.shape();
        let [nx, ny, nz] = shape;
        let n = nx * ny * nz;

        let spacing = geometry.spacing();
        let dx = spacing[0] * 1e-3;
        let dy = spacing[1] * 1e-3;
        let dz = spacing[2] * 1e-3;

        // Get property fields
        let k = geometry.thermal_conductivity_field();
        let rho = geometry.density_field();
        let c = geometry.specific_heat_field();
        let omega = geometry.perfusion_field();

        let q_m: Vec<f32> = geometry
            .data()
            .iter()
            .map(|v| {
                geometry
                    .tissue_db()
                    .get_or_default(v.label.tissue_type())
                    .metabolic_heat
            })
            .collect();

        // Initialize
        let mut result = SimulationResult::new(n);
        for (i, voxel) in geometry.data().iter().enumerate() {
            result.temperature[i] = voxel.temperature;
        }

        let rho_b_c_b = self.params.blood_density * self.params.blood_specific_heat;
        let t_b = self.params.blood_temperature;
        let dt = self.params.dt;

        let num_steps = (duration / dt).ceil() as usize;
        let mut temp_new = result.temperature.clone();

        for step in 0..num_steps {
            for z in 1..nz - 1 {
                for y in 1..ny - 1 {
                    for x in 1..nx - 1 {
                        let idx = z * nx * ny + y * nx + x;

                        let k_c = k[idx];
                        let rho_c = rho[idx];
                        let c_c = c[idx];
                        let omega_c = omega[idx];

                        // Skip air
                        if k_c < 1e-6 || rho_c < 1e-6 {
                            continue;
                        }

                        // Laplacian approximation
                        let t_c = result.temperature[idx];
                        let lap_x = (result.temperature[idx + 1] - 2.0 * t_c
                            + result.temperature[idx - 1])
                            / (dx * dx);
                        let lap_y = (result.temperature[idx + nx] - 2.0 * t_c
                            + result.temperature[idx - nx])
                            / (dy * dy);
                        let lap_z = (result.temperature[idx + nx * ny] - 2.0 * t_c
                            + result.temperature[idx - nx * ny])
                            / (dz * dz);

                        let laplacian = lap_x + lap_y + lap_z;

                        // Perfusion term
                        let perf = rho_b_c_b * omega_c * (t_b - t_c);

                        // Heat source
                        let q_ext = if self.heat_source.is_empty() {
                            0.0
                        } else {
                            self.heat_source[idx]
                        };

                        // Time derivative: dT/dt = (k∇²T + perf + Q) / (ρc)
                        let dt_dt = (k_c * laplacian + perf + q_m[idx] + q_ext) / (rho_c * c_c);

                        temp_new[idx] = t_c + dt * dt_dt;
                    }
                }
            }

            // Swap buffers
            std::mem::swap(&mut result.temperature, &mut temp_new);

            // Apply boundary conditions
            self.apply_boundary(&mut result.temperature, shape, boundary);

            // Update damage integral (Arrhenius)
            for i in 0..n {
                let temp = result.temperature[i];
                if temp > 43.0 {
                    // Simplified damage accumulation
                    let damage_rate = (0.5f32).powf(43.0 - temp);
                    result.damage[i] += damage_rate * dt;
                }
            }

            result.iterations = step + 1;
            result.time = (step + 1) as f32 * dt;
        }

        // Final statistics
        result.max_temperature = result
            .temperature
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max);

        let voxel_volume = spacing[0] * spacing[1] * spacing[2];
        result.damaged_volume =
            result.damage.iter().filter(|&&d| d > 1.0).count() as f32 * voxel_volume;

        Ok(result)
    }

    /// Apply boundary conditions.
    fn apply_boundary(&self, temp: &mut [f32], shape: [usize; 3], bc: &BoundaryCondition) {
        let [nx, ny, nz] = shape;

        match bc {
            BoundaryCondition::Temperature(t) => {
                // Set boundary voxels to fixed temperature
                for z in 0..nz {
                    for y in 0..ny {
                        // X boundaries
                        temp[z * nx * ny + y * nx] = *t;
                        temp[z * nx * ny + y * nx + (nx - 1)] = *t;
                    }
                    for x in 0..nx {
                        // Y boundaries
                        temp[z * nx * ny + x] = *t;
                        temp[z * nx * ny + (ny - 1) * nx + x] = *t;
                    }
                }
                for y in 0..ny {
                    for x in 0..nx {
                        // Z boundaries
                        temp[y * nx + x] = *t;
                        temp[(nz - 1) * nx * ny + y * nx + x] = *t;
                    }
                }
            }
            BoundaryCondition::Adiabatic => {
                // Zero gradient (copy from interior)
                for z in 0..nz {
                    for y in 0..ny {
                        temp[z * nx * ny + y * nx] = temp[z * nx * ny + y * nx + 1];
                        temp[z * nx * ny + y * nx + (nx - 1)] =
                            temp[z * nx * ny + y * nx + (nx - 2)];
                    }
                    for x in 0..nx {
                        temp[z * nx * ny + x] = temp[z * nx * ny + nx + x];
                        temp[z * nx * ny + (ny - 1) * nx + x] =
                            temp[z * nx * ny + (ny - 2) * nx + x];
                    }
                }
                for y in 0..ny {
                    for x in 0..nx {
                        temp[y * nx + x] = temp[nx * ny + y * nx + x];
                        temp[(nz - 1) * nx * ny + y * nx + x] =
                            temp[(nz - 2) * nx * ny + y * nx + x];
                    }
                }
            }
            _ => {
                // Default to body temperature for other BCs
                self.apply_boundary(temp, shape, &BoundaryCondition::Temperature(37.0));
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::geometry::TissueLabel;
    use crate::tissue::TissueType;

    #[test]
    fn test_bioheat_params_default() {
        let params = BioheatParams::default();
        assert_eq!(params.blood_temperature, 37.0);
    }

    #[test]
    fn test_steady_state_uniform() {
        // Uniform tissue, no heat source -> should stay at body temp
        let mut labels = vec![3u8; 125]; // 5x5x5 muscle

        // Set boundaries to air
        for i in 0..5 {
            for j in 0..5 {
                labels[i * 5 + j] = 0;
                labels[4 * 25 + i * 5 + j] = 0;
                labels[i * 25 + j] = 0;
                labels[i * 25 + 4 * 5 + j] = 0;
                labels[i * 25 + j * 5] = 0;
                labels[i * 25 + j * 5 + 4] = 0;
            }
        }

        let geometry = OrganGeometry::from_labels(&labels, [5, 5, 5], [1.0, 1.0, 1.0]).unwrap();

        let params = BioheatParams {
            max_iterations: 100,
            tolerance: 0.1,
            ..Default::default()
        };
        let model = BioheatModel::new(params);

        let result = model
            .solve_steady_state(&geometry, &BoundaryCondition::Temperature(37.0))
            .unwrap();

        // Interior should be close to body temperature
        let center = 2 * 25 + 2 * 5 + 2;
        assert!(
            (result.temperature[center] - 37.0).abs() < 1.0,
            "Center temp should be ~37°C, got {}",
            result.temperature[center]
        );
    }

    #[test]
    fn test_with_heat_source() {
        // Create geometry with liver
        let mut geometry = OrganGeometry::new([10, 10, 10], [1.0, 1.0, 1.0]);

        // Fill interior with liver
        for z in 1..9 {
            for y in 1..9 {
                for x in 1..9 {
                    geometry.set_label(x, y, z, TissueLabel::from(TissueType::Liver));
                }
            }
        }

        // Add heat source in a 3x3x3 region at center (more realistic)
        let mut heat_source = vec![0.0f32; 1000];
        for z in 4..7 {
            for y in 4..7 {
                for x in 4..7 {
                    let idx = z * 100 + y * 10 + x;
                    heat_source[idx] = 1e7; // 10 MW/m³ - typical for RF ablation
                }
            }
        }

        let params = BioheatParams {
            max_iterations: 1000,
            tolerance: 0.001,
            ..Default::default()
        };
        let mut model = BioheatModel::new(params);
        model.set_heat_source(heat_source);

        let result = model
            .solve_steady_state(&geometry, &BoundaryCondition::Temperature(37.0))
            .unwrap();

        // Center should be hotter than boundary
        let center_temp = result.temperature[555];
        assert!(
            center_temp > 37.0,
            "Center with heat source should be warmer than 37°C, got {}",
            center_temp
        );
        // And cooler than unrealistic values
        assert!(
            center_temp < 100.0,
            "Temperature should be physiologically reasonable, got {}",
            center_temp
        );
    }

    #[test]
    fn test_simulation_result() {
        let result = SimulationResult::new(100);
        assert_eq!(result.temperature.len(), 100);
        assert_eq!(result.damage.len(), 100);
        assert_eq!(result.max_temperature, 37.0);
    }
}