//! `CardioSim` cardiac electrophysiology simulation demo.
//!
//! This crate implements cardiac electrophysiology simulation using
//! physics-informed neural operators (PINO).
//!
//! # Architecture
//!
//! The simulation consists of:
//! - **Monodomain Solver**: Traditional PDE solver for comparison
//! - **Neural Operator**: PINO for fast inference
//! - **Ionic Models**: Mitchell-Schaeffer, FitzHugh-Nagumo, etc.
//! - **ECG Computation**: Virtual ECG from voltage fields
//!
//! # Example
//!
//! ```ignore
//! use rtx_cardiosim_demo::{run_simulation, SimulationRequest};
//!
//! let result = run_simulation(request).await?;
//! println!("Computed {} time steps in {:.2}s",
//!     result.stats.time_steps,
//!     result.stats.compute_time);
//! ```

pub mod ecg;
pub mod ionic_models;
pub mod monodomain;
pub mod neural_operator;
pub mod sample_data;

use cardiosim_shared::{
    APDMap, ActivationMap, ArrhythmiaEvent, ArrhythmiaType, ECGRequest, ECGResult, HeartMesh,
    NeuralOperatorConfig, Point3D, SimulationConfig, SimulationRequest, SimulationResult,
    SimulationStats, VoltageField,
};
use thiserror::Error;

/// Errors that can occur during cardiac simulation.
#[derive(Debug, Error)]
pub enum CardioSimError {
    /// Invalid mesh
    #[error("Invalid mesh: {0}")]
    InvalidMesh(String),

    /// Simulation error
    #[error("Simulation error: {0}")]
    SimulationError(String),

    /// Neural operator error
    #[error("Neural operator error: {0}")]
    NeuralOperatorError(String),

    /// Configuration error
    #[error("Configuration error: {0}")]
    ConfigError(String),
}

/// Run cardiac electrophysiology simulation.
pub async fn run_simulation(
    request: SimulationRequest,
) -> Result<SimulationResult, CardioSimError> {
    use std::time::Instant;

    let start = Instant::now();

    // Validate mesh
    validate_mesh(&request.mesh)?;

    tracing::info!(
        "Running cardiac simulation on mesh with {} vertices",
        request.mesh.vertices.len()
    );

    // Choose solver based on configuration
    let (voltage_fields, speedup_factor) = if request.config.use_neural_operator {
        let operator = neural_operator::CardiacPINO::new(NeuralOperatorConfig::default());
        let fields = operator.solve(&request.mesh, &request.protocol, &request.config)?;
        (fields, 1000.0) // Neural operators are ~1000x faster
    } else {
        let solver = monodomain::MonodomainSolver::new(&request.config);
        let fields = solver.solve(&request.mesh, &request.protocol)?;
        (fields, 1.0)
    };

    // Compute activation and APD maps
    let activation_map = compute_activation_map(&voltage_fields);
    let apd_map = compute_apd_map(&voltage_fields);

    // Detect arrhythmias
    let arrhythmias = detect_arrhythmias(&voltage_fields, &activation_map);

    // Extract time points
    let times: Vec<f32> = voltage_fields.iter().map(|f| f.time).collect();

    let elapsed = start.elapsed();

    let stats = SimulationStats {
        compute_time: elapsed.as_secs_f32(),
        time_steps: voltage_fields.len(),
        speedup_factor,
        inference_time: elapsed.as_millis() as f32 / voltage_fields.len() as f32,
    };

    Ok(SimulationResult {
        times,
        voltage_fields,
        activation_map,
        apd_map,
        arrhythmias,
        stats,
    })
}

/// Compute ECG from simulation result.
pub async fn compute_ecg(request: ECGRequest) -> Result<ECGResult, CardioSimError> {
    let ecg_computer = ecg::ECGComputer::new(&request.leads);
    ecg_computer.compute(&request.simulation)
}

/// Validate mesh.
fn validate_mesh(mesh: &HeartMesh) -> Result<(), CardioSimError> {
    if mesh.vertices.is_empty() {
        return Err(CardioSimError::InvalidMesh(
            "Mesh has no vertices".to_string(),
        ));
    }

    if mesh.triangles.is_empty() {
        return Err(CardioSimError::InvalidMesh(
            "Mesh has no triangles".to_string(),
        ));
    }

    if mesh.fibers.len() != mesh.vertices.len() {
        return Err(CardioSimError::InvalidMesh(
            "Fiber directions don't match vertex count".to_string(),
        ));
    }

    Ok(())
}

/// Compute activation time map.
fn compute_activation_map(voltage_fields: &[VoltageField]) -> ActivationMap {
    if voltage_fields.is_empty() {
        return ActivationMap {
            activation_times: vec![],
            conduction_velocity: vec![],
        };
    }

    let n_vertices = voltage_fields[0].voltages.len();
    let mut activation_times = vec![f32::MAX; n_vertices];

    // Find first time each vertex crosses -30 mV (activation threshold)
    let threshold = -30.0;

    for field in voltage_fields {
        for (i, &voltage) in field.voltages.iter().enumerate() {
            if voltage > threshold && activation_times[i] == f32::MAX {
                activation_times[i] = field.time;
            }
        }
    }

    // Replace unactivated vertices with NaN
    for t in &mut activation_times {
        if *t == f32::MAX {
            *t = f32::NAN;
        }
    }

    // Compute conduction velocity (simplified)
    let conduction_velocity = vec![0.8; n_vertices]; // ~0.8 m/s typical

    ActivationMap {
        activation_times,
        conduction_velocity,
    }
}

/// Compute APD map.
fn compute_apd_map(voltage_fields: &[VoltageField]) -> APDMap {
    if voltage_fields.is_empty() {
        return APDMap {
            apd50: vec![],
            apd90: vec![],
            dispersion: 0.0,
        };
    }

    let n_vertices = voltage_fields[0].voltages.len();
    let mut apd50 = vec![0.0; n_vertices];
    let mut apd90 = vec![0.0; n_vertices];

    // Track activation and repolarization times
    let mut activation_times = vec![f32::MAX; n_vertices];
    let mut repol50_times = vec![f32::MAX; n_vertices];
    let mut repol90_times = vec![f32::MAX; n_vertices];
    let mut max_voltages = vec![f32::MIN; n_vertices];

    for field in voltage_fields {
        for (i, &voltage) in field.voltages.iter().enumerate() {
            // Track max voltage
            if voltage > max_voltages[i] {
                max_voltages[i] = voltage;
            }

            // Activation (upstroke)
            if voltage > -30.0 && activation_times[i] == f32::MAX {
                activation_times[i] = field.time;
            }

            // Repolarization at 50%
            if activation_times[i] < f32::MAX && repol50_times[i] == f32::MAX {
                let v_max = max_voltages[i];
                let v_rest = -85.0;
                let v_50 = v_rest + 0.5 * (v_max - v_rest);
                if voltage < v_50 {
                    repol50_times[i] = field.time;
                }
            }

            // Repolarization at 90%
            if activation_times[i] < f32::MAX && repol90_times[i] == f32::MAX {
                let v_max = max_voltages[i];
                let v_rest = -85.0;
                let v_90 = v_rest + 0.1 * (v_max - v_rest);
                if voltage < v_90 {
                    repol90_times[i] = field.time;
                }
            }
        }
    }

    // Calculate APDs
    for i in 0..n_vertices {
        if activation_times[i] < f32::MAX {
            if repol50_times[i] < f32::MAX {
                apd50[i] = repol50_times[i] - activation_times[i];
            }
            if repol90_times[i] < f32::MAX {
                apd90[i] = repol90_times[i] - activation_times[i];
            }
        }
    }

    // Calculate dispersion
    let valid_apd90: Vec<f32> = apd90.iter().copied().filter(|&x| x > 0.0).collect();
    let dispersion = if valid_apd90.len() > 1 {
        let mean: f32 = valid_apd90.iter().sum::<f32>() / valid_apd90.len() as f32;
        let variance: f32 =
            valid_apd90.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / valid_apd90.len() as f32;
        variance.sqrt()
    } else {
        0.0
    };

    APDMap {
        apd50,
        apd90,
        dispersion,
    }
}

/// Detect arrhythmias in simulation.
fn detect_arrhythmias(
    voltage_fields: &[VoltageField],
    activation_map: &ActivationMap,
) -> Vec<ArrhythmiaEvent> {
    let mut events = Vec::new();

    if voltage_fields.is_empty() {
        return events;
    }

    // Check for normal rhythm
    let max_activation = activation_map
        .activation_times
        .iter()
        .filter(|&&t| t.is_finite())
        .fold(0.0_f32, |a, &b| a.max(b));

    let min_activation = activation_map
        .activation_times
        .iter()
        .filter(|&&t| t.is_finite())
        .fold(f32::MAX, |a, &b| a.min(b));

    let activation_spread = max_activation - min_activation;

    // Normal activation should complete within ~150ms
    if activation_spread < 150.0 {
        events.push(ArrhythmiaEvent {
            arrhythmia_type: ArrhythmiaType::Normal,
            start_time: 0.0,
            end_time: Some(voltage_fields.last().map_or(0.0, |f| f.time)),
            location: None,
            severity: 0.0,
        });
    } else if activation_spread > 500.0 {
        // Likely reentry
        events.push(ArrhythmiaEvent {
            arrhythmia_type: ArrhythmiaType::Reentry,
            start_time: min_activation,
            end_time: None,
            location: Some(Point3D::new(0.0, 0.0, 0.0)),
            severity: 0.7,
        });
    } else if activation_spread > 300.0 {
        // Possible conduction abnormality
        events.push(ArrhythmiaEvent {
            arrhythmia_type: ArrhythmiaType::Block,
            start_time: min_activation,
            end_time: Some(max_activation),
            location: None,
            severity: 0.4,
        });
    }

    events
}

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

    #[tokio::test]
    async fn test_run_simulation() {
        let mesh = cardiosim_shared::get_sample_heart_mesh();
        let protocol = cardiosim_shared::get_sample_protocol();
        let config = SimulationConfig::default();

        let request = SimulationRequest {
            mesh,
            config,
            protocol,
        };

        let result = run_simulation(request).await;
        assert!(result.is_ok());

        let sim = result.unwrap();
        assert!(!sim.voltage_fields.is_empty());
        assert!(!sim.activation_map.activation_times.is_empty());
    }

    #[tokio::test]
    async fn test_neural_operator_simulation() {
        let mesh = cardiosim_shared::get_sample_heart_mesh();
        let protocol = cardiosim_shared::get_sample_protocol();
        let config = SimulationConfig {
            use_neural_operator: true,
            total_time: 100.0,
            ..Default::default()
        };

        let request = SimulationRequest {
            mesh,
            config,
            protocol,
        };

        let result = run_simulation(request).await;
        assert!(result.is_ok());

        let sim = result.unwrap();
        assert!(sim.stats.speedup_factor > 1.0);
    }

    #[test]
    fn test_validate_mesh() {
        let mesh = cardiosim_shared::get_sample_heart_mesh();
        assert!(validate_mesh(&mesh).is_ok());
    }

    #[test]
    fn test_validate_empty_mesh() {
        let mesh = HeartMesh {
            name: "Empty".to_string(),
            vertices: vec![],
            triangles: vec![],
            tetrahedra: None,
            fibers: vec![],
            sheets: vec![],
            regions: vec![],
            vertex_regions: vec![],
        };
        assert!(validate_mesh(&mesh).is_err());
    }

    #[test]
    fn test_compute_activation_map() {
        let fields = vec![
            VoltageField {
                time: 0.0,
                voltages: vec![-85.0, -85.0, -85.0],
            },
            VoltageField {
                time: 10.0,
                voltages: vec![20.0, -85.0, -85.0],
            },
            VoltageField {
                time: 20.0,
                voltages: vec![-20.0, 20.0, -85.0],
            },
        ];

        let map = compute_activation_map(&fields);
        assert!((map.activation_times[0] - 10.0).abs() < 0.001);
        assert!((map.activation_times[1] - 20.0).abs() < 0.001);
    }

    #[test]
    fn test_detect_arrhythmias() {
        let fields = vec![
            VoltageField {
                time: 0.0,
                voltages: vec![-85.0; 100],
            },
            VoltageField {
                time: 50.0,
                voltages: vec![20.0; 100],
            },
        ];

        let activation_map = compute_activation_map(&fields);
        let events = detect_arrhythmias(&fields, &activation_map);

        assert!(!events.is_empty());
    }
}