//! # DQN Benchmark Runner - Production GPU Training Benchmarks
//!
//! This module implements comprehensive benchmarking for the DQN model
//! using real market data from DBN files. It measures actual GPU training performance,
//! memory usage, stability, and convergence metrics.
//!
//! ## Architecture
//!
//! ```text
//! ┌──────────────┐     ┌──────────────┐     ┌──────────────┐
//! │  DBN Files   │────▶│  DQN Trainer │────▶│  Benchmarks  │
//! │  (Real Data) │     │  (GPU/CPU)   │     │  (Stats)     │
//! └──────────────┘     └──────────────┘     └──────────────┘
//!        │                     │                     │
//!        ▼                     ▼                     ▼
//!   OHLCV Bars          Training Loop         Memory Profile
//!   6E.FUT Data         Experience Replay     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
//!
//! ## Usage
//!
//! ```rust,no_run
//! use ml::benchmark::dqn_benchmark::DqnBenchmarkRunner;
//! 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 = DqnBenchmarkRunner::new(gpu_manager);
//!
//! let result = runner.run_benchmark(10).await?;
//! println!("DQN 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 std::sync::Arc;
use std::time::Instant;
use tokio::sync::Mutex;
use tracing::info;

use crate::dqn::{Experience, DQN, DQNConfig};
use crate::data_loader::{FeatureMatrix, 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};

/// DQN benchmark result containing all metrics
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DqnBenchmarkResult {
    /// 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>,
    /// Average loss across all epochs
    pub avg_loss: f64,
}

/// DQN Benchmark Runner
///
/// Runs comprehensive benchmarks on DQN using real market data.
/// Measures GPU performance, memory usage, training stability, and convergence.
#[derive(Debug)]
pub struct DqnBenchmarkRunner {
    /// 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 DqnBenchmarkRunner {
    /// Create new DQN benchmark runner
    ///
    /// # Arguments
    ///
    /// * `gpu_manager` - GPU hardware manager for device management
    ///
    /// # Returns
    ///
    /// New DqnBenchmarkRunner 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 DQN benchmark
    ///
    /// # Arguments
    ///
    /// * `epochs` - Number of training epochs to run
    ///
    /// # Returns
    ///
    /// DqnBenchmarkResult with all metrics
    pub async fn run_benchmark(&mut self, epochs: usize) -> Result<DqnBenchmarkResult> {
        info!("Starting DQN benchmark with {} epochs", epochs);

        // Step 1: Load real market data from DBN files
        let (state_data, _indicators) = self.load_market_data().await?;
        let state_dim = state_data.first().map(|s| s.len()).unwrap_or(32);
        info!(
            "Loaded {} market data samples, state_dim={}",
            state_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 DQN model
        let mut dqn = self.create_dqn_model(state_dim)?;
        info!(
            "Created DQN model on device: {:?}",
            self.gpu_manager.device()
        );

        // Step 5: Populate experience replay buffer with real data
        self.populate_replay_buffer(&mut dqn, &state_data)?;
        info!(
            "Populated replay buffer with {} experiences",
            dqn.get_replay_buffer_size()?
        );

        // Step 6: Training loop with comprehensive metrics
        let mut training_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 result = dqn.train_step(None)?;
            let device = ml_core::device::MlDevice::cuda_if_available(0);
            let stream = device.cuda_stream()
                .map_err(|e| anyhow::anyhow!("No CUDA stream: {e}"))?;
            let loss = result.loss_gpu.to_scalar(stream)
                .map_err(|e| anyhow::anyhow!("Loss readback: {e}"))?;
            let epoch_time_s = epoch_start.elapsed().as_secs_f64();

            // Record metrics
            self.statistical_sampler.add_sample(epoch_time_s);
            training_losses.push(loss as f64);
            self.stability_validator.record_loss(loss as f64);
            self.stability_validator
                .record_gradient_norm(loss.abs() as f64); // 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 {}/{}: loss={:.6}, time={:.4}s",
                    epoch + 1,
                    epochs,
                    loss,
                    epoch_time_s
                );
            }
        }

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

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

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

        // Step 10: Compute average loss
        let avg_loss = training_losses.iter().sum::<f64>() / training_losses.len() as f64;

        info!("DQN 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 loss: {:.6}", avg_loss);
        info!("  Training stable: {}", stability.is_stable);

        Ok(DqnBenchmarkResult {
            model_name: "DQN".to_owned(),
            total_epochs: epochs,
            statistics,
            memory_peak_mb,
            stability,
            batch_config,
            training_losses,
            avg_loss,
        })
    }

    /// Load real market data from DBN files
    ///
    /// Loads OHLCV data and extracts features for DQN training.
    async fn load_market_data(&self) -> Result<(Vec<Vec<f32>>, Vec<Vec<f32>>)> {
        // 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")?;

        // Calculate technical indicators
        let indicators = loader
            .calculate_indicators(&bars)
            .context("Failed to calculate indicators")?;

        // Convert FeatureMatrix to state vectors
        let state_data = self.convert_features_to_states(&features)?;

        // Convert indicators to feature vectors
        let indicator_features = self.convert_indicators_to_features(&indicators)?;

        Ok((state_data, indicator_features))
    }

    /// Convert FeatureMatrix to DQN state vectors
    fn convert_features_to_states(&self, features: &FeatureMatrix) -> Result<Vec<Vec<f32>>> {
        let mut states = Vec::new();

        for i in 0..features.prices.len() {
            let mut state = Vec::new();

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

            // Add returns (1 value)
            if i < features.returns.len() {
                state.push(features.returns[i]);
            }

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

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

            states.push(state);
        }

        Ok(states)
    }

    /// Convert Indicators to feature vectors
    fn convert_indicators_to_features(
        &self,
        indicators: &crate::data_loader::Indicators,
    ) -> Result<Vec<Vec<f32>>> {
        let mut features = Vec::new();
        let len = indicators.rsi.len();

        for i in 0..len {
            let mut feature_vec = Vec::new();

            // Add all 10 technical indicators
            if i < indicators.rsi.len() {
                feature_vec.push(indicators.rsi[i]);
            }
            if i < indicators.macd.len() {
                feature_vec.push(indicators.macd[i]);
            }
            if i < indicators.macd_signal.len() {
                feature_vec.push(indicators.macd_signal[i]);
            }
            if i < indicators.bb_upper.len() {
                feature_vec.push(indicators.bb_upper[i]);
            }
            if i < indicators.bb_middle.len() {
                feature_vec.push(indicators.bb_middle[i]);
            }
            if i < indicators.bb_lower.len() {
                feature_vec.push(indicators.bb_lower[i]);
            }
            if i < indicators.atr.len() {
                feature_vec.push(indicators.atr[i]);
            }
            if i < indicators.ema_fast.len() {
                feature_vec.push(indicators.ema_fast[i]);
            }
            if i < indicators.ema_slow.len() {
                feature_vec.push(indicators.ema_slow[i]);
            }
            if i < indicators.volume_ma.len() {
                feature_vec.push(indicators.volume_ma[i]);
            }

            features.push(feature_vec);
        }

        Ok(features)
    }

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

        // Test function: Try creating a DQN model with different batch sizes
        let test_fn = |batch_size: usize| {
            // Simple heuristic: accept reasonable batch sizes for DQN
            if batch_size > 2048 {
                Ok(false) // Too large for DQN
            } else if batch_size < 8 {
                Ok(false) // Too small for effective training
            } else {
                Ok(true) // Acceptable range
            }
        };

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

    /// Create DQN model with specified state dimension
    fn create_dqn_model(&self, state_dim: usize) -> Result<DQN> {
        let config = Self::create_dqn_config(state_dim);
        DQN::new(config).context("Failed to create DQN model")
    }

    /// Create DQN configuration
    fn create_dqn_config(state_dim: usize) -> DQNConfig {
        DQNConfig {
            state_dim,
            num_actions: 3, // Buy, Sell, Hold
            hidden_dims: vec![256, 128, 64],
            learning_rate: 1e-4,
            gamma: 0.99,
            epsilon_start: 0.1,
            epsilon_end: 0.01,
            epsilon_decay: 0.995,
            replay_buffer_capacity: 10_000,
            batch_size: 32,
            min_replay_size: 100,
            target_update_freq: 100,
            huber_delta: 1.0,     // Standard Huber delta
            leaky_relu_alpha: 0.01, // Standard LeakyReLU alpha to prevent dead neurons
            gradient_clip_norm: 10.0, // Standard gradient clipping (Wave 11 Bug #1 fix)

            // WAVE 16 (Agent 36): Target update configuration (switched to Hard updates)
            tau: 1.0,
            tau_final: 1.0,          // No annealing for hard updates
            tau_anneal_steps: 0,
            use_soft_updates: false,

            // Rainbow DQN warmup period (disabled for benchmarking)
            warmup_steps: 0, // No warmup for fast benchmarks

            // Bug #33 fix: Initial capital for Q-value normalization
            initial_capital: 100_000.0, // $100k default for benchmarks

            // Wave 3: Prioritized Experience Replay (PER) - GPU PER mandatory on CUDA
            per_alpha: 0.6,
            per_beta_start: 0.4,
            per_beta_max: 1.0,
            per_beta_annealing_steps: 100000,

            // Wave 2.1: Dueling Networks - disabled for benchmarks
            dueling_hidden_dim: 64,

            // Wave 2.2: Multi-Step Returns - standard TD(0) for benchmarks
            n_steps: 1,

            // Wave 2.3: Distributional RL (always enabled)
            num_atoms: 51,
            v_min: -1000.0,
            v_max: 1000.0,

            // Wave 2.4: Noisy Networks (always enabled)
            noisy_sigma_init: 0.5,

            // BUG #37 FIX: Q-value clipping (prevents step-level explosions)
            enable_q_value_clipping: true,
            q_value_clip_min: -500.0,
            q_value_clip_max: 500.0,

            // WAVE 23 P0 Fix #1: Adaptive gradient collapse threshold
            gradient_collapse_multiplier: 100.0,
            gradient_collapse_patience: 5,

            cql_alpha: 1.0,
            use_iqn: false,
            iqn_num_quantiles: 64,
            iqn_kappa: 1.0,
            iqn_embedding_dim: 64,
            use_cvar_action_selection: false,
            cvar_alpha: 0.05,

            minimum_profit_factor: 1.5,
            weight_decay: 1e-4,
            ..Default::default()
        }
    }

    /// Populate replay buffer with real market data
    fn populate_replay_buffer(&self, dqn: &mut DQN, state_data: &[Vec<f32>]) -> Result<()> {
        let min_samples = 1000; // Minimum samples for meaningful training
        let samples_to_add = min_samples.min(state_data.len().saturating_sub(1));

        info!("Populating replay buffer with {} samples", samples_to_add);

        for i in 0..samples_to_add {
            // Current state
            let state = state_data[i].clone();

            // Random action (exploration phase)
            let action = (i % 3) as u8; // Cycle through Buy(0), Sell(1), Hold(2)

            // Simple reward: positive if price increased, negative otherwise
            let reward = if i + 1 < state_data.len() {
                let current_close = state.get(3).copied().unwrap_or(0.0); // Close price at index 3
                let next_close = state_data[i + 1].get(3).copied().unwrap_or(0.0);
                (next_close - current_close) * 100.0 // Scale reward
            } else {
                0.0
            };

            // Next state
            let next_state = if i + 1 < state_data.len() {
                state_data[i + 1].clone()
            } else {
                state.clone()
            };

            // Terminal if last sample
            let done = i + 1 >= state_data.len();

            // Create experience
            let experience = Experience::new(state, action, reward, next_state, done);

            // Store in replay buffer
            dqn.store_experience(experience)?;
        }

        Ok(())
    }
}

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

    #[tokio::test]

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

    #[tokio::test]
    async fn test_dqn_config_creation() {
        let config = DqnBenchmarkRunner::create_dqn_config(32);
        assert_eq!(config.state_dim, 32);
        assert_eq!(config.num_actions, 3);
        assert_eq!(config.learning_rate, 1e-4);
    }

    #[tokio::test]
    async fn test_training_dtype_bf16_on_cuda() {
        // BF16 is unconditional on CUDA — no config needed
        let cuda_dtype = ml_core::native_types::NativeDType::BF16;
        // On CPU runners this returns F32, on CUDA it returns BF16
        assert!(
            cuda_dtype == ml_core::native_types::NativeDType::BF16
                || cuda_dtype == ml_core::native_types::NativeDType::F32,
            "training_dtype must return BF16 (CUDA) or F32 (CPU)"
        );
    }

    #[tokio::test]

    async fn test_feature_conversion() {
        let gpu_manager = Arc::new(GpuHardwareManager::new().expect("GPU manager creation"));
        let runner = DqnBenchmarkRunner::new(gpu_manager);

        // Create mock feature matrix
        let features = FeatureMatrix {
            prices: vec![vec![1.0, 2.0, 3.0, 4.0, 5.0]],
            returns: vec![0.1],
            volume: vec![1000.0],
            indicators: vec![vec![50.0; 10]], // 10 indicators
        };

        let states = runner.convert_features_to_states(&features).unwrap();
        assert_eq!(states.len(), 1);
        assert_eq!(states[0].len(), 5 + 1 + 1 + 10); // OHLCV + return + volume + 10 indicators
    }

    #[tokio::test]
    #[ignore] // Full benchmark loads real data — too slow for CI
    async fn test_full_dqn_benchmark() {
        let gpu_manager = Arc::new(GpuHardwareManager::new().expect("GPU manager creation"));
        let mut runner = DqnBenchmarkRunner::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, "DQN");
        assert!(result.memory_peak_mb > 0.0);
        assert!(result.statistics.mean_seconds > 0.0);
    }
}
