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

//! Sample data and configurations for FederatedMed demo.
//!
//! This module provides sample federation configurations, client lists,
//! and privacy settings for different medical imaging scenarios.

use fedmed_shared::{
    AggregationStrategy, ClientInfo, ClientType, DataDistribution, FederatedConfig, MedicalDataset,
    MedicalTask, PrivacyAccountantType, PrivacyConfig,
};

// ============================================================================
// Sample Federations
// ============================================================================

/// Create a chest X-ray classification federation.
///
/// This simulates a multi-hospital network for detecting pneumonia,
/// COVID-19, and other pulmonary conditions from chest X-rays.
#[must_use]
pub fn chest_xray_federation() -> Vec<ClientInfo> {
    vec![
        ClientInfo {
            id: "metro_general".to_string(),
            name: "Metro General Hospital".to_string(),
            data_size: 15000,
            is_active: true,
            client_type: ClientType::Hospital,
            data_distribution: DataDistribution {
                class_counts: vec![5000, 4000, 3000, 2000, 1000],
                is_iid: false,
                heterogeneity: 0.25,
            },
            last_seen: 1704067200,
            latency_ms: 35,
        },
        ClientInfo {
            id: "university_med".to_string(),
            name: "University Medical Center".to_string(),
            data_size: 25000,
            is_active: true,
            client_type: ClientType::Hospital,
            data_distribution: DataDistribution {
                class_counts: vec![8000, 6000, 5000, 4000, 2000],
                is_iid: false,
                heterogeneity: 0.15,
            },
            last_seen: 1704067200,
            latency_ms: 25,
        },
        ClientInfo {
            id: "community_health".to_string(),
            name: "Community Health Center".to_string(),
            data_size: 8000,
            is_active: true,
            client_type: ClientType::Clinic,
            data_distribution: DataDistribution {
                class_counts: vec![3000, 2500, 1500, 700, 300],
                is_iid: false,
                heterogeneity: 0.35,
            },
            last_seen: 1704067200,
            latency_ms: 55,
        },
        ClientInfo {
            id: "regional_radiology".to_string(),
            name: "Regional Radiology Associates".to_string(),
            data_size: 20000,
            is_active: true,
            client_type: ClientType::Radiology,
            data_distribution: DataDistribution {
                class_counts: vec![6000, 5500, 4500, 2500, 1500],
                is_iid: false,
                heterogeneity: 0.2,
            },
            last_seen: 1704067200,
            latency_ms: 30,
        },
        ClientInfo {
            id: "research_institute".to_string(),
            name: "National Health Research Institute".to_string(),
            data_size: 12000,
            is_active: true,
            client_type: ClientType::Research,
            data_distribution: DataDistribution {
                class_counts: vec![2400, 2400, 2400, 2400, 2400],
                is_iid: true,
                heterogeneity: 0.05,
            },
            last_seen: 1704067200,
            latency_ms: 20,
        },
    ]
}

/// Create a skin lesion classification federation.
///
/// This simulates a dermatology network for melanoma detection
/// and skin condition classification.
#[must_use]
pub fn skin_lesion_federation() -> Vec<ClientInfo> {
    vec![
        ClientInfo {
            id: "derm_clinic_north".to_string(),
            name: "Northern Dermatology Clinic".to_string(),
            data_size: 5000,
            is_active: true,
            client_type: ClientType::Clinic,
            data_distribution: DataDistribution {
                class_counts: vec![2000, 1500, 800, 500, 200],
                is_iid: false,
                heterogeneity: 0.3,
            },
            last_seen: 1704067200,
            latency_ms: 40,
        },
        ClientInfo {
            id: "skin_cancer_center".to_string(),
            name: "Skin Cancer Treatment Center".to_string(),
            data_size: 8000,
            is_active: true,
            client_type: ClientType::Hospital,
            data_distribution: DataDistribution {
                class_counts: vec![1000, 2000, 2500, 1500, 1000],
                is_iid: false,
                heterogeneity: 0.4,
            },
            last_seen: 1704067200,
            latency_ms: 35,
        },
        ClientInfo {
            id: "derm_research".to_string(),
            name: "Dermatology Research Lab".to_string(),
            data_size: 10000,
            is_active: true,
            client_type: ClientType::Research,
            data_distribution: DataDistribution {
                class_counts: vec![2000, 2000, 2000, 2000, 2000],
                is_iid: true,
                heterogeneity: 0.1,
            },
            last_seen: 1704067200,
            latency_ms: 25,
        },
        ClientInfo {
            id: "derm_clinic_south".to_string(),
            name: "Southern Dermatology Associates".to_string(),
            data_size: 4000,
            is_active: true,
            client_type: ClientType::Clinic,
            data_distribution: DataDistribution {
                class_counts: vec![1800, 1000, 600, 400, 200],
                is_iid: false,
                heterogeneity: 0.35,
            },
            last_seen: 1704067200,
            latency_ms: 50,
        },
    ]
}

/// Create a retinal disease detection federation.
#[must_use]
pub fn retinal_disease_federation() -> Vec<ClientInfo> {
    vec![
        ClientInfo {
            id: "eye_institute".to_string(),
            name: "National Eye Institute".to_string(),
            data_size: 30000,
            is_active: true,
            client_type: ClientType::Research,
            data_distribution: DataDistribution {
                class_counts: vec![10000, 8000, 6000, 4000, 2000],
                is_iid: false,
                heterogeneity: 0.2,
            },
            last_seen: 1704067200,
            latency_ms: 20,
        },
        ClientInfo {
            id: "vision_care".to_string(),
            name: "Vision Care Hospital".to_string(),
            data_size: 15000,
            is_active: true,
            client_type: ClientType::Hospital,
            data_distribution: DataDistribution {
                class_counts: vec![5000, 4000, 3000, 2000, 1000],
                is_iid: false,
                heterogeneity: 0.25,
            },
            last_seen: 1704067200,
            latency_ms: 30,
        },
        ClientInfo {
            id: "optometry_network".to_string(),
            name: "Regional Optometry Network".to_string(),
            data_size: 20000,
            is_active: true,
            client_type: ClientType::Clinic,
            data_distribution: DataDistribution {
                class_counts: vec![7000, 5500, 4000, 2500, 1000],
                is_iid: false,
                heterogeneity: 0.3,
            },
            last_seen: 1704067200,
            latency_ms: 45,
        },
    ]
}

// ============================================================================
// Sample Configurations
// ============================================================================

/// Create a standard hospital network configuration.
#[must_use]
pub fn hospital_network_config() -> FederatedConfig {
    FederatedConfig {
        num_rounds: 10, // Reduced for demo to fit within privacy budget
        local_epochs: 5,
        batch_size: 32,
        aggregation_strategy: AggregationStrategy::FedAvg,
        learning_rate: 0.01,
        min_clients: 3,
        client_fraction: 0.5,
        secure_aggregation: true,
        privacy: Some(moderate_privacy_config()),
    }
}

/// Create a high-performance configuration for powerful compute.
#[must_use]
pub fn high_performance_config() -> FederatedConfig {
    FederatedConfig {
        num_rounds: 200,
        local_epochs: 10,
        batch_size: 64,
        aggregation_strategy: AggregationStrategy::FedAvg,
        learning_rate: 0.001,
        min_clients: 5,
        client_fraction: 0.8,
        secure_aggregation: true,
        privacy: Some(moderate_privacy_config()),
    }
}

/// Create a resource-constrained configuration.
#[must_use]
pub fn resource_constrained_config() -> FederatedConfig {
    FederatedConfig {
        num_rounds: 50,
        local_epochs: 2,
        batch_size: 16,
        aggregation_strategy: AggregationStrategy::FedAvg,
        learning_rate: 0.01,
        min_clients: 2,
        client_fraction: 0.3,
        secure_aggregation: false,
        privacy: Some(relaxed_privacy_config()),
    }
}

/// Create a FedProx configuration for heterogeneous data.
#[must_use]
pub fn fedprox_heterogeneous_config() -> FederatedConfig {
    FederatedConfig {
        num_rounds: 100,
        local_epochs: 5,
        batch_size: 32,
        aggregation_strategy: AggregationStrategy::FedProx,
        learning_rate: 0.01,
        min_clients: 3,
        client_fraction: 0.5,
        secure_aggregation: true,
        privacy: Some(moderate_privacy_config()),
    }
}

/// Create a Scaffold configuration for drift correction.
#[must_use]
pub fn scaffold_config() -> FederatedConfig {
    FederatedConfig {
        num_rounds: 100,
        local_epochs: 5,
        batch_size: 32,
        aggregation_strategy: AggregationStrategy::Scaffold,
        learning_rate: 0.01,
        min_clients: 3,
        client_fraction: 0.5,
        secure_aggregation: true,
        privacy: Some(moderate_privacy_config()),
    }
}

/// Create a Byzantine-resilient configuration.
#[must_use]
pub fn byzantine_resilient_config() -> FederatedConfig {
    FederatedConfig {
        num_rounds: 100,
        local_epochs: 5,
        batch_size: 32,
        aggregation_strategy: AggregationStrategy::Krum,
        learning_rate: 0.01,
        min_clients: 5, // Need more clients for Byzantine resilience
        client_fraction: 0.8,
        secure_aggregation: true,
        privacy: Some(strict_privacy_config()),
    }
}

// ============================================================================
// Privacy Configurations
// ============================================================================

/// Create a strict privacy configuration.
///
/// Low epsilon (high privacy) suitable for highly sensitive data.
#[must_use]
pub fn strict_privacy_config() -> PrivacyConfig {
    PrivacyConfig {
        epsilon: 1.0,
        delta: 1e-6,
        clip_norm: 0.5,
        noise_multiplier: 2.0,
        local_dp: false,
        accountant: PrivacyAccountantType::RDP,
        target_epsilon: Some(3.0),
    }
}

/// Create a moderate privacy configuration.
///
/// Balanced privacy-utility trade-off suitable for most medical applications.
#[must_use]
pub fn moderate_privacy_config() -> PrivacyConfig {
    PrivacyConfig {
        epsilon: 8.0,
        delta: 1e-5,
        clip_norm: 1.0,
        noise_multiplier: 1.1,
        local_dp: false,
        accountant: PrivacyAccountantType::RDP,
        target_epsilon: Some(200.0), // Higher target for demo purposes
    }
}

/// Create a relaxed privacy configuration.
///
/// Higher epsilon (lower privacy) for less sensitive data or internal use.
#[must_use]
pub fn relaxed_privacy_config() -> PrivacyConfig {
    PrivacyConfig {
        epsilon: 20.0,
        delta: 1e-4,
        clip_norm: 2.0,
        noise_multiplier: 0.5,
        local_dp: false,
        accountant: PrivacyAccountantType::Moments,
        target_epsilon: Some(50.0),
    }
}

/// Create a local differential privacy configuration.
///
/// Privacy is applied at the client level before sharing.
#[must_use]
pub fn local_dp_config() -> PrivacyConfig {
    PrivacyConfig {
        epsilon: 4.0,
        delta: 1e-5,
        clip_norm: 1.0,
        noise_multiplier: 1.5,
        local_dp: true,
        accountant: PrivacyAccountantType::RDP,
        target_epsilon: Some(8.0),
    }
}

// ============================================================================
// Medical Dataset Samples
// ============================================================================

/// Create a ChestXray14 dataset configuration.
#[must_use]
pub fn chest_xray14_dataset() -> MedicalDataset {
    MedicalDataset {
        name: "ChestXray14".to_string(),
        task: MedicalTask::ChestXrayClassification,
        num_samples: 112120,
        num_classes: 14,
        class_names: vec![
            "Atelectasis".to_string(),
            "Cardiomegaly".to_string(),
            "Effusion".to_string(),
            "Infiltration".to_string(),
            "Mass".to_string(),
            "Nodule".to_string(),
            "Pneumonia".to_string(),
            "Pneumothorax".to_string(),
            "Consolidation".to_string(),
            "Edema".to_string(),
            "Emphysema".to_string(),
            "Fibrosis".to_string(),
            "Pleural Thickening".to_string(),
            "Hernia".to_string(),
        ],
        image_dims: (224, 224, 1),
        is_labeled: true,
    }
}

/// Create a ISIC skin lesion dataset configuration.
#[must_use]
pub fn isic_skin_lesion_dataset() -> MedicalDataset {
    MedicalDataset {
        name: "ISIC 2019".to_string(),
        task: MedicalTask::SkinLesionClassification,
        num_samples: 25331,
        num_classes: 8,
        class_names: vec![
            "Melanoma".to_string(),
            "Melanocytic Nevus".to_string(),
            "Basal Cell Carcinoma".to_string(),
            "Actinic Keratosis".to_string(),
            "Benign Keratosis".to_string(),
            "Dermatofibroma".to_string(),
            "Vascular Lesion".to_string(),
            "Squamous Cell Carcinoma".to_string(),
        ],
        image_dims: (224, 224, 3),
        is_labeled: true,
    }
}

/// Create an EyePACS diabetic retinopathy dataset configuration.
#[must_use]
pub fn eyepacs_dataset() -> MedicalDataset {
    MedicalDataset {
        name: "EyePACS".to_string(),
        task: MedicalTask::RetinalDisease,
        num_samples: 88702,
        num_classes: 5,
        class_names: vec![
            "No DR".to_string(),
            "Mild".to_string(),
            "Moderate".to_string(),
            "Severe".to_string(),
            "Proliferative DR".to_string(),
        ],
        image_dims: (512, 512, 3),
        is_labeled: true,
    }
}

/// Create a BraTS brain MRI dataset configuration.
#[must_use]
pub fn brats_dataset() -> MedicalDataset {
    MedicalDataset {
        name: "BraTS 2021".to_string(),
        task: MedicalTask::BrainMRISegmentation,
        num_samples: 2000,
        num_classes: 4, // Background, NCR/NET, ED, ET
        class_names: vec![
            "Background".to_string(),
            "Necrotic/Non-Enhancing Tumor".to_string(),
            "Peritumoral Edema".to_string(),
            "Enhancing Tumor".to_string(),
        ],
        image_dims: (240, 240, 155), // 3D volume
        is_labeled: true,
    }
}

/// Create a COVID-CT dataset configuration.
#[must_use]
pub fn covid_ct_dataset() -> MedicalDataset {
    MedicalDataset {
        name: "COVID-CT".to_string(),
        task: MedicalTask::CTScanAnalysis,
        num_samples: 746,
        num_classes: 2,
        class_names: vec!["COVID-19".to_string(), "Non-COVID".to_string()],
        image_dims: (224, 224, 1),
        is_labeled: true,
    }
}

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

    #[test]
    fn test_chest_xray_federation() {
        let clients = chest_xray_federation();
        assert_eq!(clients.len(), 5);
        assert!(clients.iter().all(|c| c.is_active));
        assert!(clients.iter().all(|c| c.data_size > 0));
    }

    #[test]
    fn test_skin_lesion_federation() {
        let clients = skin_lesion_federation();
        assert_eq!(clients.len(), 4);
        assert!(clients.iter().all(|c| c.is_active));
    }

    #[test]
    fn test_retinal_disease_federation() {
        let clients = retinal_disease_federation();
        assert_eq!(clients.len(), 3);
        assert!(clients.iter().all(|c| c.is_active));
    }

    #[test]
    fn test_hospital_network_config() {
        let config = hospital_network_config();
        assert_eq!(config.num_rounds, 10);
        assert!(config.privacy.is_some());
    }

    #[test]
    fn test_high_performance_config() {
        let config = high_performance_config();
        assert!(config.num_rounds > 100);
        assert!(config.local_epochs > 5);
    }

    #[test]
    fn test_resource_constrained_config() {
        let config = resource_constrained_config();
        assert!(config.num_rounds <= 50);
        assert!(!config.secure_aggregation);
    }

    #[test]
    fn test_fedprox_config() {
        let config = fedprox_heterogeneous_config();
        assert_eq!(config.aggregation_strategy, AggregationStrategy::FedProx);
    }

    #[test]
    fn test_scaffold_config() {
        let config = scaffold_config();
        assert_eq!(config.aggregation_strategy, AggregationStrategy::Scaffold);
    }

    #[test]
    fn test_byzantine_resilient_config() {
        let config = byzantine_resilient_config();
        assert_eq!(config.aggregation_strategy, AggregationStrategy::Krum);
        assert!(config.min_clients >= 5);
    }

    #[test]
    fn test_strict_privacy() {
        let config = strict_privacy_config();
        assert!(config.epsilon <= 2.0);
        assert!(config.noise_multiplier >= 1.5);
    }

    #[test]
    fn test_moderate_privacy() {
        let config = moderate_privacy_config();
        assert!(config.epsilon > 2.0);
        assert!(config.epsilon < 15.0);
    }

    #[test]
    fn test_relaxed_privacy() {
        let config = relaxed_privacy_config();
        assert!(config.epsilon >= 15.0);
    }

    #[test]
    fn test_local_dp_config() {
        let config = local_dp_config();
        assert!(config.local_dp);
    }

    #[test]
    fn test_chest_xray14_dataset() {
        let dataset = chest_xray14_dataset();
        assert_eq!(dataset.num_classes, 14);
        assert_eq!(dataset.class_names.len(), 14);
        assert_eq!(dataset.task, MedicalTask::ChestXrayClassification);
    }

    #[test]
    fn test_isic_dataset() {
        let dataset = isic_skin_lesion_dataset();
        assert_eq!(dataset.num_classes, 8);
        assert_eq!(dataset.task, MedicalTask::SkinLesionClassification);
    }

    #[test]
    fn test_eyepacs_dataset() {
        let dataset = eyepacs_dataset();
        assert_eq!(dataset.num_classes, 5);
        assert_eq!(dataset.task, MedicalTask::RetinalDisease);
    }

    #[test]
    fn test_brats_dataset() {
        let dataset = brats_dataset();
        assert_eq!(dataset.num_classes, 4);
        assert_eq!(dataset.task, MedicalTask::BrainMRISegmentation);
    }

    #[test]
    fn test_covid_ct_dataset() {
        let dataset = covid_ct_dataset();
        assert_eq!(dataset.num_classes, 2);
        assert_eq!(dataset.task, MedicalTask::CTScanAnalysis);
    }

    #[test]
    fn test_data_distribution_heterogeneity() {
        let clients = chest_xray_federation();

        // Research institute should have lower heterogeneity (more IID)
        let research = clients
            .iter()
            .find(|c| c.id == "research_institute")
            .unwrap();
        assert!(research.data_distribution.heterogeneity < 0.15);
        assert!(research.data_distribution.is_iid);

        // Clinics typically have higher heterogeneity
        let clinic = clients
            .iter()
            .find(|c| c.client_type == ClientType::Clinic)
            .unwrap();
        assert!(clinic.data_distribution.heterogeneity > 0.2);
    }
}