//! Shared types for FederatedMed - Privacy-Preserving Medical AI.
//!
//! This crate defines the IPC types for federated learning on hospital data
//! without sharing sensitive patient information.

use serde::{Deserialize, Serialize};

// ============================================================================
// Federated Learning Configuration
// ============================================================================

/// Configuration for federated learning training.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FederatedConfig {
    /// Number of federated rounds.
    pub num_rounds: usize,
    /// Number of local epochs per round.
    pub local_epochs: usize,
    /// Batch size for local training.
    pub batch_size: usize,
    /// Aggregation strategy.
    pub aggregation_strategy: AggregationStrategy,
    /// Learning rate.
    pub learning_rate: f64,
    /// Minimum number of clients per round.
    pub min_clients: usize,
    /// Client selection fraction (0.0 to 1.0).
    pub client_fraction: f32,
    /// Whether to use secure aggregation.
    pub secure_aggregation: bool,
    /// Privacy configuration.
    pub privacy: Option<PrivacyConfig>,
}

impl Default for FederatedConfig {
    fn default() -> Self {
        Self {
            num_rounds: 100,
            local_epochs: 5,
            batch_size: 32,
            aggregation_strategy: AggregationStrategy::FedAvg,
            learning_rate: 0.01,
            min_clients: 3,
            client_fraction: 0.3,
            secure_aggregation: true,
            privacy: Some(PrivacyConfig::default()),
        }
    }
}

// ============================================================================
// Aggregation Strategy
// ============================================================================

/// Aggregation strategy for combining client updates.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum AggregationStrategy {
    /// Federated Averaging - weighted average of model updates.
    FedAvg,
    /// FedProx - adds proximal term for heterogeneous data.
    FedProx,
    /// Scaffold - uses control variates to correct client drift.
    Scaffold,
    /// FedAdam - adaptive learning rate aggregation.
    FedAdam,
    /// FedYogi - adaptive aggregation with momentum.
    FedYogi,
    /// TrimmedMean - robust aggregation against Byzantine clients.
    TrimmedMean,
    /// Median - median-based robust aggregation.
    Median,
    /// Krum - Byzantine-resilient aggregation.
    Krum,
}

/// Parameters specific to FedProx.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FedProxParams {
    /// Proximal term coefficient (mu).
    pub mu: f64,
}

impl Default for FedProxParams {
    fn default() -> Self {
        Self { mu: 0.01 }
    }
}

/// Parameters specific to Scaffold.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScaffoldParams {
    /// Server learning rate.
    pub server_lr: f64,
    /// Client learning rate.
    pub client_lr: f64,
}

impl Default for ScaffoldParams {
    fn default() -> Self {
        Self {
            server_lr: 1.0,
            client_lr: 0.1,
        }
    }
}

// ============================================================================
// Client Information
// ============================================================================

/// Information about a federated learning client (e.g., a hospital).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientInfo {
    /// Unique client identifier.
    pub id: String,
    /// Human-readable name (e.g., "General Hospital").
    pub name: String,
    /// Number of data samples at this client.
    pub data_size: usize,
    /// Whether the client is currently active.
    pub is_active: bool,
    /// Client type.
    pub client_type: ClientType,
    /// Data distribution characteristics.
    pub data_distribution: DataDistribution,
    /// Last communication timestamp (Unix epoch).
    pub last_seen: u64,
    /// Network latency in milliseconds.
    pub latency_ms: u32,
}

impl Default for ClientInfo {
    fn default() -> Self {
        Self {
            id: "client_0".to_string(),
            name: "Default Hospital".to_string(),
            data_size: 1000,
            is_active: true,
            client_type: ClientType::Hospital,
            data_distribution: DataDistribution::default(),
            last_seen: 0,
            latency_ms: 50,
        }
    }
}

/// Type of federated learning client.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum ClientType {
    /// Hospital.
    Hospital,
    /// Clinic.
    Clinic,
    /// Research institution.
    Research,
    /// Radiology center.
    Radiology,
    /// Laboratory.
    Laboratory,
}

/// Data distribution characteristics at a client.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataDistribution {
    /// Class distribution (class_id -> count).
    pub class_counts: Vec<usize>,
    /// Whether data is IID (identically and independently distributed).
    pub is_iid: bool,
    /// Data heterogeneity measure (0.0 = homogeneous, 1.0 = highly heterogeneous).
    pub heterogeneity: f32,
}

impl Default for DataDistribution {
    fn default() -> Self {
        Self {
            class_counts: vec![500, 300, 200],
            is_iid: false,
            heterogeneity: 0.3,
        }
    }
}

// ============================================================================
// Privacy Configuration
// ============================================================================

/// Privacy configuration for differential privacy.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PrivacyConfig {
    /// Epsilon - privacy budget (lower = more private).
    pub epsilon: f64,
    /// Delta - probability of privacy breach.
    pub delta: f64,
    /// Gradient clipping norm for sensitivity bounding.
    pub clip_norm: f64,
    /// Noise multiplier for Gaussian mechanism.
    pub noise_multiplier: f64,
    /// Whether to use local differential privacy.
    pub local_dp: bool,
    /// Accountant type for privacy budget tracking.
    pub accountant: PrivacyAccountantType,
    /// Target epsilon for automatic noise calibration.
    pub target_epsilon: Option<f64>,
}

impl Default for PrivacyConfig {
    fn default() -> Self {
        Self {
            epsilon: 8.0,
            delta: 1e-5,
            clip_norm: 1.0,
            noise_multiplier: 1.1,
            local_dp: false,
            accountant: PrivacyAccountantType::RDP,
            target_epsilon: Some(10.0),
        }
    }
}

/// Type of privacy accountant.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum PrivacyAccountantType {
    /// Moments Accountant.
    Moments,
    /// Renyi Differential Privacy.
    RDP,
    /// Gaussian Differential Privacy.
    GDP,
    /// Privacy Loss Distribution.
    PLD,
}

// ============================================================================
// Model Types
// ============================================================================

/// Global model maintained by the server.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GlobalModel {
    /// Model version/round number.
    pub version: u64,
    /// Model weights (flattened).
    pub weights: Vec<f32>,
    /// Weight shapes for reconstruction.
    pub shapes: Vec<Vec<usize>>,
    /// Layer names.
    pub layer_names: Vec<String>,
    /// Model architecture type.
    pub architecture: ModelArchitecture,
    /// Number of parameters.
    pub num_parameters: usize,
    /// Accuracy on validation set.
    pub validation_accuracy: Option<f64>,
    /// Loss on validation set.
    pub validation_loss: Option<f64>,
}

impl Default for GlobalModel {
    fn default() -> Self {
        // Create a simple model structure
        let shapes = vec![
            vec![784, 256], // Input to hidden1
            vec![256],      // Hidden1 bias
            vec![256, 128], // Hidden1 to hidden2
            vec![128],      // Hidden2 bias
            vec![128, 10],  // Hidden2 to output
            vec![10],       // Output bias
        ];
        let num_params: usize = shapes.iter().map(|s| s.iter().product::<usize>()).sum();

        Self {
            version: 0,
            weights: vec![0.0; num_params],
            shapes,
            layer_names: vec![
                "fc1.weight".to_string(),
                "fc1.bias".to_string(),
                "fc2.weight".to_string(),
                "fc2.bias".to_string(),
                "fc3.weight".to_string(),
                "fc3.bias".to_string(),
            ],
            architecture: ModelArchitecture::CNN,
            num_parameters: num_params,
            validation_accuracy: None,
            validation_loss: None,
        }
    }
}

/// Model architecture type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum ModelArchitecture {
    /// Convolutional neural network.
    CNN,
    /// ResNet variant.
    ResNet,
    /// DenseNet variant.
    DenseNet,
    /// Vision Transformer.
    ViT,
    /// U-Net for segmentation.
    UNet,
    /// Multi-layer perceptron.
    MLP,
    /// Transformer.
    Transformer,
}

/// Local update from a client.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LocalUpdate {
    /// Client ID.
    pub client_id: String,
    /// Round number.
    pub round: u64,
    /// Weight updates (delta from global model).
    pub weight_deltas: Vec<f32>,
    /// Number of local samples used.
    pub num_samples: usize,
    /// Local training loss.
    pub local_loss: f64,
    /// Local training accuracy.
    pub local_accuracy: f64,
    /// Training time in seconds.
    pub training_time: f64,
    /// Control variates for Scaffold (optional).
    pub control_variates: Option<Vec<f32>>,
}

impl Default for LocalUpdate {
    fn default() -> Self {
        Self {
            client_id: "client_0".to_string(),
            round: 0,
            weight_deltas: vec![],
            num_samples: 0,
            local_loss: 0.0,
            local_accuracy: 0.0,
            training_time: 0.0,
            control_variates: None,
        }
    }
}

// ============================================================================
// Training Progress
// ============================================================================

/// Training progress for federated learning.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingProgress {
    /// Current round.
    pub round: u64,
    /// Total rounds.
    pub total_rounds: u64,
    /// Global model accuracy.
    pub accuracy: f64,
    /// Global model loss.
    pub loss: f64,
    /// Privacy budget spent so far.
    pub privacy_budget_spent: f64,
    /// Number of participating clients this round.
    pub participating_clients: usize,
    /// Total number of clients.
    pub total_clients: usize,
    /// Elapsed time in seconds.
    pub elapsed_seconds: f64,
    /// Estimated time remaining in seconds.
    pub eta_seconds: Option<f64>,
    /// Current phase.
    pub phase: TrainingPhase,
}

impl Default for TrainingProgress {
    fn default() -> Self {
        Self {
            round: 0,
            total_rounds: 100,
            accuracy: 0.0,
            loss: f64::INFINITY,
            privacy_budget_spent: 0.0,
            participating_clients: 0,
            total_clients: 0,
            elapsed_seconds: 0.0,
            eta_seconds: None,
            phase: TrainingPhase::Initialization,
        }
    }
}

/// Training phase.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum TrainingPhase {
    /// Initialization phase.
    Initialization,
    /// Client selection phase.
    ClientSelection,
    /// Local training phase.
    LocalTraining,
    /// Aggregation phase.
    Aggregation,
    /// Evaluation phase.
    Evaluation,
    /// Completed.
    Completed,
}

// ============================================================================
// Medical Domain Types
// ============================================================================

/// Medical imaging task type.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum MedicalTask {
    /// Chest X-ray classification (e.g., COVID, pneumonia).
    ChestXrayClassification,
    /// Skin lesion classification (melanoma detection).
    SkinLesionClassification,
    /// Retinal disease detection.
    RetinalDisease,
    /// Brain MRI segmentation.
    BrainMRISegmentation,
    /// CT scan analysis.
    CTScanAnalysis,
    /// Mammography screening.
    MammographyScreening,
    /// Pathology slide classification.
    PathologySlide,
    /// ECG classification.
    ECGClassification,
}

/// Medical dataset information.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MedicalDataset {
    /// Dataset name.
    pub name: String,
    /// Task type.
    pub task: MedicalTask,
    /// Number of samples.
    pub num_samples: usize,
    /// Number of classes.
    pub num_classes: usize,
    /// Class names.
    pub class_names: Vec<String>,
    /// Image dimensions (height, width, channels).
    pub image_dims: (usize, usize, usize),
    /// Whether dataset is labeled.
    pub is_labeled: bool,
}

impl Default for MedicalDataset {
    fn default() -> Self {
        Self {
            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,
        }
    }
}

// ============================================================================
// Federation Results
// ============================================================================

/// Final result of federated training.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FederationResult {
    /// Final global model.
    pub final_model: GlobalModel,
    /// Training history (accuracy per round).
    pub accuracy_history: Vec<f64>,
    /// Training history (loss per round).
    pub loss_history: Vec<f64>,
    /// Total training time.
    pub total_time: f64,
    /// Number of rounds completed.
    pub rounds_completed: u64,
    /// Total privacy budget spent.
    pub total_privacy_budget: f64,
    /// Client participation statistics.
    pub client_stats: Vec<ClientStats>,
    /// Final test metrics.
    pub test_metrics: TestMetrics,
}

/// Statistics for a single client.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientStats {
    /// Client ID.
    pub client_id: String,
    /// Number of rounds participated.
    pub rounds_participated: usize,
    /// Average local accuracy.
    pub avg_local_accuracy: f64,
    /// Total samples contributed.
    pub total_samples: usize,
    /// Average training time per round.
    pub avg_training_time: f64,
}

/// Test metrics for evaluation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestMetrics {
    /// Accuracy.
    pub accuracy: f64,
    /// Precision (macro-averaged).
    pub precision: f64,
    /// Recall (macro-averaged).
    pub recall: f64,
    /// F1 score (macro-averaged).
    pub f1_score: f64,
    /// AUC-ROC (macro-averaged).
    pub auc_roc: f64,
    /// Confusion matrix (flattened).
    pub confusion_matrix: Vec<usize>,
    /// Number of classes.
    pub num_classes: usize,
}

impl Default for TestMetrics {
    fn default() -> Self {
        Self {
            accuracy: 0.0,
            precision: 0.0,
            recall: 0.0,
            f1_score: 0.0,
            auc_roc: 0.0,
            confusion_matrix: vec![],
            num_classes: 0,
        }
    }
}

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

/// Create a sample federated config.
#[must_use]
pub fn sample_federated_config() -> FederatedConfig {
    FederatedConfig {
        num_rounds: 50,
        local_epochs: 3,
        batch_size: 32,
        aggregation_strategy: AggregationStrategy::FedAvg,
        learning_rate: 0.01,
        min_clients: 3,
        client_fraction: 0.5,
        secure_aggregation: true,
        privacy: Some(PrivacyConfig {
            epsilon: 8.0,
            delta: 1e-5,
            clip_norm: 1.0,
            noise_multiplier: 1.1,
            local_dp: false,
            accountant: PrivacyAccountantType::RDP,
            target_epsilon: Some(10.0),
        }),
    }
}

/// Create sample client infos.
#[must_use]
pub fn sample_clients() -> Vec<ClientInfo> {
    vec![
        ClientInfo {
            id: "hospital_a".to_string(),
            name: "City General Hospital".to_string(),
            data_size: 5000,
            is_active: true,
            client_type: ClientType::Hospital,
            data_distribution: DataDistribution {
                class_counts: vec![1500, 1200, 800, 700, 500, 300],
                is_iid: false,
                heterogeneity: 0.25,
            },
            last_seen: 1704067200,
            latency_ms: 45,
        },
        ClientInfo {
            id: "hospital_b".to_string(),
            name: "University Medical Center".to_string(),
            data_size: 8000,
            is_active: true,
            client_type: ClientType::Hospital,
            data_distribution: DataDistribution {
                class_counts: vec![2000, 1800, 1500, 1200, 900, 600],
                is_iid: false,
                heterogeneity: 0.2,
            },
            last_seen: 1704067200,
            latency_ms: 30,
        },
        ClientInfo {
            id: "clinic_c".to_string(),
            name: "Downtown Health Clinic".to_string(),
            data_size: 2000,
            is_active: true,
            client_type: ClientType::Clinic,
            data_distribution: DataDistribution {
                class_counts: vec![600, 500, 400, 300, 150, 50],
                is_iid: false,
                heterogeneity: 0.4,
            },
            last_seen: 1704067200,
            latency_ms: 60,
        },
        ClientInfo {
            id: "research_d".to_string(),
            name: "National Research Institute".to_string(),
            data_size: 3500,
            is_active: true,
            client_type: ClientType::Research,
            data_distribution: DataDistribution {
                class_counts: vec![700, 700, 700, 700, 350, 350],
                is_iid: true,
                heterogeneity: 0.1,
            },
            last_seen: 1704067200,
            latency_ms: 25,
        },
        ClientInfo {
            id: "radiology_e".to_string(),
            name: "Regional Radiology Center".to_string(),
            data_size: 4500,
            is_active: true,
            client_type: ClientType::Radiology,
            data_distribution: DataDistribution {
                class_counts: vec![1200, 1000, 900, 800, 400, 200],
                is_iid: false,
                heterogeneity: 0.3,
            },
            last_seen: 1704067200,
            latency_ms: 40,
        },
    ]
}

/// Create a sample global model.
#[must_use]
pub fn sample_global_model() -> GlobalModel {
    GlobalModel::default()
}

/// Create sample privacy config.
#[must_use]
pub fn sample_privacy_config() -> PrivacyConfig {
    PrivacyConfig::default()
}

/// Create a sample training progress.
#[must_use]
pub fn sample_training_progress() -> TrainingProgress {
    TrainingProgress {
        round: 25,
        total_rounds: 100,
        accuracy: 0.85,
        loss: 0.45,
        privacy_budget_spent: 4.2,
        participating_clients: 4,
        total_clients: 5,
        elapsed_seconds: 3600.0,
        eta_seconds: Some(3600.0),
        phase: TrainingPhase::Aggregation,
    }
}

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

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

    #[test]
    fn test_federated_config() {
        let config = sample_federated_config();
        assert_eq!(config.num_rounds, 50);
        assert_eq!(config.aggregation_strategy, AggregationStrategy::FedAvg);
        assert!(config.privacy.is_some());
    }

    #[test]
    fn test_aggregation_strategy() {
        let strategies = vec![
            AggregationStrategy::FedAvg,
            AggregationStrategy::FedProx,
            AggregationStrategy::Scaffold,
        ];
        for strategy in strategies {
            let json = serde_json::to_string(&strategy).unwrap();
            let parsed: AggregationStrategy = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, strategy);
        }
    }

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

    #[test]
    fn test_privacy_config() {
        let config = sample_privacy_config();
        assert!(config.epsilon > 0.0);
        assert!(config.delta > 0.0);
        assert!(config.clip_norm > 0.0);
    }

    #[test]
    fn test_global_model() {
        let model = sample_global_model();
        assert_eq!(model.version, 0);
        assert!(!model.weights.is_empty());
        assert!(!model.shapes.is_empty());
    }

    #[test]
    fn test_local_update() {
        let update = LocalUpdate {
            client_id: "hospital_a".to_string(),
            round: 10,
            weight_deltas: vec![0.1, 0.2, 0.3],
            num_samples: 1000,
            local_loss: 0.5,
            local_accuracy: 0.85,
            training_time: 120.0,
            control_variates: None,
        };
        assert_eq!(update.round, 10);
        assert_eq!(update.num_samples, 1000);
    }

    #[test]
    fn test_training_progress() {
        let progress = sample_training_progress();
        assert_eq!(progress.round, 25);
        assert!(progress.accuracy > 0.0);
        assert_eq!(progress.phase, TrainingPhase::Aggregation);
    }

    #[test]
    fn test_medical_task() {
        let tasks = vec![
            MedicalTask::ChestXrayClassification,
            MedicalTask::SkinLesionClassification,
            MedicalTask::RetinalDisease,
        ];
        for task in tasks {
            let json = serde_json::to_string(&task).unwrap();
            let parsed: MedicalTask = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, task);
        }
    }

    #[test]
    fn test_medical_dataset() {
        let dataset = MedicalDataset::default();
        assert_eq!(dataset.num_classes, 14);
        assert_eq!(dataset.class_names.len(), 14);
    }

    #[test]
    fn test_test_metrics() {
        let metrics = TestMetrics {
            accuracy: 0.92,
            precision: 0.90,
            recall: 0.88,
            f1_score: 0.89,
            auc_roc: 0.95,
            confusion_matrix: vec![100, 5, 3, 92],
            num_classes: 2,
        };
        assert!(metrics.accuracy > 0.9);
        assert!(metrics.auc_roc > 0.9);
    }

    #[test]
    fn test_serialization() {
        let config = sample_federated_config();
        let json = serde_json::to_string(&config).unwrap();
        let parsed: FederatedConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.num_rounds, config.num_rounds);
    }

    #[test]
    fn test_federation_result() {
        let result = FederationResult {
            final_model: sample_global_model(),
            accuracy_history: vec![0.5, 0.6, 0.7, 0.8, 0.85],
            loss_history: vec![2.0, 1.5, 1.0, 0.7, 0.5],
            total_time: 7200.0,
            rounds_completed: 50,
            total_privacy_budget: 8.5,
            client_stats: vec![],
            test_metrics: TestMetrics::default(),
        };
        assert_eq!(result.accuracy_history.len(), 5);
        assert_eq!(result.rounds_completed, 50);
    }
}