rustytorch/demos/rtx-cardiosim-demo/src/ecg.rs

//! ECG computation from cardiac voltage fields.
//!
//! Computes virtual ECG signals using the solid angle method.

use cardiosim_shared::{ECGLeads, ECGResult, Point3D, SimulationResult};

use crate::CardioSimError;

/// ECG computer.
#[derive(Debug)]
pub struct ECGComputer {
    leads: ECGLeads,
    electrode_positions: Vec<Point3D>,
}

impl ECGComputer {
    /// Create a new ECG computer.
    #[must_use]
    pub fn new(leads: &ECGLeads) -> Self {
        let mut electrode_positions = leads.custom_electrodes.clone();

        if leads.standard_12_lead {
            // Standard 12-lead electrode positions (simplified)
            // These are relative positions on a normalized torso

            // Limb leads
            electrode_positions.push(Point3D::new(-100.0, 0.0, 0.0)); // LA
            electrode_positions.push(Point3D::new(100.0, 0.0, 0.0)); // RA
            electrode_positions.push(Point3D::new(0.0, 0.0, -100.0)); // LL
            electrode_positions.push(Point3D::new(0.0, 0.0, -100.0)); // RL (reference)

            // Precordial leads (V1-V6)
            electrode_positions.push(Point3D::new(20.0, 50.0, 30.0)); // V1
            electrode_positions.push(Point3D::new(-20.0, 50.0, 30.0)); // V2
            electrode_positions.push(Point3D::new(-50.0, 50.0, 20.0)); // V3
            electrode_positions.push(Point3D::new(-70.0, 50.0, 0.0)); // V4
            electrode_positions.push(Point3D::new(-80.0, 40.0, -20.0)); // V5
            electrode_positions.push(Point3D::new(-90.0, 30.0, -40.0)); // V6
        }

        Self {
            leads: leads.clone(),
            electrode_positions,
        }
    }

    /// Compute ECG from simulation result.
    pub fn compute(&self, sim: &SimulationResult) -> Result<ECGResult, CardioSimError> {
        if sim.voltage_fields.is_empty() {
            return Err(CardioSimError::SimulationError(
                "No voltage fields to compute ECG".to_string(),
            ));
        }

        let times: Vec<f32> = sim.voltage_fields.iter().map(|f| f.time).collect();
        let n_times = times.len();

        // Compute potentials at each electrode
        let mut potentials: Vec<Vec<f32>> =
            vec![vec![0.0; n_times]; self.electrode_positions.len()];

        for (t_idx, field) in sim.voltage_fields.iter().enumerate() {
            for (e_idx, electrode) in self.electrode_positions.iter().enumerate() {
                potentials[e_idx][t_idx] = self.compute_potential(electrode, &field.voltages, sim);
            }
        }

        // Compute lead signals
        let (lead_i, lead_ii, lead_iii) = if self.electrode_positions.len() >= 4 {
            // LA=0, RA=1, LL=2
            let lead_i: Vec<f32> = (0..n_times)
                .map(|t| potentials[0][t] - potentials[1][t])
                .collect();
            let lead_ii: Vec<f32> = (0..n_times)
                .map(|t| potentials[2][t] - potentials[1][t])
                .collect();
            let lead_iii: Vec<f32> = (0..n_times)
                .map(|t| potentials[2][t] - potentials[0][t])
                .collect();
            (lead_i, lead_ii, lead_iii)
        } else {
            (vec![0.0; n_times], vec![0.0; n_times], vec![0.0; n_times])
        };

        // Precordial leads (V1-V6)
        let precordial: Vec<Vec<f32>> = if self.electrode_positions.len() >= 10 {
            // Wilson's central terminal
            let wct: Vec<f32> = (0..n_times)
                .map(|t| (potentials[0][t] + potentials[1][t] + potentials[2][t]) / 3.0)
                .collect();

            (4..10)
                .map(|e| (0..n_times).map(|t| potentials[e][t] - wct[t]).collect())
                .collect()
        } else {
            vec![vec![0.0; n_times]; 6]
        };

        // Calculate heart rate from R-R intervals
        let heart_rate = self.calculate_heart_rate(&lead_ii);

        // Calculate QT interval
        let qt_interval = self.calculate_qt_interval(&lead_ii, &times);

        // Calculate QRS duration
        let qrs_duration = self.calculate_qrs_duration(&lead_ii, &times);

        Ok(ECGResult {
            times,
            lead_i,
            lead_ii,
            lead_iii,
            precordial,
            heart_rate,
            qt_interval,
            qrs_duration,
        })
    }

    fn compute_potential(
        &self,
        electrode: &Point3D,
        voltages: &[f32],
        sim: &SimulationResult,
    ) -> f32 {
        // Use solid angle method (simplified)
        // Φ = (1/4πσ) ∫ Vm ∇·(r/|r|³) dV

        let mut potential = 0.0;
        let n_vertices = voltages.len();

        // Need mesh info - use activation map vertex count
        let n_mesh_vertices = sim.activation_map.activation_times.len();
        if n_mesh_vertices == 0 || n_vertices != n_mesh_vertices {
            return 0.0;
        }

        // Simplified: treat each vertex as a dipole source
        for (i, &voltage) in voltages.iter().enumerate() {
            // Estimate vertex position from activation time gradient
            // (in real implementation, would use mesh vertices)
            let t = sim.activation_map.activation_times[i];
            if !t.is_finite() {
                continue;
            }

            // Create pseudo-position based on activation time
            let r = t / 100.0; // Convert to distance-like quantity
            let theta = (i as f32) * 0.1;
            let phi = (i as f32) * 0.2;

            let vertex = Point3D::new(
                r * theta.sin() * phi.cos() * 30.0,
                r * theta.sin() * phi.sin() * 25.0,
                r * theta.cos() * 50.0,
            );

            let dx = electrode.x - vertex.x;
            let dy = electrode.y - vertex.y;
            let dz = electrode.z - vertex.z;
            let dist_sq = dx * dx + dy * dy + dz * dz + 1e-6;
            let dist = dist_sq.sqrt();

            // Dipole contribution (simplified)
            let contribution = voltage / (dist * dist);
            potential += contribution;
        }

        // Scale
        potential * 0.001
    }

    fn calculate_heart_rate(&self, lead_ii: &[f32]) -> f32 {
        if lead_ii.len() < 10 {
            return 60.0; // Default
        }

        // Find R-peaks (simplified: look for local maxima above threshold)
        let threshold = lead_ii.iter().fold(f32::MIN, |a, &b| a.max(b)) * 0.5;
        let mut r_peak_indices = Vec::new();

        for i in 1..lead_ii.len() - 1 {
            if lead_ii[i] > threshold && lead_ii[i] > lead_ii[i - 1] && lead_ii[i] > lead_ii[i + 1]
            {
                // Check if far enough from last peak
                if r_peak_indices.is_empty() || i - r_peak_indices.last().unwrap() > 50 {
                    r_peak_indices.push(i);
                }
            }
        }

        if r_peak_indices.len() < 2 {
            return 60.0;
        }

        // Calculate average R-R interval
        let mut rr_sum = 0.0;
        for i in 1..r_peak_indices.len() {
            rr_sum += (r_peak_indices[i] - r_peak_indices[i - 1]) as f32;
        }
        let avg_rr = rr_sum / (r_peak_indices.len() - 1) as f32;

        // Convert to heart rate (assuming 1ms per sample)
        60000.0 / avg_rr
    }

    fn calculate_qt_interval(&self, lead_ii: &[f32], times: &[f32]) -> f32 {
        if lead_ii.len() < 100 || times.len() < 100 {
            return 400.0; // Default ~400ms
        }

        // Find Q onset and T end (simplified)
        let threshold = lead_ii.iter().fold(f32::MIN, |a, &b| a.max(b)) * 0.1;

        // Find first significant deflection (Q)
        let q_onset = lead_ii.iter().position(|&v| v.abs() > threshold);

        // Find return to baseline after T wave
        let t_end = lead_ii.iter().rposition(|&v| v.abs() > threshold);

        match (q_onset, t_end) {
            (Some(q), Some(t)) if t > q && q < times.len() && t < times.len() => {
                times[t] - times[q]
            }
            _ => 400.0,
        }
    }

    fn calculate_qrs_duration(&self, lead_ii: &[f32], times: &[f32]) -> f32 {
        if lead_ii.len() < 50 || times.len() < 50 {
            return 100.0; // Default ~100ms
        }

        // Find QRS complex (high amplitude, fast changes)
        let threshold = lead_ii.iter().fold(f32::MIN, |a, &b| a.max(b)) * 0.3;

        let qrs_start = lead_ii.iter().position(|&v| v.abs() > threshold);
        let qrs_end = lead_ii
            .iter()
            .enumerate()
            .skip(qrs_start.unwrap_or(0) + 10)
            .find(|(_, v)| v.abs() < threshold)
            .map(|(i, _)| i);

        match (qrs_start, qrs_end) {
            (Some(start), Some(end)) if end > start && start < times.len() && end < times.len() => {
                times[end] - times[start]
            }
            _ => 100.0,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cardiosim_shared::{APDMap, ActivationMap, VoltageField};

    fn create_test_simulation() -> SimulationResult {
        let n = 100;
        let n_times = 50;

        let voltage_fields: Vec<VoltageField> = (0..n_times)
            .map(|t| {
                let time = t as f32 * 2.0;
                let voltages: Vec<f32> = (0..n)
                    .map(|i| {
                        let phase = (time - i as f32) / 50.0;
                        if phase > 0.0 && phase < 0.5 {
                            20.0 * (std::f32::consts::PI * 2.0 * phase).sin()
                        } else {
                            -85.0
                        }
                    })
                    .collect();
                VoltageField { time, voltages }
            })
            .collect();

        let activation_times: Vec<f32> = (0..n).map(|i| i as f32).collect();

        SimulationResult {
            times: voltage_fields.iter().map(|f| f.time).collect(),
            voltage_fields,
            activation_map: ActivationMap {
                activation_times,
                conduction_velocity: vec![0.8; n],
            },
            apd_map: APDMap {
                apd50: vec![150.0; n],
                apd90: vec![250.0; n],
                dispersion: 20.0,
            },
            arrhythmias: vec![],
            stats: cardiosim_shared::SimulationStats {
                compute_time: 1.0,
                time_steps: n_times,
                speedup_factor: 1.0,
                inference_time: 1.0,
            },
        }
    }

    #[test]
    fn test_ecg_computer_creation() {
        let leads = ECGLeads::default();
        let computer = ECGComputer::new(&leads);
        assert!(!computer.electrode_positions.is_empty());
    }

    #[test]
    fn test_compute_ecg() {
        let leads = ECGLeads::default();
        let computer = ECGComputer::new(&leads);
        let sim = create_test_simulation();

        let result = computer.compute(&sim);
        assert!(result.is_ok());

        let ecg = result.unwrap();
        assert!(!ecg.times.is_empty());
        assert_eq!(ecg.lead_i.len(), ecg.times.len());
        assert_eq!(ecg.lead_ii.len(), ecg.times.len());
        assert_eq!(ecg.lead_iii.len(), ecg.times.len());
        assert_eq!(ecg.precordial.len(), 6);
    }

    #[test]
    fn test_heart_rate_calculation() {
        let leads = ECGLeads::default();
        let computer = ECGComputer::new(&leads);

        // Create synthetic ECG with known R-R interval
        let lead_ii: Vec<f32> = (0..1000)
            .map(|i| {
                // R-peak every 100 samples
                if i % 100 == 0 { 1.0 } else { 0.0 }
            })
            .collect();

        let hr = computer.calculate_heart_rate(&lead_ii);
        // Should be around 600 bpm (R every 100ms)
        assert!(hr > 500.0 && hr < 700.0);
    }

    #[test]
    fn test_custom_electrodes() {
        let leads = ECGLeads {
            standard_12_lead: false,
            custom_electrodes: vec![Point3D::new(0.0, 50.0, 0.0), Point3D::new(0.0, -50.0, 0.0)],
        };

        let computer = ECGComputer::new(&leads);
        assert_eq!(computer.electrode_positions.len(), 2);
    }
}