//! Ablation probe heat source modeling
//!
//! Models the power deposition from various ablation devices:
//! - Radiofrequency (RF) ablation needles
//! - Microwave ablation antennas
//! - Laser fibers (LITT)

use bioheat_shared::{Point3D, ProbeGeometry};
use serde::{Deserialize, Serialize};

/// Heat source model for ablation probes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProbeHeatSource {
    /// Probe geometry
    pub geometry: ProbeGeometry,
    /// Total power delivered by the probe (Watts)
    pub power: f32,
    /// Heat deposition model
    pub model: HeatDepositionModel,
}

/// Model for how heat is deposited around the probe
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
pub enum HeatDepositionModel {
    /// Uniform power density within a radius
    /// Q(r) = P / V for r < R, 0 otherwise
    #[default]
    Uniform,
    /// Gaussian distribution of power
    /// Q(r) = Q0 * exp(-r²/σ²)
    Gaussian {
        /// Characteristic radius (sigma)
        sigma: f32,
    },
    /// Exponential decay (good for RF)
    /// Q(r) = Q0 * exp(-r/λ)
    Exponential {
        /// Decay length
        lambda: f32,
    },
    /// SAR-based model for microwave
    /// Q(r) = Q0 * (r0/r)^n * exp(-αr)
    Microwave {
        /// Reference radius
        r0: f32,
        /// Decay exponent
        n: f32,
        /// Attenuation coefficient
        alpha: f32,
    },
}

impl ProbeHeatSource {
    /// Create a new probe heat source
    #[must_use]
    pub fn new(geometry: ProbeGeometry, power: f32) -> Self {
        Self {
            geometry,
            power,
            model: HeatDepositionModel::Gaussian { sigma: 0.005 }, // 5mm default
        }
    }

    /// Create RF ablation needle heat source
    #[must_use]
    pub fn rf_needle(position: Point3D, power: f32) -> Self {
        Self {
            geometry: ProbeGeometry::rf_needle(position),
            power,
            model: HeatDepositionModel::Gaussian { sigma: 0.008 }, // 8mm for RF
        }
    }

    /// Create microwave antenna heat source
    #[must_use]
    pub fn microwave(position: Point3D, power: f32) -> Self {
        Self {
            geometry: ProbeGeometry::microwave_antenna(position),
            power,
            model: HeatDepositionModel::Microwave {
                r0: 0.002, // 2mm reference radius
                n: 2.0,
                alpha: 50.0, // 1/m attenuation
            },
        }
    }

    /// Create laser fiber heat source
    #[must_use]
    pub fn laser(position: Point3D, power: f32) -> Self {
        Self {
            geometry: ProbeGeometry::laser_fiber(position),
            power,
            model: HeatDepositionModel::Exponential { lambda: 0.003 }, // 3mm penetration
        }
    }

    /// Compute the volumetric heat generation rate at a point (W/m³)
    ///
    /// This is the Qs term in the Pennes equation
    #[must_use]
    pub fn heat_source_at(&self, point: &Point3D) -> f32 {
        let (axial, radial) = self.geometry.distance_to_axis(point);

        // Only contribute heat within the active region (along axis)
        if axial < 0.0 || axial > self.geometry.active_length {
            return 0.0;
        }

        // Compute volumetric power density based on model
        match self.model {
            HeatDepositionModel::Uniform => self.uniform_heat_at(radial),
            HeatDepositionModel::Gaussian { sigma } => self.gaussian_heat_at(radial, sigma),
            HeatDepositionModel::Exponential { lambda } => self.exponential_heat_at(radial, lambda),
            HeatDepositionModel::Microwave { r0, n, alpha } => {
                self.microwave_heat_at(radial, r0, n, alpha)
            }
        }
    }

    /// Uniform heat distribution
    fn uniform_heat_at(&self, radial_distance: f32) -> f32 {
        // Effective heating radius (larger than probe radius)
        let effective_radius = self.geometry.radius * 5.0;
        if radial_distance > effective_radius {
            return 0.0;
        }

        // Volume of active cylindrical region
        let volume = std::f32::consts::PI
            * effective_radius
            * effective_radius
            * self.geometry.active_length;

        self.power / volume
    }

    /// Gaussian heat distribution
    /// Q(r) = Q0 * exp(-r²/(2σ²))
    /// Normalized so total power = P
    fn gaussian_heat_at(&self, radial_distance: f32, sigma: f32) -> f32 {
        // Cutoff at 4 sigma
        if radial_distance > 4.0 * sigma {
            return 0.0;
        }

        // Peak power density (derived from normalization)
        // For 2D Gaussian integrated over cylinder: Q0 = P / (2πσ²L)
        let q0 =
            self.power / (2.0 * std::f32::consts::PI * sigma * sigma * self.geometry.active_length);

        let r_normalized = radial_distance / sigma;
        q0 * (-0.5 * r_normalized * r_normalized).exp()
    }

    /// Exponential decay heat distribution
    /// Q(r) = Q0 * exp(-r/λ)
    fn exponential_heat_at(&self, radial_distance: f32, lambda: f32) -> f32 {
        // Cutoff at 5 lambda
        if radial_distance > 5.0 * lambda {
            return 0.0;
        }

        // Peak power density (approximate normalization)
        // For exponential: Q0 ≈ P / (2πλ²L)
        let q0 = self.power
            / (2.0 * std::f32::consts::PI * lambda * lambda * self.geometry.active_length);

        q0 * (-radial_distance / lambda).exp()
    }

    /// Microwave-specific SAR distribution
    /// Q(r) = Q0 * (r0/(r+ε))^n * exp(-αr)
    fn microwave_heat_at(&self, radial_distance: f32, r0: f32, n: f32, alpha: f32) -> f32 {
        // Small offset to avoid singularity at r=0
        let epsilon = 0.0001;
        let r = radial_distance + epsilon;

        // Cutoff far from source
        if radial_distance > 0.03 {
            // 3cm max
            return 0.0;
        }

        // Approximate peak density
        let q0 = self.power / (4.0 * std::f32::consts::PI * r0 * r0 * self.geometry.active_length);

        q0 * (r0 / r).powf(n) * (-alpha * radial_distance).exp()
    }

    /// Compute heat source values for a batch of points
    pub fn heat_source_batch(&self, points: &[Point3D]) -> Vec<f32> {
        points.iter().map(|p| self.heat_source_at(p)).collect()
    }

    /// Update probe position
    pub fn set_position(&mut self, position: Point3D) {
        self.geometry.position = position;
    }

    /// Update probe power
    pub fn set_power(&mut self, power: f32) {
        self.power = power.max(0.0);
    }

    /// Get maximum expected heat source (at probe surface)
    #[must_use]
    pub fn max_heat_source(&self) -> f32 {
        self.heat_source_at(&self.geometry.position)
    }
}

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

    #[test]
    fn test_rf_needle_heat_source() {
        let source = ProbeHeatSource::rf_needle(Point3D::origin(), 15.0);

        // Heat should be maximum near the probe
        let q_near = source.heat_source_at(&Point3D::new(0.001, 0.0, 0.015));
        let q_far = source.heat_source_at(&Point3D::new(0.05, 0.0, 0.015));

        assert!(q_near > q_far);
        assert!(q_near > 0.0);
    }

    #[test]
    fn test_heat_source_axial_bounds() {
        let source = ProbeHeatSource::rf_needle(Point3D::origin(), 15.0);

        // Before probe tip (negative z)
        let q_before = source.heat_source_at(&Point3D::new(0.0, 0.0, -0.01));
        assert!((q_before).abs() < 1e-6);

        // After active length
        let active_len = source.geometry.active_length;
        let q_after = source.heat_source_at(&Point3D::new(0.0, 0.0, active_len + 0.01));
        assert!((q_after).abs() < 1e-6);
    }

    #[test]
    fn test_gaussian_model() {
        let mut source = ProbeHeatSource::rf_needle(Point3D::origin(), 15.0);
        source.model = HeatDepositionModel::Gaussian { sigma: 0.005 };

        // Heat should decay with distance
        let q0 = source.heat_source_at(&Point3D::new(0.0, 0.0, 0.015));
        let q1 = source.heat_source_at(&Point3D::new(0.005, 0.0, 0.015));
        let q2 = source.heat_source_at(&Point3D::new(0.010, 0.0, 0.015));

        assert!(q0 > q1);
        assert!(q1 > q2);
    }

    #[test]
    fn test_microwave_model() {
        let source = ProbeHeatSource::microwave(Point3D::origin(), 50.0);

        // Microwave should have higher power density near probe
        let q_near = source.heat_source_at(&Point3D::new(0.003, 0.0, 0.02));
        assert!(q_near > 0.0);
    }

    #[test]
    fn test_set_power() {
        let mut source = ProbeHeatSource::rf_needle(Point3D::origin(), 15.0);
        source.set_power(30.0);
        assert!((source.power - 30.0).abs() < 1e-6);

        // Negative power should be clamped to 0
        source.set_power(-10.0);
        assert!((source.power).abs() < 1e-6);
    }

    #[test]
    fn test_batch_computation() {
        let source = ProbeHeatSource::rf_needle(Point3D::origin(), 15.0);
        let points = vec![
            Point3D::new(0.0, 0.0, 0.015),
            Point3D::new(0.01, 0.0, 0.015),
            Point3D::new(0.02, 0.0, 0.015),
        ];

        let heat_values = source.heat_source_batch(&points);
        assert_eq!(heat_values.len(), 3);
        assert!(heat_values[0] > heat_values[1]);
        assert!(heat_values[1] > heat_values[2]);
    }
}