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

//! Sample data for cardiac simulation demo.
//!
//! Provides sample heart meshes and stimulation protocols.

use cardiosim_shared::{
    HeartMesh, HeartRegion, Point3D, ProtocolType, StimulationProtocol, StimulationSite, Vector3D,
};

/// Get a detailed heart mesh.
#[must_use]
pub fn get_detailed_heart_mesh() -> HeartMesh {
    // Create a more detailed heart with proper chamber structure
    let mut vertices = Vec::new();
    let mut fibers = Vec::new();
    let mut sheets = Vec::new();
    let mut vertex_regions = Vec::new();
    let mut triangles = Vec::new();

    // Left ventricle (ellipsoid)
    let lv_base = 20.0;
    let lv_apex = -30.0;
    let lv_radius = 25.0;

    let n_circ = 24;
    let n_long = 16;

    for i in 0..n_long {
        let z = lv_base + (lv_apex - lv_base) * (i as f32) / (n_long as f32 - 1.0);
        let radius_factor = (1.0 - ((z - lv_base) / (lv_apex - lv_base)).abs()).sqrt();
        let r = lv_radius * radius_factor;

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

            let x = r * theta.cos();
            let y = r * theta.sin();

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

            // Fiber direction (helical, rotating with depth)
            let helix_angle = std::f32::consts::PI / 6.0 * (0.5 - i as f32 / n_long as f32);
            let fiber = Vector3D::new(
                -theta.sin() * helix_angle.cos() + theta.cos() * helix_angle.sin(),
                theta.cos() * helix_angle.cos() + theta.sin() * helix_angle.sin(),
                helix_angle.sin(),
            )
            .normalize();
            fibers.push(fiber);

            // Sheet direction (radial)
            sheets.push(Vector3D::new(theta.cos(), theta.sin(), 0.0).normalize());

            // Region assignment
            vertex_regions.push(2); // Left ventricle
        }
    }

    // Right ventricle (crescent attached to LV)
    let rv_start_idx = vertices.len();
    let rv_n_circ = 12;
    let rv_n_long = 12;

    for i in 0..rv_n_long {
        let z = lv_base + (lv_apex - lv_base) * (i as f32) / (rv_n_long as f32 - 1.0);
        let radius_factor = (1.0 - ((z - lv_base) / (lv_apex - lv_base)).abs()).sqrt();

        for j in 0..rv_n_circ {
            // RV wraps around LV from -90 to -270 degrees
            let theta = -std::f32::consts::PI / 2.0
                - std::f32::consts::PI * (j as f32) / (rv_n_circ as f32 - 1.0);

            let r = (lv_radius + 10.0) * radius_factor;
            let x = r * theta.cos();
            let y = r * theta.sin();

            vertices.push(Point3D::new(x, y, z));
            fibers.push(Vector3D::new(-theta.sin(), theta.cos(), 0.0).normalize());
            sheets.push(Vector3D::new(theta.cos(), theta.sin(), 0.0).normalize());
            vertex_regions.push(3); // Right ventricle
        }
    }

    // Left atrium (sphere above LV)
    let la_center = Point3D::new(-15.0, 0.0, 35.0);
    let la_radius = 18.0;
    let la_start_idx = vertices.len();
    let la_n_theta = 8;
    let la_n_phi = 12;

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

            let x = la_center.x + la_radius * theta.sin() * phi.cos();
            let y = la_center.y + la_radius * theta.sin() * phi.sin();
            let z = la_center.z + la_radius * theta.cos();

            vertices.push(Point3D::new(x, y, z));
            fibers.push(Vector3D::new(-phi.sin(), phi.cos(), 0.0).normalize());
            sheets.push(Vector3D::new(phi.cos(), phi.sin(), 0.0).normalize());
            vertex_regions.push(0); // Left atrium
        }
    }

    // Right atrium (sphere above RV)
    let ra_center = Point3D::new(15.0, -10.0, 38.0);
    let ra_radius = 16.0;

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

            let x = ra_center.x + ra_radius * theta.sin() * phi.cos();
            let y = ra_center.y + ra_radius * theta.sin() * phi.sin();
            let z = ra_center.z + ra_radius * theta.cos();

            vertices.push(Point3D::new(x, y, z));
            fibers.push(Vector3D::new(-phi.sin(), phi.cos(), 0.0).normalize());
            sheets.push(Vector3D::new(phi.cos(), phi.sin(), 0.0).normalize());
            vertex_regions.push(1); // Right atrium
        }
    }

    // Generate triangles for LV
    for i in 0..n_long - 1 {
        for j in 0..n_circ {
            let idx = |ii: usize, jj: usize| ii * n_circ + (jj % n_circ);
            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)]);
        }
    }

    // Generate triangles for RV
    for i in 0..rv_n_long - 1 {
        for j in 0..rv_n_circ - 1 {
            let idx = |ii: usize, jj: usize| rv_start_idx + ii * rv_n_circ + 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)]);
        }
    }

    // Generate triangles for atria
    for (start_idx, n_phi) in [(la_start_idx, la_n_phi)] {
        for i in 0..la_n_theta - 2 {
            for j in 0..n_phi {
                let idx = |ii: usize, jj: usize| start_idx + ii * n_phi + (jj % n_phi);
                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: "Detailed Human Heart".to_string(),
        vertices,
        triangles,
        tetrahedra: None,
        fibers,
        sheets,
        regions: vec![
            HeartRegion::LeftAtrium,
            HeartRegion::RightAtrium,
            HeartRegion::LeftVentricle,
            HeartRegion::RightVentricle,
        ],
        vertex_regions,
    }
}

/// Get S1-S2 restitution protocol.
#[must_use]
pub fn get_s1s2_protocol() -> StimulationProtocol {
    // S1 at fixed cycle length, S2 at decreasing intervals
    let s1_times: Vec<f32> = (0..8).map(|i| i as f32 * 500.0).collect();

    let mut sites = vec![StimulationSite {
        name: "S1 - Apex".to_string(),
        center: Point3D::new(0.0, 0.0, -25.0),
        radius: 8.0,
        current: 150.0,
        times: s1_times,
        duration: 2.0,
    }];

    // S2 at varying coupling intervals
    for (i, coupling) in [400.0, 350.0, 300.0, 280.0, 260.0].iter().enumerate() {
        sites.push(StimulationSite {
            name: format!("S2 - CI {coupling}ms"),
            center: Point3D::new(0.0, 0.0, -25.0),
            radius: 8.0,
            current: 200.0,
            times: vec![3500.0 + i as f32 * 2000.0 + coupling],
            duration: 2.0,
        });
    }

    StimulationProtocol {
        sites,
        protocol_type: ProtocolType::S1S2,
    }
}

/// Get burst pacing protocol (to induce arrhythmia).
#[must_use]
pub fn get_burst_pacing_protocol() -> StimulationProtocol {
    // Rapid pacing at 50ms cycle length
    let burst_times: Vec<f32> = (0..20).map(|i| i as f32 * 50.0).collect();

    StimulationProtocol {
        sites: vec![StimulationSite {
            name: "Burst - RV apex".to_string(),
            center: Point3D::new(15.0, -10.0, -20.0),
            radius: 5.0,
            current: 200.0,
            times: burst_times,
            duration: 1.0,
        }],
        protocol_type: ProtocolType::BurstPacing,
    }
}

/// Get cross-field stimulation protocol.
#[must_use]
pub fn get_cross_field_protocol() -> StimulationProtocol {
    StimulationProtocol {
        sites: vec![
            StimulationSite {
                name: "S1 - Endocardial line".to_string(),
                center: Point3D::new(0.0, 0.0, 0.0),
                radius: 30.0,
                current: 150.0,
                times: vec![0.0],
                duration: 2.0,
            },
            StimulationSite {
                name: "S2 - Epicardial field".to_string(),
                center: Point3D::new(0.0, 25.0, 0.0),
                radius: 20.0,
                current: 100.0,
                times: vec![200.0],
                duration: 5.0,
            },
        ],
        protocol_type: ProtocolType::CrossField,
    }
}

/// Get SA node pacing protocol.
#[must_use]
pub fn get_sinus_rhythm_protocol() -> StimulationProtocol {
    // Normal sinus rhythm at 75 bpm (800ms cycle)
    let sinus_times: Vec<f32> = (0..10).map(|i| i as f32 * 800.0).collect();

    StimulationProtocol {
        sites: vec![StimulationSite {
            name: "SA Node".to_string(),
            center: Point3D::new(20.0, -5.0, 45.0),
            radius: 3.0,
            current: 50.0,
            times: sinus_times,
            duration: 1.0,
        }],
        protocol_type: ProtocolType::SinusRhythm,
    }
}

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

    #[test]
    fn test_detailed_heart_mesh() {
        let mesh = get_detailed_heart_mesh();
        assert!(!mesh.vertices.is_empty());
        assert!(!mesh.triangles.is_empty());
        assert_eq!(mesh.vertices.len(), mesh.fibers.len());
        assert_eq!(mesh.vertices.len(), mesh.sheets.len());
        assert_eq!(mesh.vertices.len(), mesh.vertex_regions.len());
    }

    #[test]
    fn test_s1s2_protocol() {
        let protocol = get_s1s2_protocol();
        assert!(protocol.sites.len() >= 2);
        assert_eq!(protocol.protocol_type, ProtocolType::S1S2);
    }

    #[test]
    fn test_burst_pacing_protocol() {
        let protocol = get_burst_pacing_protocol();
        assert!(!protocol.sites.is_empty());
        assert_eq!(protocol.protocol_type, ProtocolType::BurstPacing);
        // Should have many rapid stimuli
        assert!(protocol.sites[0].times.len() >= 10);
    }

    #[test]
    fn test_cross_field_protocol() {
        let protocol = get_cross_field_protocol();
        assert_eq!(protocol.sites.len(), 2);
        assert_eq!(protocol.protocol_type, ProtocolType::CrossField);
    }

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

        // Check ~800ms cycle length
        let times = &protocol.sites[0].times;
        if times.len() >= 2 {
            let interval = times[1] - times[0];
            assert!((interval - 800.0).abs() < 10.0);
        }
    }

    #[test]
    fn test_fiber_directions_normalized() {
        let mesh = get_detailed_heart_mesh();
        for fiber in &mesh.fibers {
            let mag = (fiber.x * fiber.x + fiber.y * fiber.y + fiber.z * fiber.z).sqrt();
            assert!((mag - 1.0).abs() < 0.01);
        }
    }

    #[test]
    fn test_region_assignments() {
        let mesh = get_detailed_heart_mesh();
        // All vertex regions should be valid indices
        for &region in &mesh.vertex_regions {
            assert!(region < mesh.regions.len());
        }
    }
}