rustytorch/demos/cardiosim-shared/src/lib.rs

//! Shared IPC types for `CardioSim` cardiac electrophysiology demo.
//!
//! This crate provides data structures for communication between
//! the Tauri frontend and Rust backend for cardiac simulation.

use serde::{Deserialize, Serialize};

// ============================================================================
// Heart Geometry Types
// ============================================================================

/// 3D coordinate.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Point3D {
    /// X coordinate (mm)
    pub x: f32,
    /// Y coordinate (mm)
    pub y: f32,
    /// Z coordinate (mm)
    pub z: f32,
}

impl Point3D {
    /// Create a new point.
    #[must_use]
    pub fn new(x: f32, y: f32, z: f32) -> Self {
        Self { x, y, z }
    }

    /// Euclidean distance to another point.
    #[must_use]
    pub fn distance_to(&self, other: &Point3D) -> f32 {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        let dz = self.z - other.z;
        (dx * dx + dy * dy + dz * dz).sqrt()
    }
}

/// 3D direction vector.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Vector3D {
    /// X component
    pub x: f32,
    /// Y component
    pub y: f32,
    /// Z component
    pub z: f32,
}

impl Vector3D {
    /// Create a new vector.
    #[must_use]
    pub fn new(x: f32, y: f32, z: f32) -> Self {
        Self { x, y, z }
    }

    /// Normalize the vector.
    #[must_use]
    pub fn normalize(&self) -> Self {
        let mag = (self.x * self.x + self.y * self.y + self.z * self.z).sqrt();
        if mag > 1e-8 {
            Self {
                x: self.x / mag,
                y: self.y / mag,
                z: self.z / mag,
            }
        } else {
            Self::new(0.0, 0.0, 1.0)
        }
    }
}

/// Heart mesh geometry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeartMesh {
    /// Mesh name
    pub name: String,
    /// Vertex positions
    pub vertices: Vec<Point3D>,
    /// Triangle indices
    pub triangles: Vec<[usize; 3]>,
    /// Tetrahedra indices (for volume mesh)
    pub tetrahedra: Option<Vec<[usize; 4]>>,
    /// Fiber directions at each vertex
    pub fibers: Vec<Vector3D>,
    /// Sheet directions at each vertex
    pub sheets: Vec<Vector3D>,
    /// Region labels (atria, ventricles, etc.)
    pub regions: Vec<HeartRegion>,
    /// Vertex region assignments
    pub vertex_regions: Vec<usize>,
}

/// Heart region.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HeartRegion {
    /// Left atrium
    LeftAtrium,
    /// Right atrium
    RightAtrium,
    /// Left ventricle
    LeftVentricle,
    /// Right ventricle
    RightVentricle,
    /// Septum
    Septum,
    /// Purkinje system
    Purkinje,
    /// Sinoatrial node
    SANode,
    /// Atrioventricular node
    AVNode,
    /// Bundle of His
    BundleOfHis,
}

impl HeartRegion {
    /// Get display name.
    #[must_use]
    pub fn display_name(&self) -> &'static str {
        match self {
            HeartRegion::LeftAtrium => "Left Atrium",
            HeartRegion::RightAtrium => "Right Atrium",
            HeartRegion::LeftVentricle => "Left Ventricle",
            HeartRegion::RightVentricle => "Right Ventricle",
            HeartRegion::Septum => "Septum",
            HeartRegion::Purkinje => "Purkinje System",
            HeartRegion::SANode => "SA Node",
            HeartRegion::AVNode => "AV Node",
            HeartRegion::BundleOfHis => "Bundle of His",
        }
    }
}

// ============================================================================
// Electrophysiology Types
// ============================================================================

/// Ionic model type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[derive(Default)]
pub enum IonicModel {
    /// Mitchell-Schaeffer (simple 2-variable)
    #[default]
    MitchellSchaeffer,
    /// FitzHugh-Nagumo
    FitzHughNagumo,
    /// Aliev-Panfilov
    AlievPanfilov,
    /// ten Tusscher-Panfilov
    TenTusscherPanfilov,
    /// O'Hara-Rudy
    OHaraRudy,
}

/// Tissue conductivity parameters.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConductivityParams {
    /// Longitudinal conductivity (along fiber, mS/mm)
    pub sigma_l: f32,
    /// Transverse conductivity (mS/mm)
    pub sigma_t: f32,
    /// Normal conductivity (mS/mm)
    pub sigma_n: f32,
}

impl Default for ConductivityParams {
    fn default() -> Self {
        Self {
            sigma_l: 0.17,
            sigma_t: 0.019,
            sigma_n: 0.019,
        }
    }
}

/// Action potential state.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActionPotentialState {
    /// Transmembrane voltage (mV)
    pub voltage: f32,
    /// Recovery variable (gating)
    pub recovery: f32,
    /// Calcium concentration (mM)
    pub calcium: Option<f32>,
    /// Sodium concentration (mM)
    pub sodium: Option<f32>,
    /// Potassium concentration (mM)
    pub potassium: Option<f32>,
}

impl Default for ActionPotentialState {
    fn default() -> Self {
        Self {
            voltage: -85.0, // Resting potential
            recovery: 0.0,
            calcium: Some(0.0001),
            sodium: Some(10.0),
            potassium: Some(140.0),
        }
    }
}

/// Action potential phase.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum APPhase {
    /// Resting (phase 4)
    Resting,
    /// Upstroke (phase 0)
    Upstroke,
    /// Early repolarization (phase 1)
    EarlyRepol,
    /// Plateau (phase 2)
    Plateau,
    /// Repolarization (phase 3)
    Repolarization,
    /// Refractory
    Refractory,
}

// ============================================================================
// Simulation Types
// ============================================================================

/// Simulation request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationRequest {
    /// Heart mesh
    pub mesh: HeartMesh,
    /// Simulation configuration
    pub config: SimulationConfig,
    /// Stimulation protocol
    pub protocol: StimulationProtocol,
}

/// Simulation configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationConfig {
    /// Ionic model to use
    pub ionic_model: IonicModel,
    /// Time step (ms)
    pub dt: f32,
    /// Total simulation time (ms)
    pub total_time: f32,
    /// Output interval (ms)
    pub output_interval: f32,
    /// Use neural operator (vs traditional solver)
    pub use_neural_operator: bool,
    /// Tissue conductivity
    pub conductivity: ConductivityParams,
}

impl Default for SimulationConfig {
    fn default() -> Self {
        Self {
            ionic_model: IonicModel::default(),
            dt: 0.01,
            total_time: 500.0,
            output_interval: 1.0,
            use_neural_operator: true,
            conductivity: ConductivityParams::default(),
        }
    }
}

/// Stimulation protocol.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StimulationProtocol {
    /// Stimulation sites
    pub sites: Vec<StimulationSite>,
    /// Protocol type
    pub protocol_type: ProtocolType,
}

/// Stimulation site.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StimulationSite {
    /// Site name
    pub name: String,
    /// Center position
    pub center: Point3D,
    /// Radius (mm)
    pub radius: f32,
    /// Stimulation current (μA/mm²)
    pub current: f32,
    /// Stimulation times (ms)
    pub times: Vec<f32>,
    /// Pulse duration (ms)
    pub duration: f32,
}

/// Protocol type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProtocolType {
    /// Normal sinus rhythm
    SinusRhythm,
    /// Pacing from custom site
    Pacing,
    /// S1-S2 restitution protocol
    S1S2,
    /// Burst pacing (induce arrhythmia)
    BurstPacing,
    /// Cross-field stimulation
    CrossField,
}

/// Simulation result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationResult {
    /// Time points (ms)
    pub times: Vec<f32>,
    /// Voltage fields at each time
    pub voltage_fields: Vec<VoltageField>,
    /// Activation maps
    pub activation_map: ActivationMap,
    /// APD maps
    pub apd_map: APDMap,
    /// Detected arrhythmias
    pub arrhythmias: Vec<ArrhythmiaEvent>,
    /// Simulation statistics
    pub stats: SimulationStats,
}

/// Voltage field at one time point.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VoltageField {
    /// Time (ms)
    pub time: f32,
    /// Voltage at each vertex (mV)
    pub voltages: Vec<f32>,
}

/// Activation time map.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActivationMap {
    /// Activation time at each vertex (ms)
    pub activation_times: Vec<f32>,
    /// Conduction velocity (m/s)
    pub conduction_velocity: Vec<f32>,
}

/// Action potential duration map.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct APDMap {
    /// APD at 50% repolarization (ms)
    pub apd50: Vec<f32>,
    /// APD at 90% repolarization (ms)
    pub apd90: Vec<f32>,
    /// APD dispersion
    pub dispersion: f32,
}

/// Simulation statistics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationStats {
    /// Total computation time (s)
    pub compute_time: f32,
    /// Time steps computed
    pub time_steps: usize,
    /// Speedup vs traditional solver
    pub speedup_factor: f32,
    /// Neural operator inference time (ms)
    pub inference_time: f32,
}

// ============================================================================
// Arrhythmia Types
// ============================================================================

/// Arrhythmia event.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArrhythmiaEvent {
    /// Arrhythmia type
    pub arrhythmia_type: ArrhythmiaType,
    /// Start time (ms)
    pub start_time: f32,
    /// End time (ms)
    pub end_time: Option<f32>,
    /// Location
    pub location: Option<Point3D>,
    /// Severity (0-1)
    pub severity: f32,
}

/// Arrhythmia type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ArrhythmiaType {
    /// Normal sinus rhythm
    Normal,
    /// Premature ventricular contraction
    PVC,
    /// Atrial fibrillation
    AtrialFibrillation,
    /// Ventricular fibrillation
    VentricularFibrillation,
    /// Ventricular tachycardia
    VentricularTachycardia,
    /// Reentrant circuit
    Reentry,
    /// Conduction block
    Block,
    /// Spiral wave
    SpiralWave,
}

impl ArrhythmiaType {
    /// Get display name.
    #[must_use]
    pub fn display_name(&self) -> &'static str {
        match self {
            ArrhythmiaType::Normal => "Normal Sinus Rhythm",
            ArrhythmiaType::PVC => "Premature Ventricular Contraction",
            ArrhythmiaType::AtrialFibrillation => "Atrial Fibrillation",
            ArrhythmiaType::VentricularFibrillation => "Ventricular Fibrillation",
            ArrhythmiaType::VentricularTachycardia => "Ventricular Tachycardia",
            ArrhythmiaType::Reentry => "Reentrant Circuit",
            ArrhythmiaType::Block => "Conduction Block",
            ArrhythmiaType::SpiralWave => "Spiral Wave",
        }
    }
}

// ============================================================================
// Neural Operator Types
// ============================================================================

/// Neural operator configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NeuralOperatorConfig {
    /// Model type
    pub model_type: NeuralOperatorModel,
    /// Input channels
    pub input_channels: usize,
    /// Output channels
    pub output_channels: usize,
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Number of layers
    pub num_layers: usize,
    /// Time steps per inference
    pub time_steps: usize,
}

impl Default for NeuralOperatorConfig {
    fn default() -> Self {
        Self {
            model_type: NeuralOperatorModel::PINO,
            input_channels: 3,  // Voltage, recovery, stimulus
            output_channels: 2, // Voltage, recovery
            hidden_dim: 64,
            num_layers: 4,
            time_steps: 10,
        }
    }
}

/// Neural operator model type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum NeuralOperatorModel {
    /// Fourier Neural Operator
    FNO,
    /// Physics-Informed Neural Operator
    PINO,
    /// Deep Operator Network
    DeepONet,
    /// Convolutional Operator
    ConvOperator,
}

/// Neural operator inference result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OperatorInference {
    /// Predicted voltage field
    pub voltage: Vec<f32>,
    /// Predicted recovery field
    pub recovery: Vec<f32>,
    /// Uncertainty estimate
    pub uncertainty: Vec<f32>,
    /// Physics residual (for PINO)
    pub physics_residual: f32,
}

// ============================================================================
// Analysis Types
// ============================================================================

/// ECG analysis request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ECGRequest {
    /// Simulation result
    pub simulation: SimulationResult,
    /// Lead configuration
    pub leads: ECGLeads,
}

/// ECG lead configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ECGLeads {
    /// Standard 12-lead ECG
    pub standard_12_lead: bool,
    /// Additional electrode positions
    pub custom_electrodes: Vec<Point3D>,
}

impl Default for ECGLeads {
    fn default() -> Self {
        Self {
            standard_12_lead: true,
            custom_electrodes: vec![],
        }
    }
}

/// ECG result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ECGResult {
    /// Time points (ms)
    pub times: Vec<f32>,
    /// Lead I voltage
    pub lead_i: Vec<f32>,
    /// Lead II voltage
    pub lead_ii: Vec<f32>,
    /// Lead III voltage
    pub lead_iii: Vec<f32>,
    /// Precordial leads (V1-V6)
    pub precordial: Vec<Vec<f32>>,
    /// Computed heart rate (bpm)
    pub heart_rate: f32,
    /// QT interval (ms)
    pub qt_interval: f32,
    /// QRS duration (ms)
    pub qrs_duration: f32,
}

// ============================================================================
// Sample Data
// ============================================================================

/// Get sample heart mesh.
#[must_use]
pub fn get_sample_heart_mesh() -> HeartMesh {
    // Create a simplified ellipsoidal heart
    let mut vertices = Vec::new();
    let mut fibers = Vec::new();
    let mut sheets = Vec::new();
    let mut vertex_regions = Vec::new();

    // Generate vertices on ellipsoid
    let n_long = 20;
    let n_lat = 10;

    for i in 0..n_long {
        let phi = 2.0 * std::f32::consts::PI * (i as f32) / (n_long as f32);
        for j in 0..n_lat {
            let theta = std::f32::consts::PI * (j as f32) / (n_lat as f32);

            // Ellipsoid radii
            let a = 30.0; // x
            let b = 25.0; // y
            let c = 50.0; // z (apex to base)

            let x = a * theta.sin() * phi.cos();
            let y = b * theta.sin() * phi.sin();
            let z = c * theta.cos();

            vertices.push(Point3D::new(x, y, z));

            // Fiber direction (circumferential with transmural rotation)
            let fiber = Vector3D::new(-phi.sin(), phi.cos(), 0.0).normalize();
            fibers.push(fiber);

            // Sheet direction (radial)
            let sheet = Vector3D::new(x, y, 0.0).normalize();
            sheets.push(sheet);

            // Assign regions
            let region = if z > 25.0 {
                0 // Atria
            } else if x > 0.0 {
                2 // Left ventricle
            } else {
                3 // Right ventricle
            };
            vertex_regions.push(region);
        }
    }

    // Generate triangles
    let mut triangles = Vec::new();
    for i in 0..n_long {
        for j in 0..n_lat - 1 {
            let idx = |ii: usize, jj: usize| (ii % n_long) * n_lat + jj;

            triangles.push([idx(i, j), idx(i + 1, j), idx(i + 1, j + 1)]);
            triangles.push([idx(i, j), idx(i + 1, j + 1), idx(i, j + 1)]);
        }
    }

    HeartMesh {
        name: "Sample Heart".to_string(),
        vertices,
        triangles,
        tetrahedra: None,
        fibers,
        sheets,
        regions: vec![
            HeartRegion::LeftAtrium,
            HeartRegion::RightAtrium,
            HeartRegion::LeftVentricle,
            HeartRegion::RightVentricle,
        ],
        vertex_regions,
    }
}

/// Get sample stimulation protocol.
#[must_use]
pub fn get_sample_protocol() -> StimulationProtocol {
    StimulationProtocol {
        sites: vec![StimulationSite {
            name: "SA Node".to_string(),
            center: Point3D::new(5.0, 15.0, 40.0),
            radius: 5.0,
            current: 100.0,
            times: vec![0.0],
            duration: 1.0,
        }],
        protocol_type: ProtocolType::SinusRhythm,
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_point3d() {
        let p1 = Point3D::new(0.0, 0.0, 0.0);
        let p2 = Point3D::new(3.0, 4.0, 0.0);
        assert!((p1.distance_to(&p2) - 5.0).abs() < 0.001);
    }

    #[test]
    fn test_vector3d_normalize() {
        let v = Vector3D::new(3.0, 4.0, 0.0);
        let n = v.normalize();
        let mag = (n.x * n.x + n.y * n.y + n.z * n.z).sqrt();
        assert!((mag - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_heart_region() {
        assert_eq!(HeartRegion::LeftVentricle.display_name(), "Left Ventricle");
        assert_eq!(HeartRegion::SANode.display_name(), "SA Node");
    }

    #[test]
    fn test_ionic_model_default() {
        let model = IonicModel::default();
        assert_eq!(model, IonicModel::MitchellSchaeffer);
    }

    #[test]
    fn test_conductivity_default() {
        let cond = ConductivityParams::default();
        assert!(cond.sigma_l > cond.sigma_t);
    }

    #[test]
    fn test_ap_state_default() {
        let state = ActionPotentialState::default();
        assert!(state.voltage < -80.0); // Resting potential
    }

    #[test]
    fn test_sim_config_default() {
        let config = SimulationConfig::default();
        assert!(config.dt < 0.1);
        assert!(config.total_time > 100.0);
    }

    #[test]
    fn test_neural_operator_config() {
        let config = NeuralOperatorConfig::default();
        assert_eq!(config.model_type, NeuralOperatorModel::PINO);
    }

    #[test]
    fn test_arrhythmia_type() {
        assert_eq!(
            ArrhythmiaType::VentricularFibrillation.display_name(),
            "Ventricular Fibrillation"
        );
    }

    #[test]
    fn test_sample_heart_mesh() {
        let mesh = get_sample_heart_mesh();
        assert!(!mesh.vertices.is_empty());
        assert!(!mesh.triangles.is_empty());
        assert_eq!(mesh.vertices.len(), mesh.fibers.len());
    }

    #[test]
    fn test_sample_protocol() {
        let protocol = get_sample_protocol();
        assert!(!protocol.sites.is_empty());
        assert_eq!(protocol.protocol_type, ProtocolType::SinusRhythm);
    }

    #[test]
    fn test_serialization() {
        let mesh = get_sample_heart_mesh();
        let json = serde_json::to_string(&mesh).unwrap();
        assert!(json.contains("Sample Heart"));
    }
}