//! Unified Training Orchestrator
//!
//! Model-agnostic training loop for all ML models (MAMBA-2, DQN, PPO, TFT).
//! Handles common training logic: batch iteration, validation, checkpointing,
//! early stopping, learning rate scheduling, and metrics aggregation.

use std::path::PathBuf;
use std::time::Instant;

use serde::{Deserialize, Serialize};
use tracing::{debug, info, warn};

use super::unified_trainer::{checkpoint, UnifiedTrainable};
use crate::MLError;

/// Orchestrator configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OrchestratorConfig {
    /// Number of training epochs
    pub num_epochs: usize,
    /// Validation frequency (every N steps)
    pub validation_frequency: usize,
    /// Checkpoint frequency (every N steps)
    pub checkpoint_frequency: usize,
    /// Checkpoint directory
    pub checkpoint_dir: PathBuf,
    /// Early stopping patience (epochs without improvement)
    pub early_stopping_patience: Option<usize>,
    /// Learning rate scheduling strategy
    pub lr_schedule: LRSchedule,
    /// Enable gradient accumulation
    pub gradient_accumulation_steps: usize,
    /// Maximum gradient norm for clipping
    pub max_grad_norm: Option<f64>,
}

impl Default for OrchestratorConfig {
    fn default() -> Self {
        Self {
            num_epochs: 100,
            validation_frequency: 100,
            checkpoint_frequency: 1000,
            checkpoint_dir: PathBuf::from("checkpoints"),
            early_stopping_patience: Some(10),
            lr_schedule: LRSchedule::Constant,
            gradient_accumulation_steps: 1,
            max_grad_norm: Some(1.0),
        }
    }
}

/// Learning rate scheduling strategies
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LRSchedule {
    /// Constant learning rate (no scheduling)
    Constant,
    /// Linear warmup then constant
    WarmupConstant { warmup_steps: usize },
    /// Cosine annealing with warmup
    CosineAnnealing {
        warmup_steps: usize,
        total_steps: usize,
        min_lr: f64,
    },
    /// Step decay (reduce by factor every N steps)
    StepDecay { step_size: usize, gamma: f64 },
}

/// Training history for a single epoch
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpochHistory {
    pub epoch: usize,
    pub train_loss: f64,
    pub val_loss: Option<f64>,
    pub accuracy: Option<f64>,
    pub learning_rate: f64,
    pub duration_secs: f64,
}

/// Unified training orchestrator
#[derive(Debug)]
pub struct UnifiedTrainingOrchestrator {
    config: OrchestratorConfig,
    current_step: usize,
    current_epoch: usize,
    best_val_loss: f64,
    epochs_without_improvement: usize,
    training_history: Vec<EpochHistory>,
    initial_lr: f64,
}

impl UnifiedTrainingOrchestrator {
    /// Create new orchestrator
    pub fn new(config: OrchestratorConfig) -> Self {
        Self {
            config,
            current_step: 0,
            current_epoch: 0,
            best_val_loss: f64::INFINITY,
            epochs_without_improvement: 0,
            training_history: Vec::new(),
            initial_lr: 1e-4, // Will be set from model
        }
    }

    /// Train a model using unified interface
    ///
    /// # Arguments
    /// * `model` - Any model implementing UnifiedTrainable trait
    /// * `train_data` - Training dataset (input, target) pairs
    /// * `val_data` - Validation dataset (input, target) pairs
    ///
    /// # Returns
    /// Training history for all epochs
    #[allow(clippy::unwrap_in_result)]
    pub fn train<M: UnifiedTrainable>(
        &mut self,
        model: &mut M,
        train_data: &[(Vec<f32>, Vec<f32>)],
        val_data: &[(Vec<f32>, Vec<f32>)],
    ) -> Result<Vec<EpochHistory>, MLError> {
        info!(
            "Starting unified training for model: {}",
            model.model_type()
        );
        info!(
            "Total epochs: {}, Training samples: {}, Validation samples: {}",
            self.config.num_epochs,
            train_data.len(),
            val_data.len()
        );

        // Create checkpoint directory
        std::fs::create_dir_all(&self.config.checkpoint_dir).map_err(|e| {
            MLError::ModelError(format!("Failed to create checkpoint directory: {}", e))
        })?;

        // Store initial learning rate
        self.initial_lr = model.get_learning_rate();

        // Training loop
        for epoch in 0..self.config.num_epochs {
            self.current_epoch = epoch;
            let epoch_start = Instant::now();

            // Train for one epoch
            let train_loss = self.train_epoch(model, train_data)?;

            // Validation (compute average loss over val_data without backward pass)
            let val_loss = if epoch % (self.config.validation_frequency / 100).max(1) == 0 && !val_data.is_empty() {
                let mut val_total = 0.0_f64;
                for (input, target) in val_data {
                    val_total += model.forward_loss(input, target)?;
                }
                Some(val_total / val_data.len() as f64)
            } else {
                None
            };

            // Collect metrics
            let metrics = model.collect_metrics();
            let learning_rate = model.get_learning_rate();

            // Update learning rate schedule
            self.update_learning_rate(model)?;

            // Record epoch history
            let epoch_duration = epoch_start.elapsed().as_secs_f64();
            let history = EpochHistory {
                epoch,
                train_loss,
                val_loss,
                accuracy: metrics.accuracy,
                learning_rate,
                duration_secs: epoch_duration,
            };
            self.training_history.push(history.clone());

            info!(
                "Epoch {}/{}: train_loss={:.6}, val_loss={:?}, accuracy={:?}, lr={:.2e}, time={:.2}s",
                epoch + 1, self.config.num_epochs, train_loss, val_loss, metrics.accuracy,
                learning_rate, epoch_duration
            );

            // Early stopping check
            if let Some(vl) = val_loss {
                if vl < self.best_val_loss {
                    self.best_val_loss = vl;
                    self.epochs_without_improvement = 0;

                    // Save best checkpoint
                    let checkpoint_path = self
                        .config
                        .checkpoint_dir
                        .join(format!("{}_best", model.model_type()));
                    let path_str = checkpoint_path.to_str().ok_or_else(|| {
                        MLError::ModelError("Invalid UTF-8 in checkpoint path".to_owned())
                    })?;
                    model.save_checkpoint(path_str)?;
                    info!("New best validation loss: {:.6}, checkpoint saved", vl);
                } else {
                    self.epochs_without_improvement += 1;

                    if let Some(patience) = self.config.early_stopping_patience {
                        if self.epochs_without_improvement >= patience {
                            info!(
                                "Early stopping triggered after {} epochs without improvement",
                                patience
                            );
                            break;
                        }
                    }
                }
            }

            // Periodic checkpointing
            if (epoch + 1) % (self.config.checkpoint_frequency / 100).max(1) == 0 {
                let checkpoint_path =
                    self.config
                        .checkpoint_dir
                        .join(checkpoint::checkpoint_filename(
                            model.model_type(),
                            epoch,
                            self.current_step,
                        ));
                let path_str = checkpoint_path.to_str().ok_or_else(|| {
                    MLError::ModelError("Invalid UTF-8 in checkpoint path".to_owned())
                })?;
                model.save_checkpoint(path_str)?;
                debug!("Checkpoint saved at epoch {}", epoch + 1);
            }
        }

        info!(
            "Training completed. Best validation loss: {:.6}",
            self.best_val_loss
        );

        Ok(self.training_history.clone())
    }

    /// Train for one epoch
    fn train_epoch<M: UnifiedTrainable>(
        &mut self,
        model: &mut M,
        train_data: &[(Vec<f32>, Vec<f32>)],
    ) -> Result<f64, MLError> {
        let mut total_loss = 0.0;
        let mut batch_count = 0;

        // Gradient accumulation buffer
        let mut accumulated_loss = 0.0;
        let mut accumulation_steps = 0;

        for (batch_idx, (input, target)) in train_data.iter().enumerate() {
            // Forward pass + loss computation via unified trait
            let loss_value = model.forward_loss(input, target)?;

            // Check for NaN
            if loss_value.is_nan() || loss_value.is_infinite() {
                return Err(MLError::TrainingError(format!(
                    "NaN or Inf detected in loss at epoch {}, batch {}",
                    self.current_epoch, batch_idx
                )));
            }

            // Accumulate loss
            accumulated_loss += loss_value;
            accumulation_steps += 1;

            // Backward pass (compute gradients)
            let grad_norm = model.backward(loss_value)?;

            // Gradient clipping (if enabled)
            if let Some(max_norm) = self.config.max_grad_norm {
                if grad_norm > max_norm {
                    warn!(
                        "Gradient norm {:.3} exceeds max {:.3}, clipping applied",
                        grad_norm, max_norm
                    );
                }
            }

            // Optimizer step (with gradient accumulation)
            if accumulation_steps >= self.config.gradient_accumulation_steps {
                model.optimizer_step()?;
                model.zero_grad()?;

                total_loss += accumulated_loss / accumulation_steps as f64;
                batch_count += 1;

                // Reset accumulation
                accumulated_loss = 0.0;
                accumulation_steps = 0;
            }

            self.current_step += 1;

            // Log progress
            if batch_idx % 100 == 0 {
                debug!(
                    "Epoch {}, Batch {}/{}: loss={:.6}",
                    self.current_epoch,
                    batch_idx,
                    train_data.len(),
                    loss_value
                );
            }
        }

        // Handle remaining accumulated gradients
        if accumulation_steps > 0 {
            model.optimizer_step()?;
            model.zero_grad()?;
            total_loss += accumulated_loss / accumulation_steps as f64;
            batch_count += 1;
        }

        Ok(total_loss / batch_count.max(1) as f64)
    }

    /// Update learning rate based on schedule
    fn update_learning_rate<M: UnifiedTrainable>(&self, model: &mut M) -> Result<(), MLError> {
        let new_lr = match &self.config.lr_schedule {
            LRSchedule::Constant => {
                return Ok(()); // No change
            },
            LRSchedule::WarmupConstant { warmup_steps } => {
                if self.current_step < *warmup_steps {
                    self.initial_lr * (self.current_step as f64 / *warmup_steps as f64)
                } else {
                    self.initial_lr
                }
            },
            LRSchedule::CosineAnnealing {
                warmup_steps,
                total_steps,
                min_lr,
            } => {
                if self.current_step < *warmup_steps {
                    self.initial_lr * (self.current_step as f64 / *warmup_steps as f64)
                } else {
                    let progress = ((self.current_step - warmup_steps) as f64
                        / (*total_steps - warmup_steps) as f64)
                        .min(1.0);
                    min_lr
                        + (self.initial_lr - min_lr)
                            * 0.5
                            * (1.0 + (std::f64::consts::PI * progress).cos())
                }
            },
            LRSchedule::StepDecay { step_size, gamma } => {
                let decay_steps = self.current_step / step_size;
                self.initial_lr * gamma.powi(decay_steps as i32)
            },
        };

        model.set_learning_rate(new_lr)?;
        Ok(())
    }

    /// Get training history
    pub fn get_history(&self) -> &[EpochHistory] {
        &self.training_history
    }

    /// Get best validation loss
    pub fn get_best_val_loss(&self) -> f64 {
        self.best_val_loss
    }

    /// Get current step count
    pub fn get_current_step(&self) -> usize {
        self.current_step
    }
}

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

    #[test]
    fn test_orchestrator_creation() {
        let config = OrchestratorConfig::default();
        let orchestrator = UnifiedTrainingOrchestrator::new(config);

        assert_eq!(orchestrator.current_step, 0);
        assert_eq!(orchestrator.current_epoch, 0);
        assert_eq!(orchestrator.best_val_loss, f64::INFINITY);
    }

    #[test]
    fn test_lr_schedule_warmup() {
        let schedule = LRSchedule::WarmupConstant { warmup_steps: 100 };
        assert!(matches!(schedule, LRSchedule::WarmupConstant { .. }));
    }

    #[test]
    fn test_lr_schedule_cosine() {
        let schedule = LRSchedule::CosineAnnealing {
            warmup_steps: 100,
            total_steps: 1000,
            min_lr: 1e-6,
        };
        assert!(matches!(schedule, LRSchedule::CosineAnnealing { .. }));
    }

    #[test]
    fn test_checkpoint_dir_creation() {
        let config = OrchestratorConfig {
            checkpoint_dir: PathBuf::from("/tmp/test_checkpoints"),
            ..Default::default()
        };

        let _orchestrator = UnifiedTrainingOrchestrator::new(config);
        // Directory will be created during training
    }
}
