#![allow(unsafe_code)] // Required for safetensors f32-to-u8 slice serialization
//! DQN Trainer Implementation
//!
//! Main training loop and execution logic for Deep Q-Network.

use std::collections::VecDeque;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use ml_core::device::MlDevice;
use ml_core::cuda_autograd::GpuTensor;
use cudarc::driver::CudaStream;
use crate::cuda_pipeline::DqnGpuData;
use crate::cuda_pipeline::gpu_experience_collector::TradeStats;
use ml_core::fill_simulator::FillSimulator;
use risk::drawdown_monitor::DrawdownMonitor;
use risk::safety::position_limiter::HybridPositionLimiter;
use tokio::sync::RwLock;
use tracing::info;
use uuid::Uuid;

use crate::dqn::action_space::FactoredAction;
use crate::dqn::circuit_breaker::CircuitBreaker;
use crate::dqn::curiosity::CuriosityModule;
use crate::dqn::logging::{LoggingConfig, MetricsAggregator};
use crate::dqn::portfolio_tracker::PortfolioTracker;
use crate::dqn::reward::RewardFunction;
use crate::TrainingMetrics;
use crate::features::extraction::FeatureVector;
use crate::features::microstructure_features::*;
use crate::labeling::triple_barrier::TripleBarrierEngine;

use super::config::{DQNAgentType, DQNHyperparameters};
use super::statistics::{FeatureStatistics, QValueStats};

mod action;
mod metrics;
mod state;
mod constructor;
mod training_loop;
mod train_step;

#[cfg(test)]
mod tests;



pub struct DQNTrainer {
    /// DQN agent
    pub(crate) agent: Arc<RwLock<DQNAgentType>>,
    /// Training hyperparameters
    pub(crate) hyperparams: DQNHyperparameters,
    /// MlDevice (GPU or CPU)
    pub(crate) device: MlDevice,
    /// Single forked CudaStream shared by ALL GPU components (experience collector,
    /// fused trainer, portfolio sim, monitoring). Forked once in constructor —
    /// eliminates dual-stream event tracking conflicts and Candle default stream
    /// usage in the hot path. None on CPU devices.
    pub(crate) cuda_stream: Option<Arc<CudaStream>>,
    /// Training metrics
    pub(crate) metrics: Arc<RwLock<TrainingMetrics>>,
    /// Loss history for plateau detection
    pub(crate) loss_history: Vec<f64>,
    /// Q-value history for floor detection
    pub(crate) q_value_history: Vec<f64>,
    /// Best validation loss achieved so far
    pub(crate) best_val_loss: f64,
    /// Validation data for computing validation loss
    pub(crate) val_data: Vec<(FeatureVector, Vec<f64>)>,
    /// Validation loss history for early stopping
    pub(crate) val_loss_history: Vec<f64>,
    /// C4: Sharpe history for Sharpe-based early stopping
    pub(crate) sharpe_history: Vec<f64>,
    /// C4: Best Sharpe ratio achieved so far (higher = better)
    pub(crate) best_sharpe: f64,
    /// Epoch with best validation loss
    pub(crate) best_epoch: usize,
    /// Step counter for gradient logging (logs every 10 steps)
    pub(crate) gradient_logging_step: usize,
    /// Original buffer_size before AutoReplaySizer (for gradient collapse warmup)
    pub(crate) collapse_warmup_buffer_size: usize,
    /// Portfolio state tracker for P&L-based rewards (Bug #2 fix)
    pub portfolio_tracker: PortfolioTracker,
    /// Feature normalization statistics (WAVE 3 FIX #2)
    /// None during stats collection phase (epochs 0-10), Some during normalization phase (epochs 11+)
    pub feature_stats: Option<FeatureStatistics>,
    /// Sliding window of recent actions for reward calculation (max 100)
    pub(crate) recent_actions: VecDeque<FactoredAction>,
    /// Reward function for calculating rewards with recent actions
    pub(crate) reward_fn: RewardFunction,

    // WAVE 16S: Adaptive Risk Management Components
    /// Kelly criterion optimizer for position sizing (None if disabled)
    pub(crate) kelly_optimizer: Option<Arc<crate::risk::kelly_optimizer::KellyCriterionOptimizer>>,
    /// Trade history for Kelly calculation (wins/losses)
    pub(crate) trade_history: VecDeque<f64>,
    /// Volatility tracker for epsilon adjustment (None if disabled)
    pub(crate) volatility_returns: VecDeque<f64>,
    /// Per-epoch trade statistics from GPU portfolio_states reduction.
    /// Replaces pnl_history — real trade metrics instead of averaged rewards.
    pub(crate) trade_stats_history: Vec<TradeStats>,

    // Wave 16 Portfolio Features
    /// Enable action masking (filters invalid actions before Q-value computation)
    pub enable_action_masking: bool,
    /// Maximum position size for action masking (default: 2.0)
    pub max_position: f64,
    /// Current epoch (for Q-gap warm-up ramp)
    pub(crate) current_epoch: usize,
    /// Entropy regularizer for preventing policy collapse (None if disabled)
    // entropy_regularizer removed — SAC-style entropy is computed directly on Q-value tensors in DQN::compute_loss_internal
    /// Multi-asset portfolio tracker (None if single-asset mode)
    pub multi_asset_portfolio: Option<Arc<crate::dqn::multi_asset::MultiAssetPortfolioTracker>>,
    /// Stress tester for robustness validation (None if disabled)
    pub stress_tester: Option<Arc<crate::dqn::stress_testing::DQNStressTester>>,

    // Wave 16 Core Risk Features Integration
    /// Drawdown monitor for tracking portfolio drawdowns (15% max drawdown)
    pub drawdown_monitor: Option<Arc<DrawdownMonitor>>,
    /// Position limiter with 3-tier limits (±10.0 absolute, 1M notional, 10% concentration)
    pub position_limiter: Option<Arc<HybridPositionLimiter>>,
    /// Circuit breaker for stopping training on consecutive failures
    pub circuit_breaker: Option<Arc<CircuitBreaker>>,

    // WAVE 3.10: Microstructure Feature Calculators (12 features)
    pub(crate) micro_high_low_spread: HighLowSpread,
    pub(crate) micro_vw_spread: VolumeWeightedSpread,
    pub(crate) micro_tick_count: TickCount,
    pub(crate) micro_inter_arrival: InterArrivalTime,
    pub(crate) micro_buy_sell_imbalance: BuySellImbalance,
    pub(crate) micro_kyle_lambda: KyleLambda,
    pub(crate) micro_price_impact: PriceImpact,
    pub(crate) micro_variance_ratio: VarianceRatio,
    // Note: Roll Measure, Corwin-Schultz, Amihud, VPIN already exist in ml/src/microstructure/
    // We'll integrate those in the update logic
    /// Track last timestamp for inter-arrival time calculation
    pub(crate) last_timestamp_ns: u64,
    /// Track last close price for microstructure calculations
    pub(crate) last_close: f64,

    // WAVE 1.1: Triple Barrier Integration
    /// Triple barrier engine for position exit labeling
    pub(crate) triple_barrier: Arc<RwLock<TripleBarrierEngine>>,
    /// Active position tracker ID (None = no active position)
    pub(crate) active_position_tracker: Option<Uuid>,
    /// WAVE P3: Track previous simulated position for barrier tracking continuity
    pub(crate) previous_simulated_position: f32,

    // WAVE 1.2: Safety Infrastructure Integration (8 Systems)
    /// Loss history window for spike detection (size: 30)
    pub(crate) safety_loss_history: VecDeque<f32>,
    /// Loss plateau counter for anomaly detection
    pub(crate) safety_loss_plateau_counter: usize,
    /// Action counts for diversity monitoring (5 exposure actions)
    pub(crate) safety_action_counts: std::collections::HashMap<usize, usize>,
    /// Memory manager for GPU OOM risk monitoring
    pub(crate) safety_memory_manager: Arc<RwLock<crate::safety::memory_manager::SafeMemoryManager>>,
    /// Safety enforcement level (Strict/Normal/Permissive)
    pub(crate) safety_level: crate::safety::SafetyLevel,
    /// Step counter for periodic safety checks
    pub(crate) safety_step_counter: usize,

    /// Optional path to feature cache directory for faster hyperopt
    pub(crate) feature_cache_dir: Option<PathBuf>,

    /// Previous epoch's mean Q-value for overestimation detection
    pub(crate) prev_epoch_q_mean: f64,
    /// Current effective tau (may be temporarily increased if Q-values drift)
    pub(crate) adaptive_tau: f64,

    /// WAVE 24 (Agent 17): Patience-based early stopping for anti-overfitting
    pub(crate) early_stopping: super::early_stopping::EarlyStopping,

    /// WAVE 26 P0.6: Learning rate scheduler with warmup
    pub(crate) lr_scheduler: super::lr_scheduler::LRScheduler,

    /// WAVE 26 P1.8: Curiosity module for intrinsic rewards (None if curiosity_weight = 0.0)
    pub(crate) curiosity_module: Option<CuriosityModule>,

    // WAVE 26 P1: Advanced DQN Features Integration
    // P1.3: Sharpe Ratio Reward Component
    /// Rolling buffer of returns for Sharpe ratio calculation (max: sharpe_window)
    pub(crate) returns_history: VecDeque<f64>,
    /// Sharpe reward weight (0.0 = disabled)
    pub(crate) sharpe_weight: f64,

    // P1.6: Adaptive Dropout Scheduling
    /// Optional dropout scheduler (None if disabled)
    pub(crate) dropout_scheduler: Option<crate::dqn::network::DropoutScheduler>,

    // P1.7: Hindsight Experience Replay (HER)
    /// Optional HER buffer (None if her_ratio = 0.0)
    pub(crate) her_buffer: Option<Arc<crate::dqn::hindsight_replay::HindsightReplayBuffer>>,

    // P1.9: Generalized Advantage Estimation (GAE)
    /// Optional GAE calculator (None if disabled)
    pub(crate) gae_calculator: Option<crate::dqn::gae::GAECalculator>,

    // P1.11: Noisy Network Sigma Scheduling
    /// Optional noisy sigma scheduler (None if disabled)
    pub(crate) noisy_sigma_scheduler: Option<crate::dqn::noisy_sigma_scheduler::NoisySigmaScheduler>,

    // WAVE 30: Structured Logging Integration
    /// Logging configuration for training metrics
    pub(crate) logging_config: LoggingConfig,
    /// Metrics aggregator for windowed training statistics
    pub(crate) metrics_aggregator: MetricsAggregator,

    /// Current effective batch size (may be reduced by OOM recovery)
    pub(crate) current_batch_size: usize,

    /// Cached Q-value estimate for periodic monitoring (avoids extra forward pass every step)
    pub(crate) cached_avg_q: f64,
    /// Counter for Q-value estimation frequency (estimate every N training steps)
    pub(crate) q_estimation_counter: u64,

    /// Pre-uploaded GPU training data (set once, reused across epochs)
    pub(crate) gpu_data: Option<DqnGpuData>,

    /// Raw cudarc targets buffer for CUDA kernel (parallel to raw features buffer)
    pub(crate) targets_raw_cuda: Option<cudarc::driver::CudaSlice<f32>>,

    /// Raw cudarc features buffer for CUDA experience kernel [num_bars * 42]
    pub(crate) features_raw_cuda: Option<cudarc::driver::CudaSlice<f32>>,

    /// GPU experience collector for zero-roundtrip CUDA kernel (Phase 2b)
    pub(crate) gpu_experience_collector: Option<crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector>,

    /// GPU-fused epsilon-greedy action selector (eliminates argmax GPU->CPU sync barrier)
    pub(crate) gpu_action_selector: Option<crate::cuda_pipeline::gpu_action_selector::GpuActionSelector>,

    /// GPU training guard for zero-sync safety checks (loss clip, NaN, grad collapse)
    pub(crate) training_guard: Option<crate::cuda_pipeline::gpu_training_guard::GpuTrainingGuard>,

    /// GPU monitoring reducer — accumulates reward/action stats across kernel launches
    pub(crate) gpu_monitoring: Option<crate::cuda_pipeline::gpu_monitoring::GpuMonitoringReducer>,

    /// GPU reduction kernels for fused stats/argmax/sum (lazy-init, reused across epochs)
    pub(crate) reduction_kernels: Option<ml_core::cuda_autograd::ReductionKernels>,

    /// Reusable GPU staging buffers for zero-alloc fold transitions
    pub(crate) buffer_pool: Option<crate::cuda_pipeline::GpuBufferPool>,

    /// Double-buffered GPU data for zero-downtime fold transitions
    pub(crate) double_buffer: Option<crate::cuda_pipeline::double_buffer::DoubleBufferedLoader>,

    /// GPU-resident walk-forward data (entire dataset on GPU, per-fold views via index ranges)
    pub(crate) gpu_walk_forward: Option<crate::cuda_pipeline::gpu_walk_forward::GpuWalkForwardData>,

    /// Multi-GPU configuration for data-parallel training (None = single GPU)
    pub(crate) multi_gpu: Option<crate::cuda_pipeline::multi_gpu::MultiGpuConfig>,

    /// Cached GPU n_episodes (computed once from nvidia-smi, reused across epochs)
    /// Avoids forking nvidia-smi subprocess every epoch (~5-10ms per fork).
    pub(crate) cached_n_episodes: Option<i32>,

    /// Pre-computed OFI features per bar (indexed by global bar position).
    /// Populated during data loading when MBP-10 order book data is available.
    /// Passed as `regime_features` in `TradingState::from_normalized()`.
    /// Arc-shared to avoid 2.68 GB copy per hyperopt trial (41.9M × 8 × 8 bytes).
    pub(crate) ofi_features: Option<Arc<[[f64; 8]]>>,
    /// Number of training bars (OFI offset for validation data).
    /// val_data[i] corresponds to ofi_features[ofi_val_offset + i].
    pub(crate) ofi_val_offset: usize,

    /// CPU-cached validation features flat vec [sample_size * 42] — built once, reused per epoch
    pub(crate) val_features_gpu: Option<Vec<f32>>,
    /// CPU-cached validation close prices (current, next) — built once per training run
    pub(crate) val_closes_gpu: Option<(Vec<f32>, Vec<f32>)>,
    /// CPU-cached validation OFI features flat vec [sample_size * 8] — built once
    pub(crate) val_ofi_gpu: Option<Vec<f32>>,

    // Phase C: Fill simulation and smart order routing
    /// Fill simulator for order type-dependent execution modeling
    pub(crate) fill_simulator: FillSimulator,
    /// EMA of bar volatility (|close log return|) for OrderRouter routing decisions
    pub(crate) vol_ema: f64,
    /// Running median volatility estimate (slowly adapting EMA)
    pub(crate) median_vol: f64,

    /// Fused CUDA training context: pre-allocated batch buffers + single entry point
    /// for CUDA Graph capture. Only active for Standard DQN on CUDA devices.
    /// Lazy-initialized on first training step; dropped and recreated if batch_size changes.
    pub(crate) fused_ctx: Option<super::fused_training::FusedTrainingCtx>,
}

impl std::fmt::Debug for DQNTrainer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DQNTrainer")
            .field("hyperparams", &self.hyperparams)
            .finish_non_exhaustive()
    }
}

impl Drop for DQNTrainer {
    fn drop(&mut self) {
        // Drop GPU resources BEFORE the CudaStream and CudaContext.
        // Default field drop order follows declaration order, but GPU resources
        // (fused_ctx at line 285, gpu_experience_collector at line 228) are declared
        // AFTER cuda_stream (line 60) and device (line 55). Their Drop impls call
        // cuStreamSynchronize on the forked stream — which would be a use-after-free
        // if the stream dropped first.
        self.fused_ctx = None;
        self.gpu_experience_collector = None;
        // Now sync and drop the forked stream before the CudaContext drops.
        if let Some(ref stream) = self.cuda_stream {
            #[allow(unsafe_code)]
            unsafe {
                cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream());
            }
        }
        self.cuda_stream = None;
    }
}

impl DQNTrainer {
    /// Create new DQN trainer with hyperparameters and debug logging disabled
    pub fn new(hyperparams: DQNHyperparameters) -> Result<Self> {
        Self::new_with_debug(hyperparams, false)
    }

    /// Create new DQN trainer with a specific compute device.
    /// Used by hyperopt to share a single CUDA context across parallel trials.
    pub fn new_with_device(hyperparams: DQNHyperparameters, device: MlDevice) -> Result<Self> {
        Self::new_internal(hyperparams, false, Some(device))
    }

    /// Create new DQN trainer with hyperparameters and configurable debug logging
    ///
    /// # Arguments
    /// * `hyperparams` - DQN training hyperparameters
    /// * `debug_logging` - Enable debug logging (REWARD_DEBUG, gradient norms, etc.)
    pub fn new_with_debug(hyperparams: DQNHyperparameters, debug_logging: bool) -> Result<Self> {
        Self::new_internal(hyperparams, debug_logging, None)
    }


    /// Two-phase stress tester initialization.
    ///
    /// Call this after constructing `DQNTrainer` to resolve the circular dependency:
    /// `DQNStressTester::new()` requires a `DQNTrainer`, so the tester cannot be
    /// created *during* trainer construction.  This method builds a lightweight
    /// inner trainer (with stress testing itself disabled to avoid recursion) and
    /// hands it to `DQNStressTester::new()`.
    ///
    /// No-op if `hyperparams.enable_stress_testing` is `false`.
    pub fn init_stress_tester(&mut self) -> Result<()> {
        if !self.hyperparams.enable_stress_testing {
            return Ok(());
        }

        // Clone hyperparams with stress testing disabled so the inner trainer
        // does not recursively try to initialise its own stress tester.
        // Also disable GPU-heavy features that the stress tester doesn't need —
        // otherwise we double the VRAM usage (3 extra regime heads, 3 extra GPU PER buffers,
        // experience collector, etc.) which causes OOM on 4GB GPUs.
        let mut inner_hp = self.hyperparams.clone();
        inner_hp.enable_stress_testing = false;
        inner_hp.enable_regime_qnetwork = false;
        inner_hp.buffer_size = 1;

        let inner_trainer = Self::new_with_device(inner_hp, self.device.clone())
            .context("Failed to create inner DQNTrainer for stress tester")?;

        let tester = crate::dqn::stress_testing::DQNStressTester::new(inner_trainer)?;
        self.stress_tester = Some(Arc::new(tester));
        info!("Stress testing enabled (8 scenarios)");

        Ok(())
    }

    /// Get a reference to the double-buffered loader, if GPU is active.
    pub fn double_buffer(&self) -> Option<&crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> {
        self.double_buffer.as_ref()
    }

    /// Get a mutable reference to the double-buffered loader, if GPU is active.
    pub fn double_buffer_mut(&mut self) -> Option<&mut crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> {
        self.double_buffer.as_mut()
    }

    /// Set feature cache directory for faster hyperopt
    ///
    /// Enables loading pre-computed features from disk instead of recomputing them
    pub fn with_feature_cache(mut self, cache_dir: PathBuf) -> Self {
        self.feature_cache_dir = Some(cache_dir);
        self
    }

    /// Train DQN on market data from DBN files
    ///
    /// # Arguments
    ///
    /// * `dbn_data_dir` - Directory containing DBN files (e.g., "test_data/real/databento/ml_training/")
    /// * `checkpoint_callback` - Callback for saving checkpoints (epoch, model_data, is_final) -> `Result<String>`
    ///
    /// # Returns
    ///
    /// Training metrics (loss, accuracy, gradient norms, Q-values)
    pub async fn train<F>(
        &mut self,
        dbn_data_dir: &str,
        checkpoint_callback: F,
    ) -> Result<TrainingMetrics>
    where
        F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
    {
        info!(
            "Starting DQN training for {} epochs with batch size {}",
            self.hyperparams.epochs, self.hyperparams.batch_size
        );

        // Load market data from DBN files (ALL data for walk-forward or single-pass)
        let (training_data, val_data) = self.load_training_data(dbn_data_dir).await?;

        info!(
            "Loaded {} training samples, {} validation samples",
            training_data.len(),
            val_data.len()
        );

        // GPU walk-forward: upload ALL data to GPU, run expanding-window folds
        if self.hyperparams.enable_gpu_walk_forward && self.cuda_stream.is_some() {
            // Merge train+val into a single dataset for walk-forward splitting
            let mut all_data = training_data;
            all_data.extend(val_data);
            info!(
                "GPU walk-forward enabled: {} total bars, uploading to VRAM",
                all_data.len(),
            );
            return self.train_walk_forward(&all_data, checkpoint_callback).await;
        }

        // Standard single-pass training
        self.ofi_val_offset = training_data.len();
        self.val_data = val_data;
        self.val_features_gpu = None;
        self.val_closes_gpu = None;
        self.val_ofi_gpu = None;
        self.train_with_data_full_loop(&training_data, checkpoint_callback)
            .await
    }

    /// Train with preloaded data (skips disk I/O and feature extraction).
    ///
    /// Accepts pre-split training and validation data that was loaded once and
    /// cached across hyperopt trials. This avoids re-reading 36 `.dbn.zst` files
    /// and re-extracting 42 features on every trial, eliminating minutes of GPU
    /// idle time at each trial boundary.
    ///
    /// # Arguments
    ///
    /// * `training_data` - Pre-extracted (features, targets) for training split
    /// * `val_data` - Pre-extracted (features, targets) for validation split
    /// * `checkpoint_callback` - Checkpoint save callback
    ///
    /// # Returns
    ///
    /// Training metrics from the completed run
    ///
    /// # Errors
    ///
    /// Returns error if the training loop fails
    pub async fn train_with_preloaded_data<F>(
        &mut self,
        training_data: Vec<(FeatureVector, Vec<f64>)>,
        val_data: Vec<(FeatureVector, Vec<f64>)>,
        checkpoint_callback: F,
    ) -> Result<TrainingMetrics>
    where
        F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
    {
        // Clear stale CUDA context errors from previous training runs.
        // cudarc stores errors from CUDA Graph capture (cuStreamWaitEvent on
        // disabled events). These persist across DQNTrainer instances and
        // cause bind_to_thread() to fail on the next run.
        if let MlDevice::Cuda { ref context, .. } = self.device {
            let _ = context.check_err();
        }

        info!(
            "Starting DQN training with preloaded data: {} train, {} val samples",
            training_data.len(),
            val_data.len()
        );

        // Store validation data for loss computation
        self.ofi_val_offset = training_data.len();
        self.val_data = val_data;
        self.val_features_gpu = None;
        self.val_closes_gpu = None;
        self.val_ofi_gpu = None;

        // Use the common training loop (Wave 12 Group 3 refactor)
        self.train_with_data_full_loop(&training_data, checkpoint_callback)
            .await
    }

    /// Train with shared preloaded data (zero-copy for hyperopt).
    ///
    /// Same as [`train_with_preloaded_data`] but accepts `Arc`-wrapped data,
    /// avoiding a ~150 MB deep clone per hyperopt trial.
    pub async fn train_with_shared_data<F>(
        &mut self,
        training_data: &[(FeatureVector, Vec<f64>)],
        val_data: Vec<(FeatureVector, Vec<f64>)>,
        checkpoint_callback: F,
    ) -> Result<TrainingMetrics>
    where
        F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
    {
        info!(
            "Starting DQN training with shared data: {} train, {} val samples",
            training_data.len(),
            val_data.len()
        );

        self.ofi_val_offset = training_data.len();
        self.val_data = val_data;
        self.val_features_gpu = None;
        self.val_closes_gpu = None;
        self.val_ofi_gpu = None;
        self.train_with_data_full_loop(training_data, checkpoint_callback)
            .await
    }

    /// Train with GPU-resident walk-forward cross-validation.
    ///
    /// Uploads the ENTIRE dataset to GPU VRAM once, then runs expanding-window
    /// walk-forward: each fold trains on [0..T], validates on [T..V], tests on
    /// [V..E]. Fold transitions are zero-copy (index range changes only).
    ///
    /// Returns the metrics from the LAST fold (most data, most representative).
    pub async fn train_walk_forward<F>(
        &mut self,
        training_data: &[(FeatureVector, Vec<f64>)],
        mut checkpoint_callback: F,
    ) -> Result<TrainingMetrics>
    where
        F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
    {
        use crate::cuda_pipeline::gpu_walk_forward::{GpuWalkForwardConfig, GpuWalkForwardData};

        let wf_config = GpuWalkForwardConfig {
            initial_train_fraction: self.hyperparams.wf_initial_train_fraction,
            val_fraction: self.hyperparams.wf_val_fraction,
            test_fraction: self.hyperparams.wf_test_fraction,
            step_fraction: self.hyperparams.wf_step_fraction,
        };

        // Upload ALL data to GPU once
        let wf_stream = self.cuda_stream.as_ref()
            .ok_or_else(|| anyhow::anyhow!("CUDA stream required for walk-forward upload"))?;
        let wf_data = GpuWalkForwardData::upload(
            training_data,
            self.ofi_features.as_deref(),
            &wf_config,
            wf_stream,
        ).map_err(|e| anyhow::anyhow!("GPU walk-forward upload: {e}"))?;

        let num_folds = wf_data.num_folds();
        if num_folds == 0 {
            return Err(anyhow::anyhow!(
                "Insufficient data for walk-forward: {} bars, need at least {} for one fold",
                training_data.len(),
                ((wf_config.initial_train_fraction + wf_config.val_fraction + wf_config.test_fraction) * training_data.len() as f64) as usize,
            ));
        }

        info!(
            "GPU walk-forward: {} folds, {:.1} MB VRAM, {} total bars",
            num_folds, wf_data.vram_bytes as f64 / 1_048_576.0, wf_data.total_bars,
        );

        // Store GPU walk-forward data and cudarc buffers for the experience collector
        self.features_raw_cuda = Some(wf_data.features);
        self.targets_raw_cuda = Some(wf_data.targets);

        let mut last_metrics = TrainingMetrics::new();

        for fold_idx in 0..num_folds {
            let fold = wf_data.folds.get(fold_idx).ok_or_else(|| {
                anyhow::anyhow!("Fold {fold_idx} out of range")
            })?;

            info!(
                "=== Walk-Forward Fold {}/{} === train: {} bars, val: {} bars, test: {} bars",
                fold_idx + 1, num_folds, fold.train_len(), fold.val_len(), fold.test_len(),
            );

            // Split training_data into fold's train and val slices (for CPU-side data)
            let fold_train = &training_data[fold.train_start..fold.train_end];
            let fold_val: Vec<(FeatureVector, Vec<f64>)> =
                training_data[fold.val_start..fold.val_end].to_vec(); // cpu-side fold split

            // Store validation data for this fold
            self.ofi_val_offset = fold.train_end;
            self.val_data = fold_val;

            // Reset training state for new fold
            self.gpu_data = None; // Force re-upload via DqnGpuData for the fold's range
            self.best_sharpe = f64::NEG_INFINITY;
            self.best_val_loss = f64::INFINITY;
            self.loss_history.clear();
            self.q_value_history.clear();
            self.val_loss_history.clear();
            self.sharpe_history.clear();

            // Run training loop on this fold's data
            last_metrics = self
                .train_with_data_full_loop(fold_train, &mut checkpoint_callback)
                .await?;

            info!(
                "Fold {}/{} complete: loss={:.6}, epochs={}",
                fold_idx + 1, num_folds,
                last_metrics.loss,
                last_metrics.epochs_trained,
            );
        }

        // Clean up GPU walk-forward buffers (features/targets already stored in self)
        self.gpu_walk_forward = None;

        Ok(last_metrics)
    }


    
    /// Calculate adaptive bounds with margin
    ///
    /// # Arguments
    ///
    /// * `stats` - Q-value statistics from Phase 1
    /// * `margin` - Safety margin as fraction (e.g., 0.3 = 30%)
    ///
    /// # Returns
    ///
    /// Tuple of (v_min, v_max) with safety margin applied
    fn calculate_adaptive_bounds(stats: &QValueStats, margin: f64) -> (f64, f64) {
        let range = stats.max - stats.min;
        let v_min = stats.min - range * margin;
        let v_max = stats.max + range * margin;
        // Cap at ±10,000 to prevent explosion
        (v_min.max(-10000.0), v_max.min(10000.0))
    }
    
    /// Reinitialize categorical distribution with new bounds
    async fn reinit_categorical_distribution(&mut self, v_min: f64, v_max: f64) -> Result<(), crate::MLError> {
        let mut agent = self.agent.write().await;
        match &mut *agent {
            DQNAgentType::Standard(agent) => agent.reinit_categorical_distribution(v_min, v_max)?,
            DQNAgentType::RegimeConditional(agent) => agent.reinit_categorical_distribution(v_min, v_max)?,
        }
        Ok(())
    }















    /// Calculate reward based on price movement
    ///
    /// # Arguments
    /// * `current_close` - Current bar's close price
    /// * `next_close` - Next bar's close price (target)
    ///
    /// # Returns
    /// Normalized reward in [-1.0, 1.0] based on price change
    fn calculate_reward(&self, current_close: f64, next_close: f64) -> f32 {
        let price_change = next_close - current_close;
        // Normalize by 10.0 for ES futures typical moves (±10 points)
        // Clamp to [-1.0, 1.0] to prevent extreme rewards
        (price_change / 10.0).clamp(-1.0, 1.0) as f32
    }

    /// Check if we can train (buffer has enough samples)
    async fn can_train(&self) -> Result<bool> {
        let agent = self.agent.read().await;
        Ok(agent.can_train())
    }




    /// Get current epsilon value
    async fn get_epsilon(&self) -> Result<f64> {
        let agent = self.agent.read().await;
        Ok(agent.get_epsilon() as f64)
    }

    /// Set epsilon value (used for deterministic evaluation)
    async fn set_epsilon(&self, epsilon: f64) -> Result<()> {
        let mut agent = self.agent.write().await;
        agent.set_epsilon(epsilon);
        Ok(())
    }

    /// Get best validation loss achieved during training
    ///
    /// Returns the lowest validation loss seen across all epochs.
    /// Used by hyperopt adapter to optimize for generalization.
    pub fn get_best_val_loss(&self) -> f64 {
        self.best_val_loss
    }

    /// Get epoch number where best validation loss was achieved
    ///
    /// Returns the 1-indexed epoch number with the best validation loss.
    pub fn get_best_epoch(&self) -> usize {
        self.best_epoch
    }

    /// Get validation data for backtest integration
    ///
    /// Returns a reference to the validation dataset for hyperopt backtest evaluation.
    /// Each entry contains a FeatureVector (42 market + 3 portfolio = 45 dims) and the corresponding target values.
    /// Used by hyperopt adapter to run backtests on unseen data after training.
    pub fn get_val_data(&self) -> &[(FeatureVector, Vec<f64>)] {
        &self.val_data
    }




    /// Update portfolio tracker to reflect current position from backtest engine.
    ///
    /// Called between chunks so the next chunk's portfolio features
    /// accurately reflect the current position (direction, value, exposure).
    pub fn set_portfolio_for_backtest(
        &mut self,
        position_size: f32,
        entry_price: f32,
        current_price: f32,
    ) {
        self.portfolio_tracker = PortfolioTracker::new(
            self.portfolio_tracker.initial_capital(),
            self.portfolio_tracker.spread(),
            0.0,
        );
        if position_size.abs() > f32::EPSILON {
            self.portfolio_tracker
                .set_position_direct(position_size, entry_price, current_price);
        }
    }

    /// Get the device used by this trainer
    pub fn device(&self) -> &MlDevice {
        &self.device
    }

    /// Get access to the DQN agent
    ///
    /// Returns a reference to the Arc<RwLock<DQNAgentType>> for checkpoint saving.
    /// Used by hyperopt adapter to save model weights after training.
    pub fn get_agent(&self) -> &Arc<RwLock<DQNAgentType>> {
        &self.agent
    }

    /// Get reference to training hyperparameters
    ///
    /// Returns a reference to the DQN hyperparameters used for this trainer.
    /// Used by tests to validate configuration.
    pub fn hyperparams(&self) -> &DQNHyperparameters {
        &self.hyperparams
    }

    /// Get current learning rate from scheduler
    ///
    /// Returns the current learning rate after applying warmup and decay.
    /// Used by tests and monitoring to track LR schedule.
    pub fn get_current_lr(&self) -> f64 {
        self.lr_scheduler.get_lr()
    }

    /// Serialize model to bytes with architecture metadata embedded in safetensors header.
    ///
    /// For RegimeConditional agents, serializes ALL 3 heads (trending, ranging,
    /// volatile) into a single safetensors file using prefixed tensor names
    /// (`trending__`, `ranging__`, `volatile__`). This ensures walk-forward
    /// checkpoint restore loads all heads, not just the trending head.
    pub async fn serialize_model(&self) -> Result<Vec<u8>> {
        let agent = self.agent.read().await;

        let tensors: std::collections::HashMap<String, GpuTensor> = match &*agent {
            crate::trainers::dqn::DQNAgentType::RegimeConditional(regime) => {
                let mut all_tensors = std::collections::HashMap::new();
                for (prefix, head_opt) in [
                    ("trending__", regime.get_trending_head()),
                    ("ranging__", regime.get_ranging_head()),
                    ("volatile__", regime.get_volatile_head()),
                ] {
                    let head = head_opt.ok_or_else(|| {
                        anyhow::anyhow!("Missing {} head for serialization", prefix)
                    })?;
                    let vars = head.get_q_network_vars();
                    let vars_data = vars.data();
                    let _ser_stream = vars.stream();
                    for (name, param) in vars_data.iter() {
                        if let Ok(tensor) = GpuTensor::new(param.data.clone(), param.shape.clone()) {
                            all_tensors.insert(
                                format!("{}{}", prefix, name),
                                tensor,
                            );
                        }
                    }
                }
                all_tensors
            }
            _ => {
                let vars = agent.get_q_network_vars();
                let vars_data = vars.data();
                vars_data
                    .iter()
                    .filter_map(|(name, param)| {
                        GpuTensor::new(param.data.clone(), param.shape.clone())
                            .ok()
                            .map(|t| (name.clone(), t))
                    })
                    .collect()
            }
        };

        // Embed architecture metadata in safetensors header
        let arch_metadata = Some(agent.checkpoint_metadata());
        let ser_stream = self.cuda_stream.as_ref()
            .ok_or_else(|| anyhow::anyhow!("CUDA stream required for serialization"))?;
        // Download tensors to host for safetensors serialization
        let mut st_views = std::collections::HashMap::new();
        let mut host_bufs: std::collections::HashMap<String, Vec<f32>> = std::collections::HashMap::new();
        let mut shapes_map: std::collections::HashMap<String, Vec<usize>> = std::collections::HashMap::new();
        for (name, tensor) in &tensors {
            let host_data = tensor.to_host(ser_stream)
                .map_err(|e| anyhow::anyhow!("Checkpoint DtoH '{name}': {e}"))?;
            shapes_map.insert(name.clone(), tensor.shape().to_vec()); // cpu-side shape clone
            host_bufs.insert(name.clone(), host_data);
        }
        for (name, host_data) in &host_bufs {
            let shape = shapes_map.get(name).cloned().unwrap_or_default();
            let bytes: &[u8] = unsafe {
                std::slice::from_raw_parts(
                    host_data.as_ptr() as *const u8,
                    host_data.len() * std::mem::size_of::<f32>(),
                )
            };
            let view = safetensors::tensor::TensorView::new(
                safetensors::Dtype::F32, shape, bytes,
            ).map_err(|e| anyhow::anyhow!("TensorView '{name}': {e}"))?;
            st_views.insert(name.as_str(), view);
        }
        let data = safetensors::serialize(st_views, arch_metadata)
            .map_err(|e| anyhow::anyhow!("Failed to serialize safetensors: {}", e))?;

        Ok(data)
    }

    /// Inject pre-uploaded GPU data (e.g. from a `DoubleBufferedLoader`).
    ///
    /// The trainer's `train_epoch` lazily uploads data on first call.
    /// Use this to provide data that was uploaded in advance by a
    /// `DoubleBufferedLoader`, skipping the per-fold upload latency.
    pub fn set_gpu_data(&mut self, data: DqnGpuData) {
        info!(
            "DqnTrainer: injected pre-uploaded GPU data ({} bars, {:.1} MB)",
            data.num_bars,
            data.vram_bytes() as f64 / 1_048_576.0,
        );
        self.gpu_data = Some(data);
    }

    /// Drop cached GPU data, freeing VRAM for the next fold.
    pub fn clear_gpu_data(&mut self) {
        if self.gpu_data.is_some() {
            info!("DqnTrainer: cleared GPU data (VRAM freed)");
            self.gpu_data = None;
        }
    }

    /// BUG #38 FIX: Clear replay buffer of contaminated experiences
    pub async fn clear_replay_buffer(&mut self) -> Result<()> {
        let mut agent = self.agent.write().await;
        agent.clear_replay_buffer().map_err(|e| {
            anyhow::anyhow!("Failed to clear replay buffer: {}", e)
        })?;
        let buffer_size = agent.get_replay_buffer_size().unwrap_or(0);
        info!("Replay buffer cleared successfully. Current size: {}", buffer_size);
        Ok(())
    }

    /// BUG #38 FIX: Reset target network to match current network
    pub async fn reset_target_network(&mut self) -> Result<()> {
        let mut agent = self.agent.write().await;
        agent.reset_target_network().map_err(|e| {
            anyhow::anyhow!("Failed to reset target network: {}", e)
        })?;
        info!("Target network reset successfully");
        Ok(())
    }




    /// Get per-epoch training loss history (for smoke test verification)
    pub fn loss_history(&self) -> &[f64] {
        &self.loss_history
    }

    /// Get per-epoch validation loss history
    pub fn val_loss_history(&self) -> &[f64] {
        &self.val_loss_history
    }

    /// Get current epsilon from the DQN agent
    pub async fn get_agent_epsilon(&self) -> f32 {
        let agent_lock = self.agent.read().await;
        agent_lock.get_epsilon()
    }
}

// ---------------------------------------------------------------------------
// GPU Q-value diagnostics (Task 6)
// ---------------------------------------------------------------------------



