rustytorch/demos/rtx-digital-twin-demo/src/simulation.rs

//! Simulation module for Digital Twin demo.

use tokio::sync::mpsc;

use rtx_digital_twin::{
    AblationProbe, BioheatParams, BoundaryCondition, DigitalTwin,
    InterventionType as CoreInterventionType, TwinConfig,
};
use rtx_digital_twin_shared::{
    DigitalTwinDemoError, DigitalTwinDemoResult, GeometrySummary, InterventionType, ProbeConfig,
    SimulationConfig, SimulationProgress, SimulationResult, SliceData, SliceOrientation,
    TissueCount, TissueType, WhatIfRequest, WhatIfResult,
};

use crate::presets::GeometryPresets;

/// Digital twin simulator for the demo.
pub struct TwinSimulator {
    config: SimulationConfig,
    twin: Option<DigitalTwin>,
}

impl TwinSimulator {
    /// Create a new simulator with the given configuration.
    #[must_use]
    pub fn new(config: SimulationConfig) -> Self {
        Self { config, twin: None }
    }

    /// Initialize the digital twin from a preset.
    pub fn init_from_preset(&mut self, preset: &str) -> DigitalTwinDemoResult<GeometrySummary> {
        let resolution = [
            self.config.resolution[0] as usize,
            self.config.resolution[1] as usize,
            self.config.resolution[2] as usize,
        ];

        let geometry = GeometryPresets::create(preset, resolution, self.config.spacing)?;

        let twin_config = TwinConfig {
            bioheat_params: BioheatParams {
                blood_temperature: self.config.blood_temperature,
                dt: self.config.time_step,
                ..Default::default()
            },
            boundary_condition: BoundaryCondition::Temperature(self.config.blood_temperature),
            compute_damage: self.config.compute_damage,
        };

        let twin = DigitalTwin::with_config(geometry, twin_config);
        let summary = self.geometry_summary(&twin);

        self.twin = Some(twin);
        Ok(summary)
    }

    /// Get geometry summary.
    fn geometry_summary(&self, twin: &DigitalTwin) -> GeometrySummary {
        let core_summary = twin.geometry_summary();
        let histogram = twin.geometry().tissue_histogram();

        let total_voxels = core_summary.total_voxels as u32;
        let tissue_histogram: Vec<TissueCount> = histogram
            .iter()
            .map(|(tissue_type, &count)| {
                let demo_type = convert_tissue_type(*tissue_type);
                TissueCount {
                    tissue_type: demo_type,
                    count: count as u32,
                    percentage: (count as f32 / total_voxels as f32) * 100.0,
                }
            })
            .collect();

        GeometrySummary {
            shape: [
                core_summary.shape[0] as u32,
                core_summary.shape[1] as u32,
                core_summary.shape[2] as u32,
            ],
            spacing: core_summary.spacing,
            dimensions: core_summary.dimensions,
            total_voxels,
            tissue_voxels: core_summary.tissue_voxels as u32,
            tissue_volume: core_summary.tissue_volume,
            tissue_histogram,
        }
    }

    /// Run simulation with progress reporting.
    pub async fn run_simulation(
        &mut self,
        probe: ProbeConfig,
        duration: Option<f32>,
        steady_state: bool,
        progress_tx: mpsc::Sender<SimulationProgress>,
    ) -> DigitalTwinDemoResult<SimulationResult> {
        let twin = self
            .twin
            .as_mut()
            .ok_or(DigitalTwinDemoError::TwinNotInitialized)?;

        let ablation_probe = create_ablation_probe(&probe);
        let sim_duration = duration.unwrap_or(self.config.duration);

        // Send initial progress
        let _ = progress_tx
            .send(SimulationProgress {
                status: "Starting simulation".to_string(),
                total_duration: sim_duration,
                ..Default::default()
            })
            .await;

        let start_time = std::time::Instant::now();

        let result = if steady_state {
            twin.simulate_intervention_steady(&ablation_probe)
                .map_err(|e| DigitalTwinDemoError::SimulationError(e.to_string()))?
        } else {
            twin.simulate_intervention(&ablation_probe, sim_duration)
                .map_err(|e| DigitalTwinDemoError::SimulationError(e.to_string()))?
        };

        // Send completion progress
        let _ = progress_tx
            .send(SimulationProgress {
                current_time: sim_duration,
                total_duration: sim_duration,
                iteration: result.iterations as u32,
                total_iterations: result.iterations as u32,
                max_temperature: result.max_temperature,
                damaged_volume: result.damaged_volume,
                residual: result.residual,
                status: "Complete".to_string(),
            })
            .await;

        let elapsed = start_time.elapsed().as_secs_f32();

        // Compute severe damage volume
        let severe_damage_volume = result.damage.iter().filter(|&&d| d > 4.6).count() as f32
            * self.config.spacing.iter().product::<f32>();

        Ok(SimulationResult {
            temperature_field: result.temperature.clone(),
            damage_field: result.damage.clone(),
            dimensions: self.config.resolution,
            max_temperature: result.max_temperature,
            total_damaged_volume: result.damaged_volume,
            severe_damage_volume,
            simulation_time: elapsed,
            iterations: result.iterations as u32,
            safety_ok: result.max_temperature < 100.0,
        })
    }

    /// Run what-if analysis.
    pub async fn what_if(&mut self, request: WhatIfRequest) -> DigitalTwinDemoResult<WhatIfResult> {
        let twin = self
            .twin
            .as_mut()
            .ok_or(DigitalTwinDemoError::TwinNotInitialized)?;

        let ablation_probe = create_ablation_probe(&request.probe);

        let result = twin
            .what_if(&ablation_probe, request.duration)
            .map_err(|e| DigitalTwinDemoError::SimulationError(e.to_string()))?;

        let recommendation = if result.safety_margin_ok {
            if result.total_damaged_volume > 1000.0 {
                "Treatment parameters appear safe. Adequate ablation zone expected.".to_string()
            } else {
                "Treatment parameters appear safe but ablation zone may be insufficient. Consider increasing power or duration.".to_string()
            }
        } else {
            "WARNING: Safety margin exceeded. Reduce power or duration to avoid damage to healthy tissue.".to_string()
        };

        Ok(WhatIfResult {
            intervention_type: request.probe.intervention_type,
            power: request.probe.power,
            duration: request.duration,
            max_temperature: result.max_temperature,
            total_damaged_volume: result.total_damaged_volume,
            severe_damage_volume: result.severe_damage_volume,
            moderate_damage_volume: result.moderate_damage_volume,
            max_boundary_temperature: result.max_boundary_temperature,
            safety_ok: result.safety_margin_ok,
            recommendation,
        })
    }

    /// Get a slice of the current temperature/damage field.
    pub fn get_slice(
        &self,
        orientation: SliceOrientation,
        index: u32,
    ) -> DigitalTwinDemoResult<SliceData> {
        let twin = self
            .twin
            .as_ref()
            .ok_or(DigitalTwinDemoError::TwinNotInitialized)?;

        let geometry = twin.geometry();
        let shape = geometry.shape();
        let result = twin.last_result();

        let (width, height, slice_fn): (usize, usize, Box<dyn Fn(usize, usize) -> usize>) =
            match orientation {
                SliceOrientation::Axial => {
                    let z = index as usize;
                    (
                        shape[0],
                        shape[1],
                        Box::new(move |x, y| z * shape[0] * shape[1] + y * shape[0] + x),
                    )
                }
                SliceOrientation::Coronal => {
                    let y = index as usize;
                    (
                        shape[0],
                        shape[2],
                        Box::new(move |x, z| z * shape[0] * shape[1] + y * shape[0] + x),
                    )
                }
                SliceOrientation::Sagittal => {
                    let x = index as usize;
                    (
                        shape[1],
                        shape[2],
                        Box::new(move |y, z| z * shape[0] * shape[1] + y * shape[0] + x),
                    )
                }
            };

        let mut temperature = vec![0.0f32; width * height];
        let mut damage = vec![0.0f32; width * height];
        let mut tissue = vec![0u8; width * height];

        for h in 0..height {
            for w in 0..width {
                let idx = slice_fn(w, h);
                let slice_idx = h * width + w;

                if let Some(voxel) = geometry.data().get(idx) {
                    tissue[slice_idx] = voxel.label.value();
                    temperature[slice_idx] = voxel.temperature;
                    damage[slice_idx] = voxel.damage;
                }

                // Override with simulation result if available
                if let Some(res) = result
                    && idx < res.temperature.len()
                {
                    temperature[slice_idx] = res.temperature[idx];
                    damage[slice_idx] = res.damage[idx];
                }
            }
        }

        Ok(SliceData {
            orientation,
            index,
            temperature,
            damage,
            tissue,
            dimensions: [width as u32, height as u32],
        })
    }

    /// Reset the twin state.
    pub fn reset(&mut self) -> DigitalTwinDemoResult<()> {
        if let Some(ref mut twin) = self.twin {
            twin.reset_temperature();
        }
        Ok(())
    }
}

/// Create ablation probe from demo config.
fn create_ablation_probe(config: &ProbeConfig) -> AblationProbe {
    let ablation_type = match config.intervention_type {
        InterventionType::RadiofrequencyAblation => CoreInterventionType::RadiofrequencyAblation,
        InterventionType::MicrowaveAblation => CoreInterventionType::MicrowaveAblation,
        InterventionType::HIFU => CoreInterventionType::HIFU,
        InterventionType::Cryoablation => CoreInterventionType::Cryoablation,
        InterventionType::LaserAblation => CoreInterventionType::LaserAblation,
    };

    AblationProbe::new(config.position, config.power)
        .with_type(ablation_type)
        .with_active_length(config.active_length)
        .with_diameter(config.diameter)
}

/// Convert core tissue type to demo tissue type.
fn convert_tissue_type(core_type: rtx_digital_twin::TissueType) -> TissueType {
    match core_type {
        rtx_digital_twin::TissueType::Air => TissueType::Air,
        rtx_digital_twin::TissueType::CorticalBone
        | rtx_digital_twin::TissueType::TrabecularBone => TissueType::Bone,
        rtx_digital_twin::TissueType::Muscle => TissueType::Muscle,
        rtx_digital_twin::TissueType::Fat => TissueType::Fat,
        rtx_digital_twin::TissueType::Liver => TissueType::Liver,
        rtx_digital_twin::TissueType::Kidney => TissueType::Kidney,
        rtx_digital_twin::TissueType::Lung => TissueType::Lung,
        rtx_digital_twin::TissueType::BrainGrayMatter
        | rtx_digital_twin::TissueType::BrainWhiteMatter => TissueType::Brain,
        rtx_digital_twin::TissueType::Heart => TissueType::Heart,
        rtx_digital_twin::TissueType::Tumor => TissueType::Tumor,
        rtx_digital_twin::TissueType::Blood => TissueType::BloodVessel,
        rtx_digital_twin::TissueType::Skin => TissueType::Skin,
        _ => TissueType::Air, // Default for any other types
    }
}

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

    #[test]
    fn test_simulator_creation() {
        let config = SimulationConfig::default();
        let simulator = TwinSimulator::new(config);
        assert!(simulator.twin.is_none());
    }

    #[test]
    fn test_init_from_preset() {
        let config = SimulationConfig {
            resolution: [32, 32, 32],
            ..Default::default()
        };
        let mut simulator = TwinSimulator::new(config);

        let summary = simulator.init_from_preset("liver_tumor").unwrap();
        assert_eq!(summary.shape, [32, 32, 32]);
        assert!(simulator.twin.is_some());
    }

    #[test]
    fn test_get_slice() {
        let config = SimulationConfig {
            resolution: [32, 32, 32],
            ..Default::default()
        };
        let mut simulator = TwinSimulator::new(config);
        simulator.init_from_preset("simple_sphere").unwrap();

        let slice = simulator.get_slice(SliceOrientation::Axial, 16).unwrap();
        assert_eq!(slice.dimensions, [32, 32]);
        assert_eq!(slice.temperature.len(), 32 * 32);
    }

    #[tokio::test]
    async fn test_run_simulation() {
        let config = SimulationConfig {
            resolution: [16, 16, 16],
            duration: 1.0,
            ..Default::default()
        };
        let mut simulator = TwinSimulator::new(config);
        simulator.init_from_preset("simple_sphere").unwrap();

        let probe = ProbeConfig {
            position: [8.0, 8.0, 8.0],
            power: 30.0,
            ..Default::default()
        };

        let (tx, _rx) = mpsc::channel(100);
        let result = simulator
            .run_simulation(probe, None, true, tx)
            .await
            .unwrap();

        assert!(result.max_temperature >= 37.0);
        assert_eq!(result.dimensions, [16, 16, 16]);
    }

    #[test]
    fn test_reset() {
        let config = SimulationConfig {
            resolution: [16, 16, 16],
            ..Default::default()
        };
        let mut simulator = TwinSimulator::new(config);
        simulator.init_from_preset("simple_sphere").unwrap();
        simulator.reset().unwrap();
    }
}