//! FederatedMed - Privacy-Preserving Medical AI.
//!
//! This demo implements federated learning for training on hospital data
//! without sharing sensitive patient information.

pub mod aggregator;
pub mod privacy;
pub mod sample_data;

use rand::Rng;
use thiserror::Error;

use fedmed_shared::{
    AggregationStrategy, ClientInfo, FederatedConfig, FederationResult, GlobalModel, LocalUpdate,
    PrivacyConfig, TestMetrics, TrainingPhase, TrainingProgress,
};

use aggregator::Aggregator;
use privacy::{DifferentialPrivacy, GradientClipper, PrivacyAccountant};

/// Errors that can occur in FederatedMed.
#[derive(Debug, Error)]
pub enum FederatedError {
    /// Invalid configuration.
    #[error("Invalid configuration: {0}")]
    InvalidConfig(String),

    /// Training failed.
    #[error("Training failed: {0}")]
    TrainingFailed(String),

    /// Not enough clients.
    #[error("Not enough clients: need {needed}, have {available}")]
    NotEnoughClients { needed: usize, available: usize },

    /// Privacy budget exhausted.
    #[error("Privacy budget exhausted: spent {spent}, limit {limit}")]
    PrivacyBudgetExhausted { spent: f64, limit: f64 },

    /// Aggregation failed.
    #[error("Aggregation failed: {0}")]
    AggregationFailed(String),

    /// Client error.
    #[error("Client error: {0}")]
    ClientError(String),
}

/// Main Federated Learning Coordinator.
#[derive(Debug)]
pub struct FederatedCoordinator {
    /// Configuration for federated learning.
    config: FederatedConfig,
    /// Global model.
    global_model: GlobalModel,
    /// Registered clients.
    clients: Vec<ClientInfo>,
    /// Aggregator for combining updates.
    aggregator: Aggregator,
    /// Privacy mechanism.
    privacy: Option<DifferentialPrivacy>,
    /// Privacy accountant for budget tracking.
    accountant: Option<PrivacyAccountant>,
    /// Gradient clipper.
    clipper: Option<GradientClipper>,
    /// Current round.
    current_round: u64,
    /// Training history (accuracy per round).
    accuracy_history: Vec<f64>,
    /// Training history (loss per round).
    loss_history: Vec<f64>,
    /// Total privacy budget spent.
    privacy_spent: f64,
}

impl Default for FederatedCoordinator {
    fn default() -> Self {
        Self::new(FederatedConfig::default())
    }
}

impl FederatedCoordinator {
    /// Create a new federated coordinator.
    #[must_use]
    pub fn new(config: FederatedConfig) -> Self {
        let aggregator = Aggregator::new(config.aggregation_strategy);

        let (privacy, accountant, clipper) = if let Some(ref privacy_config) = config.privacy {
            (
                Some(DifferentialPrivacy::new(privacy_config.clone())),
                Some(PrivacyAccountant::new(privacy_config.clone())),
                Some(GradientClipper::new(privacy_config.clip_norm)),
            )
        } else {
            (None, None, None)
        };

        Self {
            config,
            global_model: GlobalModel::default(),
            clients: Vec::new(),
            aggregator,
            privacy,
            accountant,
            clipper,
            current_round: 0,
            accuracy_history: Vec::new(),
            loss_history: Vec::new(),
            privacy_spent: 0.0,
        }
    }

    /// Register a client with the federation.
    pub fn register_client(&mut self, client: ClientInfo) {
        self.clients.push(client);
    }

    /// Register multiple clients.
    pub fn register_clients(&mut self, clients: Vec<ClientInfo>) {
        self.clients.extend(clients);
    }

    /// Get the number of registered clients.
    #[must_use]
    pub fn num_clients(&self) -> usize {
        self.clients.len()
    }

    /// Get active clients.
    #[must_use]
    pub fn active_clients(&self) -> Vec<&ClientInfo> {
        self.clients.iter().filter(|c| c.is_active).collect()
    }

    /// Select clients for a round of training.
    fn select_clients(&self) -> Vec<&ClientInfo> {
        let active: Vec<_> = self.active_clients();
        let num_to_select =
            ((active.len() as f32 * self.config.client_fraction).ceil() as usize).max(1);

        let mut rng = rand::thread_rng();
        let mut indices: Vec<usize> = (0..active.len()).collect();

        // Fisher-Yates shuffle for random selection
        for i in (1..indices.len()).rev() {
            let j = rng.gen_range(0..=i);
            indices.swap(i, j);
        }

        indices
            .into_iter()
            .take(num_to_select)
            .map(|i| active[i])
            .collect()
    }

    /// Simulate local training at a client.
    fn simulate_local_training(&self, client: &ClientInfo) -> LocalUpdate {
        let mut rng = rand::thread_rng();

        // Simulate weight updates
        let num_params = self.global_model.num_parameters;
        let weight_deltas: Vec<f32> = (0..num_params)
            .map(|_| rng.gen_range(-0.1_f32..0.1_f32))
            .collect();

        // Simulate local metrics
        let base_accuracy = 0.7 + (client.data_size as f64 / 10000.0).min(0.2);
        let noise: f64 = rng.gen_range(-0.05..0.05);

        LocalUpdate {
            client_id: client.id.clone(),
            round: self.current_round,
            weight_deltas,
            num_samples: client.data_size,
            local_loss: 0.5 - (self.current_round as f64 * 0.005),
            local_accuracy: base_accuracy + noise + (self.current_round as f64 * 0.002),
            training_time: (client.data_size as f64 / 100.0) + rng.gen_range(0.0..5.0),
            control_variates: None,
        }
    }

    /// Run one federated round.
    pub fn run_round(&mut self) -> Result<TrainingProgress, FederatedError> {
        // Check if we have enough clients
        let active_count = self.active_clients().len();
        if active_count < self.config.min_clients {
            return Err(FederatedError::NotEnoughClients {
                needed: self.config.min_clients,
                available: active_count,
            });
        }

        // Check privacy budget
        if let Some(ref privacy_config) = self.config.privacy
            && let Some(target) = privacy_config.target_epsilon
            && self.privacy_spent >= target
        {
            return Err(FederatedError::PrivacyBudgetExhausted {
                spent: self.privacy_spent,
                limit: target,
            });
        }

        self.current_round += 1;

        // Select clients
        let selected_clients = self.select_clients();
        let num_participating = selected_clients.len();

        // Collect local updates
        let mut updates: Vec<LocalUpdate> = selected_clients
            .iter()
            .map(|client| {
                let mut update = self.simulate_local_training(client);

                // Apply gradient clipping if enabled
                if let Some(ref clipper) = self.clipper {
                    update.weight_deltas = clipper.clip(&update.weight_deltas);
                }

                // Add noise if differential privacy is enabled
                if let Some(ref dp) = self.privacy {
                    update.weight_deltas = dp.add_noise(&update.weight_deltas);
                }

                update
            })
            .collect();

        // Track privacy budget
        if let Some(ref mut accountant) = self.accountant {
            let budget_spent = accountant.step(updates.len());
            self.privacy_spent += budget_spent;
        }

        // Aggregate updates
        let aggregated = self
            .aggregator
            .aggregate(&mut updates, &self.global_model)?;

        // Update global model
        self.update_global_model(&aggregated);

        // Evaluate
        let (accuracy, loss) = self.evaluate();
        self.accuracy_history.push(accuracy);
        self.loss_history.push(loss);

        // Update global model metrics
        self.global_model.validation_accuracy = Some(accuracy);
        self.global_model.validation_loss = Some(loss);

        Ok(TrainingProgress {
            round: self.current_round,
            total_rounds: self.config.num_rounds as u64,
            accuracy,
            loss,
            privacy_budget_spent: self.privacy_spent,
            participating_clients: num_participating,
            total_clients: self.clients.len(),
            elapsed_seconds: self.current_round as f64 * 60.0, // Simulated
            eta_seconds: Some((self.config.num_rounds as u64 - self.current_round) as f64 * 60.0),
            phase: TrainingPhase::Aggregation,
        })
    }

    /// Update global model with aggregated weights.
    fn update_global_model(&mut self, aggregated_deltas: &[f32]) {
        for (weight, delta) in self
            .global_model
            .weights
            .iter_mut()
            .zip(aggregated_deltas.iter())
        {
            *weight += delta * self.config.learning_rate as f32;
        }
        self.global_model.version = self.current_round;
    }

    /// Evaluate the global model.
    fn evaluate(&self) -> (f64, f64) {
        // Simulated evaluation
        let base_accuracy = 0.5;
        let improvement = (self.current_round as f64 * 0.008).min(0.45);
        let accuracy = base_accuracy + improvement;

        let base_loss = 2.0;
        let loss_reduction = (self.current_round as f64 * 0.03).min(1.5);
        let loss = base_loss - loss_reduction;

        (accuracy, loss)
    }

    /// Run complete federated training.
    pub fn train(
        &mut self,
        progress_callback: Option<Box<dyn Fn(TrainingProgress) + Send>>,
    ) -> Result<FederationResult, FederatedError> {
        // Validate configuration
        if self.clients.is_empty() {
            return Err(FederatedError::InvalidConfig(
                "No clients registered".to_string(),
            ));
        }

        for round in 0..self.config.num_rounds {
            let progress = self.run_round()?;

            if let Some(ref callback) = progress_callback {
                callback(progress);
            }

            // Early stopping check
            if round > 10 && self.accuracy_history.len() >= 5 {
                let recent: Vec<_> = self.accuracy_history.iter().rev().take(5).collect();
                let improvement = recent[0] - recent[4];
                if improvement < 0.001 {
                    break; // Convergence
                }
            }
        }

        // Calculate final test metrics
        let test_metrics = self.calculate_test_metrics();

        // Calculate client statistics
        let client_stats = self.calculate_client_stats();

        Ok(FederationResult {
            final_model: self.global_model.clone(),
            accuracy_history: self.accuracy_history.clone(),
            loss_history: self.loss_history.clone(),
            total_time: self.current_round as f64 * 60.0,
            rounds_completed: self.current_round,
            total_privacy_budget: self.privacy_spent,
            client_stats,
            test_metrics,
        })
    }

    /// Calculate test metrics.
    fn calculate_test_metrics(&self) -> TestMetrics {
        let accuracy = self.accuracy_history.last().copied().unwrap_or(0.0);

        TestMetrics {
            accuracy,
            precision: accuracy * 0.98,
            recall: accuracy * 0.96,
            f1_score: accuracy * 0.97,
            auc_roc: accuracy * 1.02, // AUC typically slightly higher
            confusion_matrix: vec![85, 5, 3, 7], // Example 2-class confusion matrix
            num_classes: 2,
        }
    }

    /// Calculate client statistics.
    fn calculate_client_stats(&self) -> Vec<fedmed_shared::ClientStats> {
        self.clients
            .iter()
            .map(|client| {
                let rounds_participated = (self.current_round as f64
                    * self.config.client_fraction as f64)
                    .ceil() as usize;

                fedmed_shared::ClientStats {
                    client_id: client.id.clone(),
                    rounds_participated,
                    avg_local_accuracy: 0.85,
                    total_samples: client.data_size * rounds_participated,
                    avg_training_time: client.data_size as f64 / 100.0,
                }
            })
            .collect()
    }

    /// Get current global model.
    #[must_use]
    pub fn global_model(&self) -> &GlobalModel {
        &self.global_model
    }

    /// Get current round.
    #[must_use]
    pub fn current_round(&self) -> u64 {
        self.current_round
    }

    /// Get privacy budget spent.
    #[must_use]
    pub fn privacy_spent(&self) -> f64 {
        self.privacy_spent
    }

    /// Get accuracy history.
    #[must_use]
    pub fn accuracy_history(&self) -> &[f64] {
        &self.accuracy_history
    }

    /// Get loss history.
    #[must_use]
    pub fn loss_history(&self) -> &[f64] {
        &self.loss_history
    }

    /// Get configuration.
    #[must_use]
    pub fn config(&self) -> &FederatedConfig {
        &self.config
    }

    /// Update aggregation strategy.
    pub fn set_aggregation_strategy(&mut self, strategy: AggregationStrategy) {
        self.aggregator = Aggregator::new(strategy);
    }

    /// Update privacy configuration.
    pub fn set_privacy_config(&mut self, config: PrivacyConfig) {
        self.privacy = Some(DifferentialPrivacy::new(config.clone()));
        self.accountant = Some(PrivacyAccountant::new(config.clone()));
        self.clipper = Some(GradientClipper::new(config.clip_norm));
    }
}

/// Run the demo.
pub fn run_demo() -> Result<FederationResult, FederatedError> {
    // Create coordinator with sample config
    let config = sample_data::hospital_network_config();
    let mut coordinator = FederatedCoordinator::new(config);

    // Register sample clients
    let clients = sample_data::chest_xray_federation();
    coordinator.register_clients(clients);

    // Run training
    coordinator.train(None)
}

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

    #[test]
    fn test_coordinator_creation() {
        let coordinator = FederatedCoordinator::default();
        assert_eq!(coordinator.current_round(), 0);
        assert_eq!(coordinator.num_clients(), 0);
    }

    #[test]
    fn test_register_clients() {
        let mut coordinator = FederatedCoordinator::default();
        let clients = fedmed_shared::sample_clients();
        coordinator.register_clients(clients);
        assert_eq!(coordinator.num_clients(), 5);
    }

    #[test]
    fn test_run_round() {
        let mut coordinator = FederatedCoordinator::default();
        coordinator.register_clients(fedmed_shared::sample_clients());

        let result = coordinator.run_round();
        assert!(result.is_ok());

        let progress = result.unwrap();
        assert_eq!(progress.round, 1);
        assert!(progress.participating_clients > 0);
    }

    #[test]
    fn test_not_enough_clients() {
        let config = FederatedConfig {
            min_clients: 10,
            ..Default::default()
        };
        let mut coordinator = FederatedCoordinator::new(config);
        coordinator.register_clients(vec![ClientInfo::default()]);

        let result = coordinator.run_round();
        assert!(matches!(
            result,
            Err(FederatedError::NotEnoughClients { .. })
        ));
    }

    #[test]
    fn test_training() {
        let config = FederatedConfig {
            num_rounds: 5,
            min_clients: 2,
            privacy: None, // Disable privacy for this test to avoid budget exhaustion
            ..Default::default()
        };
        let mut coordinator = FederatedCoordinator::new(config);
        coordinator.register_clients(fedmed_shared::sample_clients());

        let result = coordinator.train(None);
        assert!(result.is_ok());

        let federation_result = result.unwrap();
        assert!(federation_result.rounds_completed > 0);
        assert!(!federation_result.accuracy_history.is_empty());
    }

    #[test]
    fn test_privacy_tracking() {
        let config = FederatedConfig {
            num_rounds: 3,
            privacy: Some(fedmed_shared::sample_privacy_config()),
            ..Default::default()
        };
        let mut coordinator = FederatedCoordinator::new(config);
        coordinator.register_clients(fedmed_shared::sample_clients());

        let _ = coordinator.run_round();
        assert!(coordinator.privacy_spent() > 0.0);
    }

    #[test]
    fn test_set_aggregation_strategy() {
        let mut coordinator = FederatedCoordinator::default();
        coordinator.set_aggregation_strategy(AggregationStrategy::FedProx);
        // Strategy is set internally
    }

    #[test]
    fn test_run_demo() {
        let result = run_demo();
        if let Err(ref e) = result {
            eprintln!("run_demo error: {:?}", e);
        }
        assert!(result.is_ok());
    }
}