//! Ensemble Training Coordinator
//!
//! Coordinates training of multiple ML models (DQN, PPO, MAMBA-2, TFT) as an ensemble
//! with weight optimization, checkpoint synchronization, and performance-based adaptation.
//!
//! Key Features:
//! - Multi-model training coordination (sequential or parallel)
//! - Dynamic weight optimization based on validation performance
//! - Synchronized checkpoint management for all models
//! - Failure recovery and retry mechanisms
//! - Integration with existing ML Training Service infrastructure

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use anyhow::{anyhow, Result};
use chrono::{DateTime, Utc};
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
use uuid::Uuid;

use ml::training_pipeline::ProductionTrainingConfig;

/// Status of individual model training within ensemble
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelTrainingStatus {
    /// Not yet started
    Pending,
    /// Currently training
    Training,
    /// Successfully completed
    Completed,
    /// Training failed
    Failed,
    /// Paused by user
    Paused,
}

/// Configuration for ensemble training
#[derive(Debug, Clone)]
pub struct EnsembleTrainingConfig {
    /// Unique job identifier
    pub job_id: Uuid,

    /// Training configuration for each model
    pub model_configs: HashMap<String, ProductionTrainingConfig>,

    /// Initial weights for each model (must sum to 1.0)
    pub model_weights: HashMap<String, f64>,

    /// Enable dynamic weight optimization based on performance
    pub enable_weight_optimization: bool,

    /// Optimize weights every N epochs
    pub weight_optimization_interval_epochs: u32,

    /// Save checkpoints every N epochs
    pub checkpoint_interval_epochs: u32,

    /// Maximum number of epochs for training
    pub max_epochs: u32,

    /// Train models in parallel (true) or sequentially (false)
    pub parallel_training: bool,

    /// Configuration created timestamp
    pub created_at: DateTime<Utc>,
}

impl EnsembleTrainingConfig {
    /// Validate ensemble configuration
    pub fn validate(&self) -> Result<()> {
        // Check all 4 required models present
        let required_models = ["DQN", "PPO", "MAMBA2", "TFT"];
        for model in &required_models {
            if !self.model_configs.contains_key(*model) {
                return Err(anyhow!("Missing configuration for model: {}", model));
            }
            if !self.model_weights.contains_key(*model) {
                return Err(anyhow!("Missing weight for model: {}", model));
            }
        }

        // Check weights sum to 1.0
        let weight_sum: f64 = self.model_weights.values().sum();
        if (weight_sum - 1.0).abs() > 1e-6 {
            return Err(anyhow!("Model weights must sum to 1.0, got {}", weight_sum));
        }

        // Validate each model has valid config
        for (model_name, config) in &self.model_configs {
            if config.model_config.input_dim == 0 {
                return Err(anyhow!("{}: Invalid input dimension (0)", model_name));
            }
            if config.model_config.hidden_dims.is_empty() {
                return Err(anyhow!("{}: No hidden layers specified", model_name));
            }
            if config.training_params.max_epochs == 0 {
                return Err(anyhow!("{}: Invalid max_epochs (0)", model_name));
            }
        }

        Ok(())
    }

    /// Check if configuration is valid (boolean version)
    pub fn is_valid(&self) -> bool {
        self.validate().is_ok()
    }

    /// Get total number of models
    pub fn model_count(&self) -> usize {
        self.model_configs.len()
    }

    /// Check if specific model exists
    pub fn has_model(&self, name: &str) -> bool {
        self.model_configs.contains_key(name)
    }

    /// Get total weight
    pub fn total_weight(&self) -> f64 {
        self.model_weights.values().sum()
    }

    /// Check if model has configuration
    pub fn has_model_config(&self, name: &str) -> bool {
        self.model_configs.contains_key(name)
    }

    /// Check if model has weight
    pub fn has_model_weight(&self, name: &str) -> bool {
        self.model_weights.contains_key(name)
    }
}

/// Performance metrics for a single model
#[derive(Debug, Clone)]
pub struct ModelPerformance {
    pub accuracy: f64,
    pub loss: f64,
    pub sharpe_ratio: f64,
    pub validation_loss: f64,
    pub epoch: u32,
    pub updated_at: DateTime<Utc>,
}

/// State of a single model within the ensemble
#[derive(Debug, Clone)]
struct ModelState {
    status: ModelTrainingStatus,
    performance: Option<ModelPerformance>,
    latest_checkpoint: Option<String>,
    current_epoch: u32,
    error_message: Option<String>,
}

impl Default for ModelState {
    fn default() -> Self {
        Self {
            status: ModelTrainingStatus::Pending,
            performance: None,
            latest_checkpoint: None,
            current_epoch: 0,
            error_message: None,
        }
    }
}

/// Ensemble Training Coordinator
///
/// Manages the lifecycle of training multiple models as an ensemble,
/// including weight optimization and checkpoint synchronization.
pub struct EnsembleTrainingCoordinator {
    /// Ensemble configuration
    config: EnsembleTrainingConfig,

    /// Current state of each model
    model_states: Arc<RwLock<HashMap<String, ModelState>>>,

    /// Current (possibly optimized) weights
    current_weights: Arc<RwLock<HashMap<String, f64>>>,

    /// Training start time
    started_at: Option<Instant>,
}

impl EnsembleTrainingCoordinator {
    /// Create new ensemble training coordinator
    pub async fn new(config: EnsembleTrainingConfig) -> Result<Self> {
        // Validate configuration
        config.validate()?;

        // Initialize model states
        let mut model_states = HashMap::new();
        for model_name in config.model_configs.keys() {
            model_states.insert(model_name.clone(), ModelState::default());
        }

        // Initialize weights from config
        let current_weights = config.model_weights.clone();

        info!(
            "Created ensemble training coordinator for job {} with {} models",
            config.job_id,
            config.model_count()
        );

        Ok(Self {
            config,
            model_states: Arc::new(RwLock::new(model_states)),
            current_weights: Arc::new(RwLock::new(current_weights)),
            started_at: None,
        })
    }

    /// Start ensemble training
    pub async fn start_ensemble_training(&mut self) -> Result<Uuid> {
        info!(
            "Starting ensemble training for job {} with {} models",
            self.config.job_id,
            self.config.model_count()
        );

        self.started_at = Some(Instant::now());

        // Set all models to Training status
        {
            let mut states = self.model_states.write().await;
            for state in states.values_mut() {
                state.status = ModelTrainingStatus::Training;
            }
        }

        Ok(self.config.job_id)
    }

    /// Get status of specific model
    pub async fn get_model_status(&self, model_name: &str) -> Result<ModelTrainingStatus> {
        let states = self.model_states.read().await;
        states
            .get(model_name)
            .map(|s| s.status.clone())
            .ok_or_else(|| anyhow!("Model {} not found", model_name))
    }

    /// Get current ensemble weights
    pub async fn get_current_weights(&self) -> Result<HashMap<String, f64>> {
        let weights = self.current_weights.read().await;
        Ok(weights.clone())
    }

    /// Simulate training epochs (for testing)
    pub async fn simulate_training_epochs(&self, epochs: u32) -> Result<()> {
        debug!("Simulating {} training epochs", epochs);

        for epoch in 1..=epochs {
            // Update model states
            let mut states = self.model_states.write().await;
            for (model_name, state) in states.iter_mut() {
                state.current_epoch = epoch;

                // Simulate checkpoint creation
                let checkpoint_path = format!(
                    "models/{}/checkpoints/{}_epoch_{}.safetensors",
                    self.config.job_id, model_name, epoch
                );
                state.latest_checkpoint = Some(checkpoint_path);

                // Simulate performance metrics
                if state.performance.is_none() {
                    state.performance = Some(ModelPerformance {
                        accuracy: 0.5 + (epoch as f64 * 0.01),
                        loss: 1.0 - (epoch as f64 * 0.05),
                        sharpe_ratio: 0.5 + (epoch as f64 * 0.05),
                        validation_loss: 1.0 - (epoch as f64 * 0.04),
                        epoch,
                        updated_at: Utc::now(),
                    });
                }
            }
            drop(states);

            // Trigger weight optimization if needed
            if self.config.enable_weight_optimization
                && epoch % self.config.weight_optimization_interval_epochs == 0
            {
                self.optimize_weights().await?;
            }
        }

        Ok(())
    }

    /// Get latest checkpoint for a model
    pub async fn get_latest_checkpoint(&self, model_name: &str) -> Result<Option<String>> {
        let states = self.model_states.read().await;
        states
            .get(model_name)
            .map(|s| s.latest_checkpoint.clone())
            .ok_or_else(|| anyhow!("Model {} not found", model_name))
    }

    /// Get all model checkpoints
    pub async fn get_all_checkpoints(&self) -> Result<Vec<(String, String)>> {
        let states = self.model_states.read().await;
        let mut checkpoints = Vec::new();

        for (model_name, state) in states.iter() {
            if let Some(checkpoint) = &state.latest_checkpoint {
                checkpoints.push((model_name.clone(), checkpoint.clone()));
            }
        }

        Ok(checkpoints)
    }

    /// Load synchronized ensemble from specific epoch
    pub async fn load_synchronized_ensemble(&self, epoch: u32) -> Result<()> {
        info!(
            "Loading synchronized ensemble from epoch {} for job {}",
            epoch, self.config.job_id
        );

        // Verify all models have checkpoints from this epoch
        let states = self.model_states.read().await;
        for (model_name, state) in states.iter() {
            if let Some(checkpoint) = &state.latest_checkpoint {
                if !checkpoint.contains(&format!("epoch_{}", epoch)) {
                    return Err(anyhow!(
                        "Model {} checkpoint not from epoch {}",
                        model_name,
                        epoch
                    ));
                }
            } else {
                return Err(anyhow!(
                    "Model {} has no checkpoint for epoch {}",
                    model_name,
                    epoch
                ));
            }
        }

        info!("All models synchronized at epoch {}", epoch);
        Ok(())
    }

    /// Set model performance metrics
    pub async fn set_model_performance(
        &self,
        model_name: &str,
        accuracy: f64,
        loss: f64,
    ) -> Result<()> {
        let mut states = self.model_states.write().await;

        if let Some(state) = states.get_mut(model_name) {
            state.performance = Some(ModelPerformance {
                accuracy,
                loss,
                sharpe_ratio: accuracy * 2.0 - 1.0, // Simple approximation
                validation_loss: loss * 1.1,
                epoch: state.current_epoch,
                updated_at: Utc::now(),
            });
            debug!(
                "Updated performance for {}: accuracy={:.3}, loss={:.3}",
                model_name, accuracy, loss
            );
            Ok(())
        } else {
            Err(anyhow!("Model {} not found", model_name))
        }
    }

    /// Optimize ensemble weights based on model performance
    pub async fn optimize_weights(&self) -> Result<()> {
        info!("Optimizing ensemble weights based on model performance");

        let states = self.model_states.read().await;

        // Calculate performance-based weights
        let mut new_weights = HashMap::new();
        let mut total_score = 0.0;

        for (model_name, state) in states.iter() {
            if let Some(perf) = &state.performance {
                // Performance score: accuracy weighted by inverse loss
                let score = perf.accuracy / (1.0 + perf.loss);
                new_weights.insert(model_name.clone(), score);
                total_score += score;
            } else {
                // Keep original weight if no performance data
                let weights = self.current_weights.read().await;
                new_weights.insert(
                    model_name.clone(),
                    *weights.get(model_name).unwrap_or(&0.25),
                );
            }
        }

        drop(states);

        // Normalize weights to sum to 1.0
        if total_score > 0.0 {
            for weight in new_weights.values_mut() {
                *weight /= total_score;
            }
        }

        // Update current weights
        {
            let mut weights = self.current_weights.write().await;
            *weights = new_weights.clone();
        }

        info!("Updated ensemble weights: {:?}", new_weights);
        Ok(())
    }

    /// Simulate model failure (for testing)
    pub async fn simulate_model_failure(&self, model_name: &str) -> Result<()> {
        warn!("Simulating failure for model: {}", model_name);

        let mut states = self.model_states.write().await;
        if let Some(state) = states.get_mut(model_name) {
            state.status = ModelTrainingStatus::Failed;
            state.error_message = Some("Simulated failure for testing".to_string());
            Ok(())
        } else {
            Err(anyhow!("Model {} not found", model_name))
        }
    }

    /// Retry failed model
    pub async fn retry_failed_model(&self, model_name: &str) -> Result<()> {
        info!("Retrying model with previous training issues: {}", model_name);

        let mut states = self.model_states.write().await;
        if let Some(state) = states.get_mut(model_name) {
            if state.status == ModelTrainingStatus::Failed {
                state.status = ModelTrainingStatus::Training;
                state.error_message = None;
                info!("Model {} reset to Training status", model_name);
                Ok(())
            } else {
                Err(anyhow!(
                    "Model {} is not in Failed state (current: {:?})",
                    model_name,
                    state.status
                ))
            }
        } else {
            Err(anyhow!("Model {} not found", model_name))
        }
    }

    /// Get ensemble-level metrics
    pub async fn get_ensemble_metrics(&self) -> Result<HashMap<String, f64>> {
        let states = self.model_states.read().await;
        let weights = self.current_weights.read().await;

        let mut metrics = HashMap::new();

        // Calculate weighted ensemble metrics
        let mut weighted_loss = 0.0;
        let mut weighted_val_loss = 0.0;
        let mut weighted_accuracy = 0.0;
        let mut total_weight = 0.0;

        for (model_name, state) in states.iter() {
            if let Some(perf) = &state.performance {
                let weight = weights.get(model_name).unwrap_or(&0.25);
                weighted_loss += perf.loss * weight;
                weighted_val_loss += perf.validation_loss * weight;
                weighted_accuracy += perf.accuracy * weight;
                total_weight += weight;
            }
        }

        if total_weight > 0.0 {
            metrics.insert("ensemble_train_loss".to_string(), weighted_loss);
            metrics.insert("ensemble_val_loss".to_string(), weighted_val_loss);
            metrics.insert("ensemble_accuracy".to_string(), weighted_accuracy);
        }

        // Calculate prediction diversity (variance in predictions)
        let performances: Vec<_> = states
            .values()
            .filter_map(|s| s.performance.as_ref())
            .collect();

        if performances.len() > 1 {
            let mean_accuracy: f64 =
                performances.iter().map(|p| p.accuracy).sum::<f64>() / performances.len() as f64;

            let variance: f64 = performances
                .iter()
                .map(|p| (p.accuracy - mean_accuracy).powi(2))
                .sum::<f64>()
                / performances.len() as f64;

            let diversity = variance.sqrt().min(1.0); // Normalize to [0, 1]
            metrics.insert("prediction_diversity".to_string(), diversity);
        } else {
            metrics.insert("prediction_diversity".to_string(), 0.0);
        }

        Ok(metrics)
    }

    /// Get model training configuration
    pub async fn get_model_training_config(
        &self,
        model_name: &str,
    ) -> Result<ProductionTrainingConfig> {
        self.config
            .model_configs
            .get(model_name)
            .cloned()
            .ok_or_else(|| anyhow!("Model {} not found in configuration", model_name))
    }

    /// Get job ID
    pub fn job_id(&self) -> Uuid {
        self.config.job_id
    }

    /// Get number of models
    pub fn model_count(&self) -> usize {
        self.config.model_count()
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use ml::safety::{GradientSafetyConfig, MLSafetyConfig};
    use ml::training_pipeline::FinancialValidationConfig;
    use ml::training_pipeline::{
        ModelArchitectureConfig, PerformanceConfig, TrainingHyperparameters,
    };

    fn create_test_model_config(_model_type: &str) -> ProductionTrainingConfig {
        ProductionTrainingConfig {
            model_config: ModelArchitectureConfig {
                input_dim: 64,
                hidden_dims: vec![256, 128],
                output_dim: 32,
                dropout_rate: 0.1,
                activation: "relu".to_string(),
                batch_norm: true,
                residual_connections: false,
            },
            training_params: TrainingHyperparameters {
                learning_rate: 0.001,
                batch_size: 64,
                max_epochs: 100,
                patience: 10,
                validation_split: 0.2,
                l2_regularization: 0.0001,
                lr_decay_factor: 0.5,
                lr_decay_patience: 5,
            },
            safety_config: MLSafetyConfig {
                safety_enabled: true,
                max_tensor_elements: 100_000_000,
                max_inference_timeout_ms: 5000,
                max_gpu_memory_bytes: 2_000_000_000,
                drift_sensitivity: 0.5,
                financial_precision: 2,
                nan_infinity_checks: true,
                max_prediction_value: 100.0,
                min_prediction_value: -100.0,
                bounds_checking: true,
                auto_fallback: true,
                max_retries: 3,
            },
            gradient_config: GradientSafetyConfig {
                max_gradient_norm: 1.0,
                min_gradient_norm: 1e-8,
                max_individual_gradient: 5.0,
                enable_norm_clipping: true,
                enable_value_clipping: true,
                enable_nan_detection: true,
                gradient_history_size: 100,
                explosion_threshold: 2.0,
                min_gradient_history: 10,
                enable_adaptive_scaling: true,
                lr_adjustment_factor: 0.5,
                base_learning_rate: 0.001,
            },
            financial_config: FinancialValidationConfig {
                max_prediction_multiple: 2.0,
                min_prediction_confidence: 0.6,
                validate_position_sizing: true,
                max_position_fraction: 0.2,
                min_sharpe_threshold: 0.5,
            },
            performance_config: PerformanceConfig {
                device_preference: "cpu".to_string(),
                max_memory_bytes: 4_000_000_000,
                num_workers: 2,
                gradient_accumulation_steps: 1,
            },
        }
    }

    fn create_test_ensemble_config() -> EnsembleTrainingConfig {
        let mut model_configs = HashMap::new();
        let mut model_weights = HashMap::new();

        for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
            model_configs.insert(model_name.to_string(), create_test_model_config(model_name));
        }

        model_weights.insert("DQN".to_string(), 0.33);
        model_weights.insert("PPO".to_string(), 0.33);
        model_weights.insert("MAMBA2".to_string(), 0.17);
        model_weights.insert("TFT".to_string(), 0.17);

        EnsembleTrainingConfig {
            job_id: Uuid::new_v4(),
            model_configs,
            model_weights,
            enable_weight_optimization: true,
            weight_optimization_interval_epochs: 5,
            checkpoint_interval_epochs: 1,
            max_epochs: 100,
            parallel_training: false,
            created_at: Utc::now(),
        }
    }

    #[tokio::test]
    async fn test_coordinator_creation() {
        let config = create_test_ensemble_config();
        let coordinator = EnsembleTrainingCoordinator::new(config).await;

        assert!(coordinator.is_ok());
        let coord = coordinator.unwrap();
        assert_eq!(coord.model_count(), 4);
    }

    #[tokio::test]
    async fn test_config_validation() {
        let config = create_test_ensemble_config();
        assert!(config.is_valid());
        assert_eq!(config.model_count(), 4);
    }

    #[tokio::test]
    async fn test_weights_sum_to_one() {
        let config = create_test_ensemble_config();
        let weight_sum = config.total_weight();
        assert!((weight_sum - 1.0).abs() < 1e-6);
    }
}
