//! # MAMBA-2 Benchmark Runner - Production GPU Training Benchmarks
//!
//! This module implements comprehensive benchmarking for the MAMBA-2 state-space model
//! using real market data from DBN files. It measures actual GPU training performance,
//! memory usage, stability, and convergence metrics for sequence modeling.
//!
//! ## Architecture
//!
//! ```text
//! ┌──────────────┐     ┌──────────────┐     ┌──────────────┐
//! │  DBN Files   │────▶│ MAMBA Trainer│────▶│  Benchmarks  │
//! │  (Real Data) │     │  (GPU/CPU)   │     │  (Stats)     │
//! └──────────────┘     └──────────────┘     └──────────────┘
//!        │                     │                     │
//!        ▼                     ▼                     ▼
//!   OHLCV Bars          Training Loop         Memory Profile
//!   6E.FUT Data         State Space Model     Stability Stats
//! ```
//!
//! ## Features
//!
//! - Real market data loading from DBN files
//! - GPU warmup and optimal batch size finding
//! - Memory profiling per epoch
//! - Statistical sampling for latency measurements
//! - Stability validation (gradient norms, loss variance)
//! - Comprehensive benchmark results
//! - Sequence-to-sequence training with time series
//!
//! ## Memory Configuration
//!
//! - Target VRAM: 150-500MB
//! - Gradient accumulation for 4GB GPU
//! - State dimension: Dynamic (from data)
//! - Sequence length: 128 (for 1-minute bars)
//! - Batch size: 8-16 (with gradient accumulation to simulate 64)
//!
//! ## Usage
//!
//! ```rust,no_run
//! use ml::benchmark::mamba2_benchmark::Mamba2BenchmarkRunner;
//! use ml::benchmark::gpu_hardware::GpuHardwareManager;
//! use std::sync::Arc;
//!
//! # async fn example() -> anyhow::Result<()> {
//! let gpu_manager = Arc::new(GpuHardwareManager::new()?);
//! let mut runner = Mamba2BenchmarkRunner::new(gpu_manager);
//!
//! let result = runner.run_benchmark(10).await?;
//! println!("MAMBA-2 Benchmark: {:.2}s avg epoch time", result.statistics.mean_seconds);
//! println!("Peak memory: {:.2}MB", result.memory_peak_mb);
//! # Ok(())
//! # }
//! ```

use anyhow::{Context, Result};
use ml_core::cuda_autograd::stream_ops::StreamTensor;
use ml_core::device::MlDevice;

// Type alias: Mamba2SSM uses StreamTensor (aliased as GpuTensor in ml-supervised).
type GpuTensor = StreamTensor;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::Mutex;
use tracing::info;

use crate::mamba::{Mamba2Config, Mamba2SSM, OptimizerType};
use crate::data_loader::RealDataLoader;

use super::batch_size_finder::{BatchSizeConfig, BatchSizeFinder};
use super::gpu_hardware::GpuHardwareManager;
use super::memory_profiler::MemoryProfiler;
use super::stability_validator::{StabilityMetrics, StabilityValidator};
use super::statistical_sampler::{BenchmarkStatistics, StatisticalSampler};

/// MAMBA-2 benchmark result containing all metrics
#[derive(Debug, Clone)]
pub struct Mamba2BenchmarkResult {
    /// Model name identifier
    pub model_name: String,
    /// Total epochs trained
    pub total_epochs: usize,
    /// Statistical summary of epoch times
    pub statistics: BenchmarkStatistics,
    /// Peak memory usage in MB
    pub memory_peak_mb: f64,
    /// Stability metrics (gradient norms, loss variance)
    pub stability: StabilityMetrics,
    /// Optimal batch size configuration
    pub batch_config: BatchSizeConfig,
    /// Training losses per epoch
    pub training_losses: Vec<f64>,
    /// Validation losses per epoch
    pub validation_losses: Vec<f64>,
    /// Average training loss across all epochs
    pub avg_train_loss: f64,
    /// Average validation loss
    pub avg_val_loss: f64,
    /// Final model accuracy (within 10% error threshold)
    pub final_accuracy: f64,
}

/// MAMBA-2 Benchmark Runner
///
/// Runs comprehensive benchmarks on MAMBA-2 SSM using real market data.
/// Measures GPU performance, memory usage, training stability, and convergence.
#[derive(Debug)]
pub struct Mamba2BenchmarkRunner {
    /// GPU hardware manager
    gpu_manager: Arc<GpuHardwareManager>,
    /// Memory profiler
    memory_profiler: Arc<Mutex<MemoryProfiler>>,
    /// Statistical sampler
    statistical_sampler: StatisticalSampler,
    /// Stability validator
    stability_validator: StabilityValidator,
}

impl Mamba2BenchmarkRunner {
    /// Create new MAMBA-2 benchmark runner
    ///
    /// # Arguments
    ///
    /// * `gpu_manager` - GPU hardware manager for device management
    ///
    /// # Returns
    ///
    /// New Mamba2BenchmarkRunner instance
    pub fn new(gpu_manager: Arc<GpuHardwareManager>) -> Self {
        Self {
            gpu_manager,
            memory_profiler: Arc::new(Mutex::new(MemoryProfiler::new(0))), // Device 0
            statistical_sampler: StatisticalSampler::new(3),               // 3 warmup epochs
            stability_validator: StabilityValidator::new(),
        }
    }

    /// Run complete MAMBA-2 benchmark
    ///
    /// # Arguments
    ///
    /// * `epochs` - Number of training epochs to run
    ///
    /// # Returns
    ///
    /// Mamba2BenchmarkResult with all metrics
    pub async fn run_benchmark(&mut self, epochs: usize) -> Result<Mamba2BenchmarkResult> {
        info!("Starting MAMBA-2 benchmark with {} epochs", epochs);

        // Step 1: Load real market data from DBN files
        let (train_data, val_data) = self.load_market_data().await?;
        let state_dim = train_data
            .first()
            .and_then(|(s, _)| s.shape.get(1).copied())
            .unwrap_or(32);
        info!(
            "Loaded {} training samples, {} validation samples, state_dim={}",
            train_data.len(),
            val_data.len(),
            state_dim
        );

        // Step 2: Find optimal batch size
        let batch_config = self.find_optimal_batch_size(state_dim)?;
        info!(
            "Optimal batch size: {}, gradient accumulation: {}",
            batch_config.batch_size, batch_config.gradient_accumulation_steps
        );

        // Step 3: GPU warmup
        let warmup_duration = self.gpu_manager.warmup()?;
        info!("GPU warmup complete in {:?}", warmup_duration);

        // Step 4: Create MAMBA-2 model
        let mut mamba = self.create_mamba_model(state_dim, batch_config.batch_size)?;
        info!(
            "Created MAMBA-2 model on device: {:?}",
            self.gpu_manager.device()
        );

        // Step 5: Training loop with comprehensive metrics
        let mut training_losses = Vec::new();
        let mut validation_losses = Vec::new();

        for epoch in 0..epochs {
            // Take memory snapshot before training
            {
                let mut profiler = self.memory_profiler.lock().await;
                profiler
                    .take_snapshot()
                    .map_err(|e| anyhow::anyhow!("Memory snapshot failed: {}", e))?;
            }

            // Train for one epoch
            let epoch_start = Instant::now();
            let epoch_result = mamba
                .train(&train_data, &val_data, 1, None)
                .await
                .context("Training step failed")?;
            let epoch_time_s = epoch_start.elapsed().as_secs_f64();

            // Extract metrics from training result
            let train_loss = epoch_result.last().map(|e| e.loss).unwrap_or(0.0);

            // Compute validation loss
            let val_loss = self.compute_validation_loss(&mut mamba, &val_data)?;

            // Record metrics
            self.statistical_sampler.add_sample(epoch_time_s);
            training_losses.push(train_loss);
            validation_losses.push(val_loss);
            self.stability_validator.record_loss(train_loss);
            self.stability_validator
                .record_gradient_norm(train_loss.abs()); // Estimate gradient norm from loss

            // Take memory snapshot after training
            {
                let mut profiler = self.memory_profiler.lock().await;
                profiler
                    .take_snapshot()
                    .map_err(|e| anyhow::anyhow!("Memory snapshot failed: {}", e))?;
            }

            // Log progress every 10 epochs
            if epoch % 10 == 0 || epoch == epochs - 1 {
                info!(
                    "Epoch {}/{}: train_loss={:.6}, val_loss={:.6}, time={:.4}s",
                    epoch + 1,
                    epochs,
                    train_loss,
                    val_loss,
                    epoch_time_s
                );
            }
        }

        // Step 6: Compute statistics
        let statistics = self
            .statistical_sampler
            .compute_statistics()
            .map_err(|e| anyhow::anyhow!("Failed to compute statistics: {}", e))?;

        // Step 7: Memory analysis
        let memory_peak_mb = {
            let profiler = self.memory_profiler.lock().await;
            profiler.peak_usage_mb()
        };

        // Step 8: Stability validation
        let stability = self.stability_validator.validate();

        // Step 9: Compute average losses
        let avg_train_loss = training_losses.iter().sum::<f64>() / training_losses.len() as f64;
        let avg_val_loss = validation_losses.iter().sum::<f64>() / validation_losses.len() as f64;

        // Step 10: Compute final accuracy
        let final_accuracy = self.compute_accuracy(&mut mamba, &val_data)?;

        info!("MAMBA-2 Benchmark Complete:");
        info!("  Mean epoch time: {:.4}s", statistics.mean_seconds);
        info!("  Median epoch time (P50): {:.4}s", statistics.p50_median);
        info!("  P95 epoch time: {:.4}s", statistics.p95);
        info!("  P99 epoch time: {:.4}s", statistics.p99);
        info!("  Peak memory: {:.2}MB", memory_peak_mb);
        info!("  Average train loss: {:.6}", avg_train_loss);
        info!("  Average val loss: {:.6}", avg_val_loss);
        info!("  Final accuracy: {:.2}%", final_accuracy * 100.0);
        info!("  Training stable: {}", stability.is_stable);

        Ok(Mamba2BenchmarkResult {
            model_name: "MAMBA-2-SSM".to_owned(),
            total_epochs: epochs,
            statistics,
            memory_peak_mb,
            stability,
            batch_config,
            training_losses,
            validation_losses,
            avg_train_loss,
            avg_val_loss,
            final_accuracy,
        })
    }

    /// Load real market data from DBN files
    ///
    /// Loads OHLCV data and creates training/validation sequences for MAMBA-2.
    async fn load_market_data(&self) -> Result<(Vec<(GpuTensor, GpuTensor)>, Vec<(GpuTensor, GpuTensor)>)> {
        // Create data loader
        let mut loader = RealDataLoader::new_from_workspace()
            .context("Failed to create data loader from workspace")?;

        // Load 6E.FUT data (Euro futures)
        let bars = loader
            .load_symbol_data("6E.FUT")
            .await
            .context("Failed to load 6E.FUT data")?;

        if bars.is_empty() {
            return Err(anyhow::anyhow!("No market data loaded"));
        }

        info!("Loaded {} OHLCV bars from DBN files", bars.len());

        // Extract features
        let features = loader
            .extract_features(&bars)
            .context("Failed to extract features")?;

        // Convert to sequence pairs (input, target) for time series prediction
        let sequences = self.create_sequences(&features)?;

        // Split 80/20 train/val
        let split_idx = (sequences.len() as f64 * 0.8) as usize;
        let train_data = sequences[..split_idx].to_vec(); // cpu-side split
        let val_data = sequences[split_idx..].to_vec(); // cpu-side split

        info!(
            "Created {} training sequences, {} validation sequences",
            train_data.len(),
            val_data.len()
        );

        Ok((train_data, val_data))
    }

    /// Create sequences for MAMBA-2 training
    ///
    /// Converts feature matrix to (input_sequence, target) pairs.
    /// Each sequence is 128 timesteps, predicting the next value.
    fn create_sequences(
        &self,
        features: &crate::data_loader::FeatureMatrix,
    ) -> Result<Vec<(GpuTensor, GpuTensor)>> {
        let seq_len = 128; // Sequence length for MAMBA-2
        let mut sequences = Vec::new();

        // Combine all features into a single feature vector per timestep
        let num_timesteps = features.prices.len();
        if num_timesteps < seq_len + 1 {
            return Err(anyhow::anyhow!(
                "Not enough data for sequence creation: {} timesteps, need {}",
                num_timesteps,
                seq_len + 1
            ));
        }

        // Create sequences
        for i in 0..num_timesteps - seq_len {
            let mut sequence_data = Vec::new();

            // Build sequence of feature vectors
            for t in i..i + seq_len {
                let mut feature_vec = Vec::new();

                // Add OHLCV prices (5 values)
                if t < features.prices.len() {
                    feature_vec.extend_from_slice(&features.prices[t]);
                }

                // Add return (1 value)
                if t < features.returns.len() {
                    feature_vec.push(features.returns[t]);
                }

                // Add volume (1 value)
                if t < features.volume.len() {
                    feature_vec.push(features.volume[t]);
                }

                // Add technical indicators (10 values)
                if t < features.indicators.len() {
                    feature_vec.extend_from_slice(&features.indicators[t]);
                }

                sequence_data.push(feature_vec);
            }

            // Target is the close price at the next timestep
            let target_idx = i + seq_len;
            let target_value = if target_idx < features.prices.len() {
                features.prices[target_idx][3] // Close price at index 3
            } else {
                features.prices[features.prices.len() - 1][3]
            };

            // Convert to GPU tensors
            let ml_dev = MlDevice::cuda(0)
                .map_err(|e| anyhow::anyhow!("CUDA device required: {}", e))?;
            let stream = ml_dev.cuda_stream()
                .map_err(|e| anyhow::anyhow!("No CUDA stream: {}", e))?;

            // Flatten sequence data
            let feature_dim = sequence_data[0].len();
            let flat_data: Vec<f32> = sequence_data.into_iter().flatten().collect();

            let input = GpuTensor::from_vec(flat_data, &[1, seq_len, feature_dim], stream)
                .context("Failed to create input tensor")?;

            let target = GpuTensor::from_vec(vec![target_value], &[1, 1], stream)
                .context("Failed to create target tensor")?;

            sequences.push((input, target));
        }

        Ok(sequences)
    }

    /// Find optimal batch size for MAMBA-2 training
    fn find_optimal_batch_size(&self, _state_dim: usize) -> Result<BatchSizeConfig> {
        let finder = BatchSizeFinder::new(*self.gpu_manager.device());

        // Test function: Try creating a MAMBA-2 model with different batch sizes
        let test_fn = |batch_size: usize| {
            // MAMBA-2 memory heuristic
            if batch_size > 32 {
                Ok(false) // Too large for MAMBA-2 (high memory usage)
            } else if batch_size < 4 {
                Ok(false) // Too small for effective training
            } else {
                Ok(true) // Acceptable range: 4-32
            }
        };

        finder
            .find_optimal_batch_size(test_fn)
            .map_err(|e| anyhow::anyhow!("Batch size finding failed: {}", e))
    }

    /// Create MAMBA-2 model with specified configuration
    fn create_mamba_model(&self, state_dim: usize, batch_size: usize) -> Result<Mamba2SSM> {
        let config = Self::create_mamba_config(state_dim, batch_size);
        let ml_dev = MlDevice::cuda(0)
            .map_err(|e| anyhow::anyhow!("CUDA device required: {}", e))?;
        let stream = ml_dev.cuda_stream()
            .map_err(|e| anyhow::anyhow!("No CUDA stream: {}", e))?;
        Mamba2SSM::new(config, stream)
            .map_err(|e| anyhow::anyhow!("Failed to create MAMBA-2 model: {}", e))
    }

    /// Create MAMBA-2 configuration for benchmarking
    fn create_mamba_config(state_dim: usize, batch_size: usize) -> Mamba2Config {
        Mamba2Config {
            d_model: state_dim,        // Model dimension matches state dimension
            d_state: 16,               // State space dimension (compact for 4GB GPU)
            d_head: 16,                // Head dimension
            num_heads: 2,              // Number of attention heads (reduced for memory)
            expand: 1,                 // No expansion to minimize memory
            num_layers: 4,             // 4 layers for expressiveness
            dropout: 0.1,              // Light dropout for regularization
            use_ssd: true,             // Enable SSD layers
            use_selective_state: true, // Enable selective state mechanism
            hardware_aware: true,      // Enable hardware optimizations
            target_latency_us: 5000,   // 5ms target latency (relaxed for benchmarking)
            max_seq_len: 128,          // 128 timestep sequences
            learning_rate: 1e-4,       // Standard learning rate
            weight_decay: 1e-4,        // Light weight decay
            grad_clip: 1.0,            // Gradient clipping threshold
            warmup_steps: 100,         // 100 step warmup
            adam_beta1: 0.9,           // P0: Adam beta1
            adam_beta2: 0.999,         // P1: Adam beta2
            adam_epsilon: 1e-8,        // P1: Adam epsilon
            total_decay_steps: 10000,  // P1: Decay schedule
            optimizer_type: OptimizerType::Adam,
            sgd_momentum: 0.9,
            batch_size,   // From batch size finder
            seq_len: 128, // Match max_seq_len
            shuffle_batches: false,
            sequence_stride: 1,            // P2: No overlapping
            norm_eps: 1e-5,                // P2: Layer norm epsilon
            early_stopping_enabled: false, // Disable for benchmarking
            early_stopping_patience: 20,
            early_stopping_min_delta: 1e-4,
            early_stopping_min_epochs: 20,
            spsa_epsilon: 0.01, // SPSA perturbation magnitude (Spall 1992)
        }
    }

    /// Compute validation loss on validation set
    fn compute_validation_loss(
        &self,
        mamba: &mut Mamba2SSM,
        val_data: &[(GpuTensor, GpuTensor)],
    ) -> Result<f64> {
        let mut total_loss = 0.0;
        let max_samples = 100.min(val_data.len()); // Limit for speed

        let ml_dev = MlDevice::cuda(0)
            .map_err(|e| anyhow::anyhow!("CUDA device required: {}", e))?;
        let _stream = ml_dev.cuda_stream()
            .map_err(|e| anyhow::anyhow!("No CUDA stream: {}", e))?;

        for (input, target) in val_data.iter().take(max_samples) {
            // Forward pass (no gradients)
            let output = mamba
                .forward(input)
                .map_err(|e| anyhow::anyhow!("Forward pass failed: {}", e))?;

            // Compute MSE loss on CPU (readback acceptable for validation)
            let out_host = output.to_vec() // gpu-exit: validation readback
                .map_err(|e| anyhow::anyhow!("Output readback failed: {}", e))?;
            let tgt_host = target.to_vec() // gpu-exit: validation readback
                .map_err(|e| anyhow::anyhow!("Target readback failed: {}", e))?;

            let n = out_host.len().min(tgt_host.len());
            let mse: f64 = (0..n)
                .map(|i| {
                    let d = (out_host[i] - tgt_host[i]) as f64;
                    d * d
                })
                .sum::<f64>()
                / n.max(1) as f64;

            total_loss += mse;
        }

        Ok(total_loss / max_samples as f64)
    }

    /// Compute accuracy on validation set
    ///
    /// Accuracy is defined as percentage of predictions within 10% of actual value.
    fn compute_accuracy(
        &self,
        mamba: &mut Mamba2SSM,
        val_data: &[(GpuTensor, GpuTensor)],
    ) -> Result<f64> {
        let mut correct = 0;
        let mut total = 0;
        let max_samples = 100.min(val_data.len()); // Limit for speed

        let ml_dev = MlDevice::cuda(0)
            .map_err(|e| anyhow::anyhow!("CUDA device required: {}", e))?;
        let _stream = ml_dev.cuda_stream()
            .map_err(|e| anyhow::anyhow!("No CUDA stream: {}", e))?;

        for (input, target) in val_data.iter().take(max_samples) {
            // Forward pass
            let output = mamba
                .forward(input)
                .map_err(|e| anyhow::anyhow!("Forward pass failed: {}", e))?;

            // Extract scalar values via CPU readback (validation path)
            let out_host = output.to_vec() // gpu-exit: validation readback
                .map_err(|e| anyhow::anyhow!("Output readback failed: {}", e))?;
            let tgt_host = target.to_vec() // gpu-exit: validation readback
                .map_err(|e| anyhow::anyhow!("Target readback failed: {}", e))?;

            let pred = out_host.first().copied().unwrap_or(0.0);
            let actual = tgt_host.first().copied().unwrap_or(1.0);

            // Check if within 10% error
            let error = ((pred - actual) / actual).abs();
            if error < 0.1 {
                correct += 1;
            }
            total += 1;
        }

        Ok(correct as f64 / total as f64)
    }
}

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

    #[tokio::test]

    async fn test_mamba2_benchmark_runner_creation() {
        let gpu_manager = Arc::new(GpuHardwareManager::new().expect("GPU manager creation"));
        let _runner = Mamba2BenchmarkRunner::new(gpu_manager);
        // Runner created successfully
    }

    #[tokio::test]
    async fn test_mamba2_config_creation() {
        let config = Mamba2BenchmarkRunner::create_mamba_config(32, 8);
        assert_eq!(config.d_model, 32);
        assert_eq!(config.d_state, 16);
        assert_eq!(config.batch_size, 8);
        assert_eq!(config.seq_len, 128);
        assert!(config.use_ssd);
        assert!(config.use_selective_state);
        assert!(config.hardware_aware);
    }

    #[tokio::test]
    #[ignore] // Full benchmark — runs 2+ epochs of Mamba2 training, too slow for CI
    async fn test_full_mamba2_benchmark() {
        let gpu_manager = Arc::new(GpuHardwareManager::new().expect("GPU manager creation"));
        let mut runner = Mamba2BenchmarkRunner::new(gpu_manager);

        // Run short benchmark (2 epochs for testing)
        let result = runner.run_benchmark(2).await.expect("Benchmark execution");

        // Validate results
        assert_eq!(result.total_epochs, 2);
        assert_eq!(result.model_name, "MAMBA-2-SSM");
        assert!(result.memory_peak_mb > 0.0);
        assert!(result.statistics.mean_seconds > 0.0);
        assert!(result.final_accuracy >= 0.0 && result.final_accuracy <= 1.0);
    }
}
