//! MAMBA-2 Trainer with gRPC Integration
//!
//! Production-ready MAMBA-2 trainer optimized for 4GB VRAM constraint with:
//! - GPU acceleration (RTX 3050 Ti)
//! - Gradient checkpointing for memory efficiency
//! - MinIO checkpoint management
//! - Real-time metrics streaming
//! - Sequence modeling on market data
//!
//! This trainer wraps the existing MAMBA-2 training implementation and provides
//! gRPC-compatible interfaces for the ML Training Service.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Instant, SystemTime};

use ml_core::device::MlDevice;
use ml_core::cuda_autograd::stream_ops::StreamTensor;
// Mamba2SSM uses StreamTensor internally (aliased as GpuTensor in ml-supervised).
type GpuTensor = StreamTensor;
use serde::{Deserialize, Serialize};
use tracing::info;
use uuid::Uuid;

use crate::mamba::{Mamba2Config, Mamba2SSM, OptimizerType, TrainingEpoch};
use crate::MLError;

/// MAMBA-2 Hyperparameters for gRPC interface
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Mamba2Hyperparameters {
    /// Learning rate (1e-6 to 1e-3)
    pub learning_rate: f64,
    /// Batch size (1-16 for 4GB VRAM)
    pub batch_size: usize,
    /// Hidden dimension (256, 512, 1024)
    pub d_model: usize,
    /// Number of layers (4-12)
    pub n_layers: usize,
    /// SSM state dimension (16-64)
    pub state_size: usize,
    /// Dropout rate (0.0-0.3)
    pub dropout: f64,
    /// Number of training epochs
    pub epochs: usize,
    /// Sequence length for training
    pub seq_len: usize,
    /// Gradient clipping threshold
    pub grad_clip: f64,
    /// Weight decay for regularization
    pub weight_decay: f64,
    /// Warmup steps for learning rate
    pub warmup_steps: usize,
}

impl Default for Mamba2Hyperparameters {
    fn default() -> Self {
        Self {
            learning_rate: 1e-4,
            batch_size: 8, // Conservative for 4GB VRAM
            d_model: 256,  // Small model for memory efficiency
            n_layers: 6,
            state_size: 32,
            dropout: 0.1,
            epochs: 100,
            seq_len: 128, // Short sequences for memory
            grad_clip: 1.0,
            weight_decay: 1e-4,
            warmup_steps: 1000,
        }
    }
}

impl Mamba2Hyperparameters {
    /// Validate hyperparameters against detected GPU capacity.
    ///
    /// Replaces the hardcoded 4GB/16-batch limits with dynamic detection via
    /// `HardwareBudget::detect()`, so the same code works on RTX 3050 Ti (4GB),
    /// L40S (48GB), and H100 (80GB) without manual tuning.
    pub fn validate(&self) -> Result<(), MLError> {
        if self.batch_size == 0 {
            return Err(MLError::InvalidInput(
                "Batch size must be greater than 0".to_owned(),
            ));
        }

        let budget = crate::hyperopt::HardwareBudget::detect();
        let estimated_memory_mb = self.estimate_memory_usage();

        // Memory check: 85% safety margin of detected VRAM (or 3500MB for CPU-only)
        let safe_limit_mb = if budget.gpu_memory_mb > 0 {
            (budget.gpu_memory_mb as f64 * 0.85) as usize
        } else {
            3500 // CPU-only fallback — no VRAM constraint, but keep sane limit
        };

        if estimated_memory_mb > safe_limit_mb {
            return Err(MLError::InvalidInput(format!(
                "Estimated memory {}MB exceeds safe limit {}MB ({} {}MB VRAM)",
                estimated_memory_mb,
                safe_limit_mb,
                if budget.gpu_name.is_empty() { "CPU" } else { &budget.gpu_name },
                budget.gpu_memory_mb,
            )));
        }

        // Dynamic batch size limit based on GPU
        if let Some(max_batch) = budget.max_batch_size(
            estimated_memory_mb as f64,
            0.5, // ~0.5MB per sample for Mamba2 (seq_len=128, d_model=256)
            1.0,
            256.0,
        ) {
            if self.batch_size > max_batch as usize {
                return Err(MLError::InvalidInput(format!(
                    "Batch size {} exceeds GPU-scaled max {} for {} ({}MB VRAM)",
                    self.batch_size,
                    max_batch as usize,
                    budget.gpu_name,
                    budget.gpu_memory_mb,
                )));
            }
        }

        if !(1e-6..=1e-3).contains(&self.learning_rate) {
            return Err(MLError::InvalidInput(
                "Learning rate must be between 1e-6 and 1e-3".to_owned(),
            ));
        }

        if ![256, 512, 1024].contains(&self.d_model) {
            return Err(MLError::InvalidInput(
                "d_model must be 256, 512, or 1024".to_owned(),
            ));
        }

        if !(4..=12).contains(&self.n_layers) {
            return Err(MLError::InvalidInput(
                "n_layers must be between 4 and 12".to_owned(),
            ));
        }

        if !(16..=64).contains(&self.state_size) {
            return Err(MLError::InvalidInput(
                "state_size must be between 16 and 64".to_owned(),
            ));
        }

        if !(0.0..=0.3).contains(&self.dropout) {
            return Err(MLError::InvalidInput(
                "dropout must be between 0.0 and 0.3".to_owned(),
            ));
        }

        Ok(())
    }

    /// Estimate memory usage in MB for VRAM constraint validation
    pub fn estimate_memory_usage(&self) -> usize {
        // Model parameters: d_model * n_layers * state_size * 4 bytes (f32)
        let model_params = self.d_model * self.n_layers * self.state_size * 4;

        // Activations: batch_size * seq_len * d_model * n_layers * 4 bytes
        let activations = self.batch_size * self.seq_len * self.d_model * self.n_layers * 4;

        // Gradients (same as model params)
        let gradients = model_params;

        // Optimizer states (momentum + variance for Adam = 2x params)
        let optimizer_states = model_params * 2;

        // Total with 20% overhead for misc tensors
        let total_bytes = (model_params + activations + gradients + optimizer_states) * 12 / 10;

        // Convert to MB
        total_bytes / (1024 * 1024)
    }

    /// Convert to MAMBA-2 config
    pub fn to_mamba_config(&self) -> Mamba2Config {
        Mamba2Config {
            d_model: self.d_model,
            d_state: self.state_size,
            d_head: self.d_model / 8, // 8 heads by default
            num_heads: 8,
            expand: 2, // Standard expansion factor
            num_layers: self.n_layers,
            dropout: self.dropout,
            use_ssd: true,
            use_selective_state: true,
            hardware_aware: true,
            target_latency_us: 5,
            max_seq_len: self.seq_len * 2, // Allow some flexibility
            learning_rate: self.learning_rate,
            weight_decay: self.weight_decay,
            grad_clip: self.grad_clip,
            warmup_steps: self.warmup_steps,
            adam_beta1: 0.9,
            adam_beta2: 0.999,        // P1: Standard Adam beta2
            adam_epsilon: 1e-8,       // P1: Standard Adam epsilon
            total_decay_steps: 10000, // P1: Standard decay schedule
            optimizer_type: OptimizerType::Adam,
            sgd_momentum: 0.9,
            batch_size: self.batch_size,
            seq_len: self.seq_len,
            shuffle_batches: false,
            sequence_stride: 1,             // P2: No overlapping (safe default)
            norm_eps: 1e-5,                 // P2: Standard layer norm epsilon
            early_stopping_enabled: true,   // Enable early stopping by default
            early_stopping_patience: 20,    // 20 epochs patience
            early_stopping_min_delta: 1e-4, // Minimum improvement threshold
            early_stopping_min_epochs: 20,  // Minimum 20 epochs before stopping
            spsa_epsilon: 0.01,              // SPSA perturbation magnitude (Spall 1992)
        }
    }
}

/// Training metrics for real-time monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingMetrics {
    /// Current training loss
    pub loss: f64,
    /// Validation loss
    pub val_loss: f64,
    /// Perplexity (exp(loss))
    pub perplexity: f64,
    /// Learning rate
    pub learning_rate: f64,
    /// Sequence loss (sequence modeling specific)
    pub sequence_loss: f64,
    /// Average SSM state magnitude
    pub state_magnitude: f64,
    /// Memory usage (GB)
    pub memory_usage_gb: f64,
    /// Training throughput (samples/sec)
    pub throughput: f64,
    /// Epoch duration (seconds)
    pub epoch_duration: f64,
}

impl Default for TrainingMetrics {
    fn default() -> Self {
        Self {
            loss: 0.0,
            val_loss: 0.0,
            perplexity: 1.0,
            learning_rate: 0.0,
            sequence_loss: 0.0,
            state_magnitude: 0.0,
            memory_usage_gb: 0.0,
            throughput: 0.0,
            epoch_duration: 0.0,
        }
    }
}

/// Training progress callback for real-time updates
pub type ProgressCallback = Arc<dyn Fn(TrainingProgress) + Send + Sync>;

/// Training progress update
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingProgress {
    pub job_id: String,
    pub epoch: usize,
    pub total_epochs: usize,
    pub progress_percentage: f64,
    pub metrics: TrainingMetrics,
    pub timestamp: SystemTime,
}

/// MAMBA-2 Trainer with GPU optimization and checkpoint management
///
/// This trainer provides a gRPC-compatible interface around the MAMBA-2 model's
/// existing training implementation. It handles hyperparameter validation,
/// progress callbacks, and checkpoint management.
pub struct Mamba2Trainer {
    /// Unique training job ID
    pub job_id: String,
    /// MAMBA-2 model
    pub model: Mamba2SSM,
    /// Hyperparameters
    pub hyperparameters: Mamba2Hyperparameters,
    /// GPU device (if available)
    pub device: MlDevice,
    /// Training history
    pub training_history: Vec<TrainingEpoch>,
    /// Progress callback for real-time updates
    pub progress_callback: Option<ProgressCallback>,
    /// Checkpoint storage path (MinIO)
    pub checkpoint_path: String,
    /// Training start time
    pub start_time: Option<Instant>,
    /// Best validation loss
    pub best_val_loss: f64,
}

impl std::fmt::Debug for Mamba2Trainer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Mamba2Trainer")
            .field("job_id", &self.job_id)
            .field("model", &"<Mamba2SSM>")
            .field("hyperparameters", &self.hyperparameters)
            .field("device", &self.device)
            .field("training_history", &self.training_history)
            .field(
                "progress_callback",
                &self.progress_callback.as_ref().map(|_| "<callback>"),
            )
            .field("checkpoint_path", &self.checkpoint_path)
            .field("start_time", &self.start_time)
            .field("best_val_loss", &self.best_val_loss)
            .finish()
    }
}

impl Mamba2Trainer {
    /// Create new MAMBA-2 trainer with GPU support
    ///
    /// # Errors
    ///
    /// Returns `MLError` if:
    /// - Hyperparameters are invalid
    /// - Model creation fails
    /// - GPU initialization fails (falls back to CPU)
    pub fn new(
        hyperparameters: Mamba2Hyperparameters,
        checkpoint_path: Option<String>,
    ) -> Result<Self, MLError> {
        // Validate hyperparameters for VRAM constraint
        hyperparameters.validate()?;

        let estimated_memory = hyperparameters.estimate_memory_usage();
        info!(
            "Creating MAMBA-2 trainer (estimated VRAM: {}MB)",
            estimated_memory
        );

        let device = MlDevice::cuda(0)
            .map_err(|e| MLError::ConfigError(format!("CUDA GPU required for Mamba2 training: {}", e)))?;
        info!("Using CUDA device for MAMBA-2 training");

        // Create MAMBA-2 model
        let config = hyperparameters.to_mamba_config();
        let stream = device.cuda_stream()?;
        let model = Mamba2SSM::new(config, stream)?;

        let job_id = Uuid::new_v4().to_string();
        let checkpoint_path =
            checkpoint_path.unwrap_or_else(|| format!("s3://foxhunt-ml-models/mamba2/{}", job_id));

        Ok(Self {
            job_id,
            model,
            hyperparameters,
            device,
            training_history: Vec::new(),
            progress_callback: None,
            checkpoint_path,
            start_time: None,
            best_val_loss: f64::INFINITY,
        })
    }

    /// Set progress callback for real-time updates
    pub fn set_progress_callback(&mut self, callback: ProgressCallback) {
        self.progress_callback = Some(callback);
    }

    /// Train model with market data sequences
    ///
    /// This method delegates to the MAMBA-2 model's existing training implementation
    /// and provides progress callbacks for gRPC streaming.
    ///
    /// # Arguments
    ///
    /// * `train_data` - Training sequences (input, target) pairs
    /// * `val_data` - Validation sequences (input, target) pairs
    ///
    /// # Returns
    ///
    /// Returns training history with metrics for each epoch
    ///
    /// # Errors
    ///
    /// Returns `MLError` if training fails
    pub async fn train(
        &mut self,
        train_data: &[(GpuTensor, GpuTensor)],
        val_data: &[(GpuTensor, GpuTensor)],
    ) -> Result<Vec<TrainingEpoch>, MLError> {
        info!(
            "Starting MAMBA-2 training: {} epochs, {} training samples, {} validation samples",
            self.hyperparameters.epochs,
            train_data.len(),
            val_data.len()
        );

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

        // Delegate to MAMBA-2 model's training implementation
        let training_history = self
            .model
            .train(train_data, val_data, self.hyperparameters.epochs, None)
            .await?;

        // Update trainer state
        self.training_history = training_history.clone();
        if let Some(last_epoch) = training_history.last() {
            self.best_val_loss = last_epoch.loss;
        }

        // Send final progress update
        if let Some(callback) = &self.progress_callback {
            let final_metrics = self.get_final_metrics();
            let progress = TrainingProgress {
                job_id: self.job_id.clone(),
                epoch: training_history.len(),
                total_epochs: self.hyperparameters.epochs,
                progress_percentage: 100.0,
                metrics: final_metrics,
                timestamp: SystemTime::now(),
            };
            callback(progress);
        }

        info!(
            "Training completed: {} epochs, best loss: {:.6}",
            training_history.len(),
            self.best_val_loss
        );

        Ok(training_history)
    }

    /// Get final training metrics
    fn get_final_metrics(&self) -> TrainingMetrics {
        let model_metrics = self.model.get_performance_metrics();

        TrainingMetrics {
            loss: self.best_val_loss,
            val_loss: self.best_val_loss,
            perplexity: self.best_val_loss.exp(),
            learning_rate: self.hyperparameters.learning_rate,
            sequence_loss: self.best_val_loss,
            state_magnitude: model_metrics
                .get("state_compression_ratio")
                .copied()
                .unwrap_or(1.0),
            memory_usage_gb: self.hyperparameters.estimate_memory_usage() as f64 / 1024.0,
            throughput: model_metrics.get("throughput_pps").copied().unwrap_or(0.0),
            epoch_duration: self
                .start_time
                .map(|t| t.elapsed().as_secs_f64())
                .unwrap_or(0.0),
        }
    }

    /// Get training statistics for gRPC response
    pub fn get_training_statistics(&self) -> HashMap<String, f64> {
        let mut stats = HashMap::new();

        if let Some(start_time) = self.start_time {
            let elapsed = start_time.elapsed().as_secs_f64();
            stats.insert("training_duration_seconds".to_owned(), elapsed);
        }

        stats.insert("best_val_loss".to_owned(), self.best_val_loss);
        stats.insert("best_perplexity".to_owned(), self.best_val_loss.exp());
        stats.insert(
            "total_epochs".to_owned(),
            self.training_history.len() as f64,
        );
        stats.insert(
            "estimated_memory_mb".to_owned(),
            self.hyperparameters.estimate_memory_usage() as f64,
        );

        // Model-specific statistics
        let model_metrics = self.model.get_performance_metrics();
        for (key, value) in model_metrics {
            stats.insert(key, value);
        }

        stats
    }
}

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

    #[test]
    fn test_hyperparameters_validation() {
        let params = Mamba2Hyperparameters::default();
        assert!(params.validate().is_ok());

        // Test invalid learning rate
        let mut invalid_params = params.clone();
        invalid_params.learning_rate = 1.0; // Too high
        assert!(invalid_params.validate().is_err());

        // Test zero batch size is always rejected
        let mut invalid_params = params.clone();
        invalid_params.batch_size = 0;
        assert!(invalid_params.validate().is_err());

        // Test invalid d_model
        let mut invalid_params = params.clone();
        invalid_params.d_model = 333; // Not in {256, 512, 1024}
        assert!(invalid_params.validate().is_err());
    }

    #[test]
    fn test_memory_estimation() {
        let params = Mamba2Hyperparameters {
            d_model: 256,
            n_layers: 6,
            state_size: 32,
            batch_size: 8,
            seq_len: 128,
            ..Default::default()
        };

        let memory_mb = params.estimate_memory_usage();
        assert!(
            memory_mb < 3500,
            "Memory usage {}MB exceeds 4GB constraint",
            memory_mb
        );
        assert!(memory_mb > 0, "Memory estimation returned 0");
    }

    #[test]
    fn test_config_conversion() {
        let params = Mamba2Hyperparameters::default();
        let config = params.to_mamba_config();

        assert_eq!(config.d_model, params.d_model);
        assert_eq!(config.num_layers, params.n_layers);
        assert_eq!(config.d_state, params.state_size);
        assert_eq!(config.learning_rate, params.learning_rate);
        assert_eq!(config.batch_size, params.batch_size);
    }

    #[tokio::test]
    async fn test_trainer_creation() {
        let params = Mamba2Hyperparameters::default();
        let trainer = Mamba2Trainer::new(params, None);

        assert!(trainer.is_ok());
        let trainer = trainer.unwrap();
        assert!(!trainer.job_id.is_empty());
        assert!(trainer.checkpoint_path.contains("s3://"));
    }

    #[tokio::test]
    async fn test_zero_batch_size_handling() {
        // Test MAMBA-2 rejects zero batch size
        let mut params = Mamba2Hyperparameters::default();
        params.batch_size = 0;

        let result = params.validate();

        // Should fail with descriptive error
        assert!(
            result.is_err(),
            "MAMBA-2 should reject zero batch size, but got: {:?}",
            result
        );

        // Error message should mention batch size
        let error_msg = result.unwrap_err().to_string();
        assert!(
            error_msg.to_lowercase().contains("batch"),
            "Error message should mention batch size, got: {}",
            error_msg
        );
    }
}
