//! DQN Training with Full Rainbow Architecture (DEFAULT)
//!
//! This implementation includes ALL Rainbow DQN components by default:
//! - Double DQN (always enabled) - reduces overestimation bias
//! - Dueling Networks (separate value/advantage streams) - better credit assignment
//! - Prioritized Experience Replay (PER) - samples high TD-error transitions
//! - Multi-Step Returns (n=3) - balances bias vs variance
//! - Distributional RL (C51 with 51 atoms) - models full return distribution
//! - Noisy Networks (learnable exploration) - replaces epsilon-greedy
//!
//! # Usage
//!
//! ```bash
//! # Train with default parameters (100 epochs, ALL Rainbow features enabled)
//! cargo run -p ml --example train_dqn --release --features cuda
//!
//! # Disable specific components (opt-out)
//! cargo run -p ml --example train_dqn --release --features cuda -- \
//!   --no-dueling --no-distributional --no-noisy-nets
//!
//! # Train vanilla DQN (all Rainbow features disabled)
//! cargo run -p ml --example train_dqn --release --features cuda -- \
//!   --no-dueling --no-distributional --no-noisy-nets --n-steps 1
//!
//! # Custom hyperparameters with Rainbow
//! cargo run -p ml --example train_dqn --release --features cuda -- \
//!   --epochs 500 \
//!   --n-steps 5 \
//!   --num-atoms 101 \
//!   --dueling-hidden-dim 256
//! ```

// Use mimalloc allocator for 10-25% performance improvement
#[cfg(feature = "mimalloc-allocator")]
use mimalloc::MiMalloc;
#[cfg(feature = "mimalloc-allocator")]
#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

use anyhow::{Context, Result};
use clap::Parser;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use clap::ValueEnum;
use tokio::signal;
use tracing::{info, warn};
use tracing_subscriber::FmtSubscriber;

use ml::checkpoint::{CheckpointConfig, CheckpointManager};
use ml::data_loaders::BarSamplingMethod;
use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
use ml::evaluation::{EvaluationEngine, PerformanceMetrics};
use ml::trainers::TargetUpdateMode;

/// Barrier presets for common HFT trading strategies
///
/// Each preset configures the triple barrier method (profit target, stop loss, max hold time)
/// for a specific HFT strategy profile.
#[derive(Debug, Clone, Copy, ValueEnum)]
enum BarrierPreset {
    /// Generic preset: Wide targets, long holds (100/50 bps, 1 hour)
    /// Use for: Trend-following, position trading
    Generic,
    
    /// Scalping preset: Tight targets, quick exits (7/4 bps, 30 seconds)
    /// Use for: High-frequency scalping, micro-movements
    Scalping,
    
    /// Mean-reversion preset: Moderate targets, patient holds (25/12 bps, 2.5 minutes)
    /// Use for: Statistical arbitrage, contrarian strategies
    MeanReversion,
}

impl BarrierPreset {
    /// Convert preset to barrier parameters: (profit_bps, stop_bps, max_hold_seconds)
    pub fn to_barrier_params(self) -> (u32, u32, u64) {
        match self {
            Self::Scalping => (
                7,   // 7 bps profit target (0.07%, tight for quick scalps)
                4,   // 4 bps stop loss (0.04%, tight risk control)
                30,  // 30 seconds max hold (fast in/out)
            ),
            Self::MeanReversion => (25, 12, 150),  // 25/12 bps, 2.5 minutes
            Self::Generic => (100, 50, 3600),      // 100/50 bps, 1 hour (current defaults)
        }
    }
}

/// Train DQN model on market data
#[derive(Debug, Parser)]
#[command(name = "train_dqn", about = "Train DQN model on market data")]
struct Opts {
    /// Number of training epochs
    #[arg(long, default_value = "100")]
    epochs: usize,

    /// Learning rate
    /// BUG #18 FIX (Wave 16S-V17): Reduced 10× to 0.00001 to prevent Q-value explosion
    /// after Bug #16 fix increased reward scale (raw portfolio values vs normalized).
    /// Previous: 0.0001 caused gradient collapse (Q-values hit 1000.0 clamp, grad_norm → 0).
    #[arg(long, default_value = "0.00001")]
    learning_rate: f64,

    /// Batch size (max 230 for RTX 3050 Ti 4GB)
    /// Optimal value from hyperopt (42 trials, 2025-10-31): 32
    #[arg(long, default_value = "32")]
    batch_size: usize,

    /// Discount factor (gamma)
    /// Optimal value from hyperopt (42 trials, 2025-10-31): 0.9626
    #[arg(long, default_value = "0.9626")]
    gamma: f64,

    /// Checkpoint save frequency (epochs)
    #[arg(long, default_value = "10")]
    checkpoint_frequency: usize,

    /// Output directory for trained model
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: String,

    /// Data directory containing DBN files
    #[arg(long, default_value = "test_data/real/databento/ml_training")]
    data_dir: String,

    /// Parquet file path (overrides data_dir if specified)
    #[arg(long)]
    parquet_file: Option<String>,

    /// Feature cache directory (optional, enables faster hyperopt by caching computed features)
    /// Example: --cache-dir /tmp/dqn_feature_cache
    #[arg(long)]
    cache_dir: Option<PathBuf>,

    /// Verbose logging
    #[arg(short, long)]
    verbose: bool,

    /// Enable debug logging for detailed diagnostics (REWARD_DEBUG, gradient norms, etc.)
    /// Disabled by default to reduce log noise during hyperopt campaigns
    #[arg(long)]
    debug_logging: bool,

    /// Enable early stopping (recommended, use --no-early-stopping to disable)
    #[arg(long)]
    early_stopping: bool,

    /// Disable early stopping
    #[arg(long)]
    no_early_stopping: bool,

    /// Q-value floor threshold for early stopping
    #[arg(long, default_value = "-5.0")]
    q_value_floor: f64,

    /// Minimum loss improvement percentage for plateau detection
    #[arg(long, default_value = "2.0")]
    min_loss_improvement: f64,

    /// Plateau detection window size (epochs)
    /// Optimal value from hyperopt: 5
    #[arg(long, default_value = "5")]
    plateau_window: usize,

    /// Minimum epochs before early stopping can trigger
    /// Updated to 50 to prevent premature stopping (was 10)
    #[arg(long, default_value = "50")]
    min_epochs_before_stopping: usize,

    /// Initial exploration rate (epsilon start)
    /// Updated to 0.3 for more initial exploration (was 1.0)
    #[arg(long, default_value = "0.3")]
    epsilon_start: f64,

    /// Final exploration rate (epsilon end)
    /// Updated to 0.05 to maintain exploration (was 0.01)
    #[arg(long, default_value = "0.05")]
    epsilon_end: f64,

    /// Exploration decay rate
    /// Updated to 0.995 for slower decay (was 0.9968)
    #[arg(long, default_value = "0.995")]
    epsilon_decay: f64,

    /// Replay buffer capacity
    /// Optimal value from hyperopt: 104346
    #[arg(long, default_value = "104346")]
    buffer_size: usize,

    /// Minimum replay buffer size before training starts
    /// Updated to 500 for more diverse experiences (was auto-calculated as batch_size * 2 = 64)
    #[arg(long, default_value = "500")]
    min_replay_size: usize,

    /// Checkpoint directory (overrides output_dir for checkpoints)
    #[arg(long)]
    checkpoint_dir: Option<String>,

    /// Alternative bar sampling method (time, tick, volume, dollar, imbalance, run)
    #[arg(long, default_value = "time")]
    bar_method: String,

    /// HOLD penalty weight (higher = stronger penalty for holding)
    #[arg(long, default_value = "0.01")]
    hold_penalty_weight: f64,

    /// Price movement threshold for dynamic HOLD reward (as fraction, e.g., 0.01 = 1%)
    #[arg(long, default_value = "0.01")]
    bar_threshold: Option<f64>,

    /// Disable preprocessing (use raw non-stationary prices - NOT RECOMMENDED)
    #[arg(long)]
    no_preprocessing: bool,

    /// Preprocessing window size (default: 50 bars)
    #[arg(long, default_value = "50")]
    preprocessing_window: i64,

    /// Preprocessing clip sigma (default: 5.0σ)
    #[arg(long, default_value = "5.0")]
    preprocessing_clip_sigma: f64,

    /// Warmup steps for random exploration before training (Rainbow DQN: 80K)
    /// Default: ADAPTIVE (0 for <200K steps, scaled 200K-1M, 80K for >1M)
    /// Set explicitly to override adaptive behavior
    #[arg(long)]
    warmup_steps: Option<usize>,

    /// Initial capital for portfolio trading (default: $100,000, minimum: $1,000)
    #[arg(long, default_value = "100000.0")]
    initial_capital: f32,

    /// Cash reserve requirement as a percentage of portfolio value (0.0-100.0)
    #[arg(long, default_value = "0.0")]
    cash_reserve_percent: f64,

    /// Polyak averaging coefficient (tau) for soft target updates (default: 0.001 = soft updates)
    /// Rainbow DQN standard: tau=0.001 (693-step convergence half-life)
    /// Lower values = slower convergence, higher values = faster convergence
    #[arg(long, default_value = "0.001")]
    tau: f64,

    /// Disable soft target updates (Polyak averaging) and use hard updates instead
    /// Soft updates blend target network gradually with main network (default)
    /// Default: soft updates (tau=0.001, gradual blending every step)
    #[arg(long)]
    no_soft_updates: bool,

    /// Maximum absolute position size for action masking (1.0-10.0 contracts)
    /// Default: 10.0 (matches hyperopt Trial #26)
    #[arg(long, default_value = "10.0")]
    max_position: f64,

    /// Entropy regularization coefficient (0.0-0.1)
    /// Controls exploration diversity via action entropy bonus
    /// Default: 0.01 (1% entropy bonus from Wave 9-13)
    #[arg(long, default_value = "0.01")]
    entropy_coefficient: f64,

    /// Transaction cost multiplier (0.5-2.0)
    /// Multiplies base transaction fees: Market 0.15%, LimitMaker 0.05%, IoC 0.10%
    /// Default: 1.0 (100% of base fees, production default)
    #[arg(long, default_value = "1.0")]
    transaction_cost_multiplier: f64,

    /// Barrier preset for triple barrier method (generic, scalping, mean-reversion)
    /// Provides strategy-specific defaults that can be overridden with explicit args
    #[arg(long, value_enum, default_value = "generic")]
    barrier_preset: BarrierPreset,

    /// Huber loss delta parameter (10.0-200.0)
    /// Controls transition from quadratic (MSE) to linear (MAE) loss
    /// Default: 10.0 (conservative starting point, hyperopt can scale to 15-40 range)
    #[arg(long, default_value = "10.0")]
    huber_delta: f64,

    // Wave 3 (Phase 2): Prioritized Experience Replay (PER) Arguments
    /// Enable Prioritized Experience Replay (PER)
    /// Samples high TD-error transitions more frequently for faster convergence
    /// Expected improvement: 25-40% fewer epochs to reach target performance
    #[arg(long, default_value_t = true)]
    use_per: bool,

    /// PER alpha parameter: prioritization exponent (0.0-1.0)
    /// Controls how much prioritization to use (0.0 = uniform, 1.0 = full priority)
    /// Rainbow DQN standard: 0.6 (balanced prioritization)
    #[arg(long, default_value = "0.6")]
    per_alpha: f64,

    /// PER beta start: importance sampling correction exponent (0.0-1.0)
    /// Anneals from beta_start to 1.0 over training to remove bias
    /// Rainbow DQN standard: 0.4 (start) → 1.0 (end)
    #[arg(long, default_value = "0.4")]
    per_beta_start: f64,

    // Wave 3 (Phase 2): Triple Barrier Method Arguments
    /// Enable triple barrier method for multi-step reward labeling
    #[arg(long)]
    enable_triple_barrier: bool,

    /// Triple barrier profit target in basis points (overrides preset)
    #[arg(long)]
    profit_target: Option<u32>,

    /// Triple barrier stop loss in basis points (overrides preset)
    #[arg(long)]
    stop_loss: Option<u32>,

    /// Triple barrier time limit in seconds (e.g., 3600 = 1 hour)
    #[arg(long)]
    time_limit: Option<u64>,

    // Regime-Conditional DQN Arguments
    /// Enable regime-conditional Q-network (3 heads: Trending, Ranging, Volatile)
    /// When enabled, routes actions and training through regime-specific Q-networks
    /// Expected improvement: +10-15% Sharpe ratio via regime-adaptive strategies
    #[arg(long)]
    enable_regime_qnetwork: bool,

    // Wave 6.4: Rainbow DQN Opt-Out Flags (ALL ENABLED BY DEFAULT)
    /// Disable dueling networks (enabled by default)
    /// Dueling networks separate value and advantage streams for better credit assignment
    #[arg(long)]
    no_dueling: bool,

    /// Disable distributional RL / C51 algorithm (enabled by default)
    /// Distributional RL models full return distribution instead of scalar Q-values
    #[arg(long)]
    no_distributional: bool,

    /// Disable noisy networks (enabled by default)
    /// Noisy networks add learnable noise to network parameters for exploration
    #[arg(long)]
    no_noisy_nets: bool,

    /// Set n_steps for multi-step returns (default: 3)
    /// Higher values reduce bias but increase variance
    #[arg(long, default_value = "3")]
    n_steps: usize,

    /// Number of atoms for C51 distributional RL (default: 51)
    #[arg(long, default_value = "51")]
    num_atoms: usize,

    /// Minimum value for C51 distribution support (default: -2.0)
    /// Bug #5 fix: Corrected from -1000.0 (500x too large)
    #[arg(long, default_value = "-2.0")]
    v_min: f64,

    /// Maximum value for C51 distribution support (default: 2.0)
    /// Bug #5 fix: Corrected from 1000.0 (500x too large)
    #[arg(long, default_value = "2.0")]
    v_max: f64,

    /// Noisy network sigma init (default: 0.5, Rainbow DQN standard)
    #[arg(long, default_value = "0.5")]
    noisy_sigma_init: f64,

    /// Dueling hidden dimension (default: 128)
    #[arg(long, default_value = "128")]
    dueling_hidden_dim: usize,

    /// Feature stats collection ratio (0.0-1.0, default: 0.3 = 30% of epochs)
    /// Determines what percentage of training epochs to use for collecting normalization statistics
    /// Example: 100 epochs * 0.3 = 30 epochs (capped by max-feature-stats-epochs)
    #[arg(long, default_value = "0.3")]
    feature_stats_collection_ratio: f32,

    /// Maximum epochs for feature statistics collection (default: 10)
    /// Acts as a cap: min(epochs * ratio, max_epochs)
    /// Set to 0 to disable cap (use pure percentage-based calculation)
    #[arg(long, default_value = "10")]
    max_feature_stats_epochs: usize,

    /// Load hyperparameters from hyperopt JSON file
    /// Path to JSON file containing best trial hyperparameters
    /// When specified, overrides CLI defaults with saved optimal values
    #[arg(long)]
    load_hyperopt_json: Option<String>,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Parse CLI options
    let mut opts = Opts::parse();

    // Load hyperparameters from JSON if specified
    if let Some(ref json_path) = opts.load_hyperopt_json {
        use ml::hyperopt::adapters::dqn::BestTrialExport;

        let json_str = std::fs::read_to_string(json_path)
            .with_context(|| format!("Failed to read hyperopt JSON from: {}", json_path))?;

        let best_trial: BestTrialExport = serde_json::from_str(&json_str)
            .with_context(|| format!("Failed to parse hyperopt JSON from: {}", json_path))?;

        println!("📥 Loading hyperparameters from: {}", json_path);
        println!("   Trial #{} (Sharpe: {:.4}, Win Rate: {:.2}%)",
                 best_trial.trial_number, best_trial.sharpe, best_trial.win_rate);

        // Override CLI defaults with JSON values
        let params = &best_trial.hyperparameters;
        opts.learning_rate = params.learning_rate;
        opts.batch_size = params.batch_size;
        opts.gamma = params.gamma;
        opts.buffer_size = params.buffer_size;
        opts.hold_penalty_weight = params.hold_penalty_weight;
        opts.max_position = params.max_position_absolute;
        opts.huber_delta = params.huber_delta;
        opts.entropy_coefficient = params.entropy_coefficient;
        opts.transaction_cost_multiplier = params.transaction_cost_multiplier;
        opts.per_alpha = params.per_alpha;
        opts.per_beta_start = params.per_beta_start;
        opts.v_min = params.v_min;
        opts.v_max = params.v_max;
        opts.noisy_sigma_init = params.noisy_sigma_init;
        opts.dueling_hidden_dim = params.dueling_hidden_dim;
        opts.n_steps = params.n_steps;
        opts.num_atoms = params.num_atoms;

        println!("✅ Hyperparameters loaded from JSON");
    }

    // Setup logging
    let level = if opts.verbose {
        tracing::Level::DEBUG
    } else {
        tracing::Level::INFO
    };

    let subscriber = FmtSubscriber::builder().with_max_level(level).finish();
    tracing::subscriber::set_global_default(subscriber)
        .context("Failed to set tracing subscriber")?;

    #[cfg(feature = "mimalloc-allocator")]
    info!("🚀 Using mimalloc allocator for improved performance");
    #[cfg(not(feature = "mimalloc-allocator"))]
    info!("ℹ️  Using system allocator (consider --features mimalloc-allocator for 10-25% speedup)");
    info!("🚀 Starting DQN Training");
    info!("Configuration:");
    info!("  • Epochs: {}", opts.epochs);
    info!("  • Learning rate: {}", opts.learning_rate);
    info!("  • Batch size: {}", opts.batch_size);
    info!("  • Gamma: {}", opts.gamma);
    info!(
        "  • Checkpoint frequency: {} epochs",
        opts.checkpoint_frequency
    );
    info!("  • Output directory: {}", opts.output_dir);
    info!("  • Data directory: {}", opts.data_dir);
    info!("  • Bar sampling method: {}", opts.bar_method);
    if let Some(threshold) = opts.bar_threshold {
        info!("  • Bar threshold: {}", threshold);
    }
    info!("  • Epsilon start: {}", opts.epsilon_start);
    info!("  • Epsilon end: {}", opts.epsilon_end);
    info!("  • Epsilon decay: {}", opts.epsilon_decay);
    info!("  • Buffer size: {}", opts.buffer_size);
    info!("  • Min replay size: {}", opts.min_replay_size);

    info!("  • Initial capital: ${:.2}", opts.initial_capital);
    info!("  • Cash reserve: {}%", opts.cash_reserve_percent);

    // Log target update configuration
    if opts.no_soft_updates {
        info!("  • Target update mode: Hard (complete replacement every 10K steps)");
        info!("  • Tau (τ): {} (no blending, hard copy)", opts.tau);
    } else {
        info!("  • Target update mode: Soft (Polyak averaging) [DEFAULT]");
        info!(
            "  • Tau (τ): {} (convergence half-life: {:.0} steps)",
            opts.tau,
            (-0.5_f64.ln()) / (-(1.0 - opts.tau).ln())
        );
    }

    // ═══════════════════════════════════════════════════════════════════════════
    // ADAPTIVE WARMUP CALCULATION
    // ═══════════════════════════════════════════════════════════════════════════

    // Estimate total training steps for adaptive warmup
    let avg_steps_per_epoch = 1392; // Empirical: ES_FUT_180d.parquet has 1,392 steps/epoch
    let total_steps_estimate = opts.epochs * avg_steps_per_epoch;

    info!(
        "📊 Training length estimate: {}K steps ({} epochs × {} steps/epoch)",
        total_steps_estimate / 1000,
        opts.epochs,
        avg_steps_per_epoch
    );

    // Adaptive warmup calculation
    let effective_warmup = if let Some(explicit_warmup) = opts.warmup_steps {
        // User explicitly set warmup - respect it
        info!(
            "🎯 Using EXPLICIT warmup: {}K steps (user override)",
            explicit_warmup / 1000
        );
        explicit_warmup
    } else {
        // Adaptive warmup based on training length
        let adaptive_warmup = match total_steps_estimate {
            0..=200_000 => {
                info!("⚡ ADAPTIVE WARMUP: 0 steps (short training <200K steps)");
                0
            },
            200_001..=500_000 => {
                let warmup = total_steps_estimate / 20; // 5% warmup
                info!(
                    "⚡ ADAPTIVE WARMUP: {}K steps (~5% of {}K total)",
                    warmup / 1000,
                    total_steps_estimate / 1000
                );
                warmup
            },
            500_001..=1_000_000 => {
                let warmup = total_steps_estimate / 12; // ~8% warmup
                info!(
                    "⚡ ADAPTIVE WARMUP: {}K steps (~8% of {}K total)",
                    warmup / 1000,
                    total_steps_estimate / 1000
                );
                warmup
            },
            _ => {
                info!("⚡ ADAPTIVE WARMUP: 80K steps (Rainbow DQN standard for >1M steps)");
                80_000 // Full Rainbow warmup for very long runs
            },
        };
        adaptive_warmup
    };

    // Warning: warmup consuming too much of training
    if effective_warmup > 0 {
        let warmup_ratio = effective_warmup as f64 / total_steps_estimate as f64;
        if warmup_ratio > 0.10 {
            warn!(
                "⚠️  Warmup period ({}K steps) is {:.1}% of total training ({}K steps)",
                effective_warmup / 1000,
                warmup_ratio * 100.0,
                total_steps_estimate / 1000
            );
            warn!("⚠️  This may significantly delay learning. Consider:");
            warn!(
                "    • Increase --epochs to lengthen training (recommended: {}+)",
                (effective_warmup * 10) / avg_steps_per_epoch
            );
            warn!(
                "    • Reduce --warmup-steps to {} (10% of total)",
                total_steps_estimate / 10
            );
            warn!("    • Set --warmup-steps 0 to disable warmup entirely");
        }

        info!(
            "  • Warmup steps: {}K ({:.1}% of training, Rainbow DQN random exploration)",
            effective_warmup / 1000,
            warmup_ratio * 100.0
        );
    } else {
        info!("  • Warmup steps: 0 (disabled for short training runs)");
    }

    // Validate initial capital
    if opts.initial_capital < 1000.0 {
        eprintln!("❌ Error: initial_capital must be >= $1,000 (got: ${:.2})", opts.initial_capital);
        eprintln!("   Use --initial-capital <amount> to specify a valid amount");
        std::process::exit(1);
    }

    // Validate cash reserve percent
    if !(0.0..=100.0).contains(&opts.cash_reserve_percent) {
        eprintln!(
            "❌ Error: cash_reserve_percent must be between 0.0 and 100.0 (got: {})",
            opts.cash_reserve_percent
        );
        eprintln!("   Use --cash-reserve-percent <percentage> to specify a valid amount");
        std::process::exit(1);
    }

    // Apply barrier preset with explicit arg override logic
    let (preset_profit, preset_stop, preset_time) = opts.barrier_preset.to_barrier_params();
    let effective_profit_target = opts.profit_target.unwrap_or(preset_profit);
    let effective_stop_loss = opts.stop_loss.unwrap_or(preset_stop);
    let effective_time_limit = opts.time_limit.unwrap_or(preset_time);

    info!("  • Barrier preset: {:?}", opts.barrier_preset);
    info!("    - Profit target: {} bps ({}%)", effective_profit_target, effective_profit_target as f64 / 100.0);
    info!("    - Stop loss: {} bps ({}%)", effective_stop_loss, effective_stop_loss as f64 / 100.0);
    info!("    - Max hold time: {}s", effective_time_limit);
    if opts.profit_target.is_some() || opts.stop_loss.is_some() || opts.time_limit.is_some() {
        info!("    (Preset overridden by explicit CLI args)");
    }

    // Setup graceful shutdown handler for containerized environments (RunPod, Docker, K8s)
    let shutdown_flag = Arc::new(AtomicBool::new(false));
    let shutdown_clone = shutdown_flag.clone();

    tokio::spawn(async move {
        let ctrl_c = signal::ctrl_c();

        #[cfg(unix)]
        {
            use tokio::signal::unix::{signal, SignalKind};
            let mut sigterm =
                signal(SignalKind::terminate()).expect("Failed to setup SIGTERM handler");

            tokio::select! {
                _ = ctrl_c => {
                    info!("🛑 Received Ctrl+C, initiating graceful shutdown...");
                }
                _ = sigterm.recv() => {
                    info!("🛑 Received SIGTERM, initiating graceful shutdown...");
                }
            }
        }

        #[cfg(not(unix))]
        {
            ctrl_c.await.expect("Failed to listen for Ctrl+C");
            info!("🛑 Received Ctrl+C, initiating graceful shutdown...");
        }

        shutdown_clone.store(true, Ordering::Relaxed);
    });

    info!("✅ Graceful shutdown handler registered (Ctrl+C / SIGTERM)");

    // Determine early stopping (enabled by default, unless --no-early-stopping is specified)
    let early_stopping_enabled = !opts.no_early_stopping;
    info!(
        "  • Early stopping: {}",
        if early_stopping_enabled {
            "enabled"
        } else {
            "disabled"
        }
    );
    if early_stopping_enabled {
        info!("    - Q-value floor: {}", opts.q_value_floor);
        info!("    - Min loss improvement: {}%", opts.min_loss_improvement);
        info!("    - Plateau window: {} epochs", opts.plateau_window);
        info!(
            "    - Min epochs before stopping: {}",
            opts.min_epochs_before_stopping
        );
    }

    // Create output and checkpoint directories
    let output_path = PathBuf::from(&opts.output_dir);
    let checkpoint_path = if let Some(ref dir) = opts.checkpoint_dir {
        PathBuf::from(dir)
    } else {
        output_path.clone()
    };

    if !output_path.exists() {
        std::fs::create_dir_all(&output_path).context("Failed to create output directory")?;
        info!("✅ Created output directory: {}", opts.output_dir);
    }

    if !checkpoint_path.exists() && checkpoint_path != output_path {
        std::fs::create_dir_all(&checkpoint_path)
            .context("Failed to create checkpoint directory")?;
        info!(
            "✅ Created checkpoint directory: {}",
            checkpoint_path.display()
        );
    }

    if opts.checkpoint_dir.is_some() {
        info!("  • Checkpoint directory: {}", checkpoint_path.display());
    }

    // Configure DQN hyperparameters with optimal values from hyperopt (42 trials, 2025-10-31)
    let hyperparams = DQNHyperparameters {
        learning_rate: opts.learning_rate,
        batch_size: opts.batch_size,
        gamma: opts.gamma,
        epsilon_start: opts.epsilon_start,
        epsilon_end: opts.epsilon_end,
        epsilon_decay: opts.epsilon_decay,
        buffer_size: opts.buffer_size,
        min_replay_size: opts.min_replay_size, // Configurable min replay size
        epochs: opts.epochs,
        checkpoint_frequency: opts.checkpoint_frequency,
        early_stopping_enabled,
        q_value_floor: opts.q_value_floor,
        min_loss_improvement_pct: opts.min_loss_improvement,
        plateau_window: opts.plateau_window,
        min_epochs_before_stopping: opts.min_epochs_before_stopping, // NOW CONFIGURABLE!
        hold_penalty: -0.001,
        // Fix #5: Increase gradient clipping threshold (clipping should be rare, not every step)
        gradient_clip_norm: Some(100.0), // Was 10.0, increase 10x to allow healthy gradients
        // Fix #4: Enable Huber loss with scaled delta
        use_huber_loss: true,
        huber_delta: opts.huber_delta,
        // Enable Double DQN to reduce overestimation bias
        use_double_dqn: true,
        // HOLD penalty weight (Bug #3 fix)
        hold_penalty_weight: opts.hold_penalty_weight, // Configurable via CLI
        // Price movement threshold for HOLD penalty (2% price movement)
        movement_threshold: 0.02,
        // Wave 14 Agent 32: Preprocessing configuration
        enable_preprocessing: !opts.no_preprocessing, // Enabled by default, disable with --no-preprocessing
        preprocessing_window: opts.preprocessing_window,
        preprocessing_clip_sigma: opts.preprocessing_clip_sigma,

        // Target update configuration (soft updates enabled by default)
        tau: opts.tau, // CLI-configurable (default: 0.001 = soft updates)
        target_update_mode: if opts.no_soft_updates {
            TargetUpdateMode::Hard
        } else {
            TargetUpdateMode::Soft
        },

        // P2-B Enhancement: Cash reserve requirement
        cash_reserve_percent: opts.cash_reserve_percent, // Configurable via CLI
        target_update_frequency: 10000, // Hard update frequency (every 10K steps)

        // Rainbow DQN warmup
        warmup_steps: effective_warmup, // Adaptive warmup (0 for <200K, scaled 200K-1M, 80K for >1M)

        // P2-A Enhancement
        initial_capital: opts.initial_capital,

        // WAVE 16S: Adaptive Risk Management (ALL ENABLED BY DEFAULT)
        enable_kelly_sizing: true,
        enable_volatility_epsilon: true,
        enable_risk_adjusted_rewards: true,
        kelly_fractional: 0.5,
        kelly_max_fraction: 0.25,
        kelly_min_trades: 20,
        volatility_window: 20,

        // WAVE 35: Advanced Features
        enable_regime_qnetwork: opts.enable_regime_qnetwork, // CLI-configurable via --enable-regime-qnetwork
        enable_compliance: true, // Enabled by default

        // WAVE 16: Core Risk Management (ALL ENABLED BY DEFAULT)
        enable_drawdown_monitoring: true,
        enable_position_limits: true,
        enable_circuit_breaker: true,

        // Wave 16 Portfolio Features (ALL ENABLED BY DEFAULT)
        enable_action_masking: true,
        enable_entropy_regularization: true,
        enable_stress_testing: true,
        max_position_absolute: opts.max_position,

        // Wave 17: Hyperopt 9D Search Space Extensions
        entropy_coefficient: Some(opts.entropy_coefficient),
        transaction_cost_multiplier: opts.transaction_cost_multiplier,

        // Wave 3 (Phase 2): Prioritized Experience Replay Configuration
        use_per: opts.use_per,
        per_alpha: opts.per_alpha,
        per_beta_start: opts.per_beta_start,

        // Wave 3 (Phase 2): Triple Barrier Method Configuration
        enable_triple_barrier: opts.enable_triple_barrier,
        triple_barrier_profit_target_bps: effective_profit_target,
        triple_barrier_stop_loss_bps: effective_stop_loss,
        triple_barrier_max_holding_seconds: effective_time_limit,

        // Wave 6.4: Rainbow DQN Features (ALL ENABLED BY DEFAULT)
        // Wave 2.1: Dueling Networks
        use_dueling: !opts.no_dueling,  // Default: enabled, opt-out with --no-dueling
        dueling_hidden_dim: opts.dueling_hidden_dim,  // Default: 128

        // Wave 2.2: Multi-Step Returns
        n_steps: opts.n_steps,  // Default: 3

        // Wave 2.3: Distributional RL (C51)
        use_distributional: !opts.no_distributional,  // Default: enabled, opt-out with --no-distributional
        num_atoms: opts.num_atoms,  // Default: 51
        v_min: opts.v_min,  // Default: -1000.0
        v_max: opts.v_max,  // Default: 1000.0

        // Wave 2.4: Noisy Networks
        use_noisy_nets: !opts.no_noisy_nets,  // Default: enabled, opt-out with --no-noisy-nets
        noisy_sigma_init: opts.noisy_sigma_init,  // Default: 0.5

        // Two-Phase Feature Normalization Configuration
        feature_stats_collection_ratio: opts.feature_stats_collection_ratio,  // Default: 0.3 (30%)
        max_feature_stats_epochs: if opts.max_feature_stats_epochs == 0 {
            None  // 0 means no cap
        } else {
            Some(opts.max_feature_stats_epochs)  // Default: Some(10)
        },

        // WAVE 23 P0: Early Stopping for Gradient Collapse
        gradient_collapse_multiplier: 100.0,  // Adaptive threshold (LR × 100)
        gradient_collapse_patience: 5,        // 5 consecutive epochs before early stop

        // Fill remaining fields from defaults
        ..Default::default()
    };

    // Configure alternative bar sampling (Wave B)
    let bar_sampling = match opts.bar_method.as_str() {
        "tick" => BarSamplingMethod::TickBars(opts.bar_threshold.unwrap_or(100.0) as usize),
        "volume" => BarSamplingMethod::VolumeBars(opts.bar_threshold.unwrap_or(10000.0)),
        "dollar" => BarSamplingMethod::DollarBars(opts.bar_threshold.unwrap_or(2_000_000.0)),
        "imbalance" => BarSamplingMethod::ImbalanceBars(opts.bar_threshold.unwrap_or(1000.0)),
        "run" => BarSamplingMethod::RunBars(opts.bar_threshold.unwrap_or(50.0) as usize),
        _ => BarSamplingMethod::TimeBars,
    };

    info!("✅ Bar sampling configured: {:?}", bar_sampling);

    // Create DQN trainer with debug logging flag
    let mut trainer = DQNTrainer::new_with_debug(hyperparams, opts.debug_logging)
        .context("Failed to create DQN trainer")?;

    // Wire feature cache if provided
    if let Some(cache_dir) = opts.cache_dir {
        info!("🗂️  Feature cache enabled: {:?}", cache_dir);
        trainer = trainer.with_feature_cache(cache_dir);
    } else {
        info!("📊 Feature cache disabled, computing features from scratch");
    }

    // Note: DQN trainer will need to accept bar_sampling parameter
    // This requires updating DQNTrainer to use DbnSequenceLoader
    info!("✅ DQN trainer initialized");

    // Setup checkpoint manager
    let checkpoint_config = CheckpointConfig {
        base_dir: output_path.clone(),
        max_checkpoints_per_model: 10,
        auto_cleanup: true,
        validate_checksums: true,
        ..Default::default()
    };

    let _checkpoint_manager =
        CheckpointManager::new(checkpoint_config).context("Failed to create checkpoint manager")?;

    info!("✅ Checkpoint manager initialized (max 10 checkpoints, auto-cleanup enabled)");

    // Create checkpoint callback with interruption handling
    let checkpoint_dir_for_callback = opts
        .checkpoint_dir
        .clone()
        .unwrap_or_else(|| opts.output_dir.clone());
    let shutdown_check = shutdown_flag.clone();

    let checkpoint_callback =
        move |epoch: usize, model_data: Vec<u8>, is_best: bool| -> Result<String> {
            // Check if shutdown was requested
            let interrupted = shutdown_check.load(Ordering::Relaxed);

            let filename = if is_best {
                // Best model checkpoint (overwrites previous best)
                "dqn_best_model.safetensors".to_string()
            } else if interrupted {
                format!("dqn_interrupted_epoch{}.safetensors", epoch)
            } else {
                // Periodic checkpoint
                format!("dqn_epoch_{}.safetensors", epoch)
            };

            let checkpoint_path = PathBuf::from(&checkpoint_dir_for_callback).join(filename);

            // Save checkpoint to disk
            std::fs::write(&checkpoint_path, &model_data)
                .context(format!("Failed to save checkpoint: {:?}", checkpoint_path))?;

            let checkpoint_type = if is_best {
                "🎉 BEST"
            } else if interrupted {
                "⚠️  INTERRUPTED"
            } else {
                "💾 PERIODIC"
            };

            info!(
                "{} Checkpoint saved: {} ({} bytes)",
                checkpoint_type,
                checkpoint_path.display(),
                model_data.len()
            );

            Ok(checkpoint_path.to_string_lossy().to_string())
        };

    // Train the model
    info!("\n🏋️  Starting training...\n");
    let start_time = std::time::Instant::now();

    let metrics = if let Some(ref parquet_path) = opts.parquet_file {
        info!("Using Parquet file: {}", parquet_path);
        trainer
            .train_from_parquet(parquet_path, checkpoint_callback)
            .await
            .context("Training from Parquet failed")?
    } else {
        info!("Using DBN directory: {}", opts.data_dir);
        trainer
            .train(&opts.data_dir, checkpoint_callback)
            .await
            .context("Training failed")?
    };

    let training_duration = start_time.elapsed();

    // Check if training was interrupted
    if shutdown_flag.load(Ordering::Relaxed) {
        info!("\n⚠️  Training was interrupted by shutdown signal");
        info!("💾 Interrupted checkpoint saved, safe to terminate");
        info!("📊 Partial training metrics:");
        info!("  • Epochs completed: {}", metrics.epochs_trained);
        info!(
            "  • Training time: {:.1}s ({:.1} min)",
            metrics.training_time_seconds,
            metrics.training_time_seconds / 60.0
        );
        return Ok(());
    }

    // Print final metrics
    info!("\n✅ Training completed successfully!");
    info!("\n📊 Final Metrics:");
    info!("  • Final loss: {:.6}", metrics.loss);
    info!("  • Epochs trained: {}", metrics.epochs_trained);
    info!(
        "  • Training time: {:.1}s ({:.1} min)",
        metrics.training_time_seconds,
        metrics.training_time_seconds / 60.0
    );
    info!(
        "  • Actual elapsed time: {:.1}s (includes data loading + overhead)",
        training_duration.as_secs_f64()
    );
    info!(
        "  • Convergence: {}",
        if metrics.convergence_achieved {
            "✅ Yes"
        } else {
            "❌ No"
        }
    );

    // Additional metrics from training
    if let Some(avg_q_value) = metrics.additional_metrics.get("avg_q_value") {
        info!("  • Average Q-value: {:.4}", avg_q_value);
    }
    if let Some(final_epsilon) = metrics.additional_metrics.get("final_epsilon") {
        info!("  • Final epsilon: {:.4}", final_epsilon);
    }
    if let Some(grad_norm) = metrics.additional_metrics.get("avg_gradient_norm") {
        info!("  • Average gradient norm: {:.6}", grad_norm);
    }

    // Always run backtest evaluation
    info!("\n📈 Running backtest evaluation...");
        
        // Get validation data from trainer
        let val_data = trainer.get_val_data();
        
        if val_data.is_empty() {
            warn!("⚠️  No validation data available for backtest, skipping evaluation");
        } else {
            info!("  • Validation samples: {}", val_data.len());
            
            // Get Kelly fraction for position sizing
            let kelly_fraction = trainer.get_kelly_fraction();
            info!("  • Kelly fraction: {:.4}", kelly_fraction);
            
            // Create evaluation engine with $10K initial capital
            let mut engine = EvaluationEngine::new_with_kelly(10000.0, kelly_fraction);
            
            // Get trained agent
            let agent_arc = trainer.get_agent();
            
            // Collect OHLCV bars for metrics calculation
            let mut ohlcv_bars = Vec::with_capacity(val_data.len());
            
            // Run backtest loop
            for (bar_idx, (feature_vec, target)) in val_data.iter().enumerate() {
                // Extract close price from target
                let close_price = if target.len() >= 2 {
                    target[0]
                } else {
                    feature_vec[3] // Fallback to close log return feature
                };
                
                // Convert feature vector to state tensor
                let state_tensor = match trainer.convert_to_state(feature_vec, close_price) {
                    Ok(tensor) => tensor,
                    Err(_) => continue,
                };
                
                // Extract state as Vec<f32>
                let state_vec: Vec<f32> = match state_tensor.to_vec1() {
                    Ok(vec) => vec,
                    Err(_) => continue,
                };
                
                // Select action (greedy, no epsilon)
                let trading_action = {
                    let mut agent = agent_arc.write().await;
                    match agent.select_action(&state_vec) {
                        Ok(action) => action,
                        Err(_) => continue,
                    }
                };
                
                // Convert FactoredAction to evaluation::Action
                use ml::dqn::TradingAction;
                use ml::evaluation::engine::Action;
                let legacy_action = trading_action.to_legacy_action();
                let action = match legacy_action {
                    TradingAction::Buy => Action::Buy,
                    TradingAction::Sell => Action::Sell,
                    TradingAction::Hold => Action::Hold,
                };
                
                // Create OHLCV bar
                use ml::evaluation::metrics::OHLCVBarF32;
                let bar = OHLCVBarF32 {
                    timestamp: bar_idx as i64,
                    open: close_price as f32,
                    high: close_price as f32,
                    low: close_price as f32,
                    close: close_price as f32,
                    volume: 0.0,
                };
                
                // Process bar in evaluation engine
                engine.process_bar(bar_idx, &bar, action);
                ohlcv_bars.push(bar);
            }
            
            // Close any open position
            if let Some(last_bar) = ohlcv_bars.last() {
                engine.close_position(ohlcv_bars.len() - 1, last_bar);
            }
            
            // Calculate performance metrics
            let perf_metrics = PerformanceMetrics::from_trades(
                &engine.trades,
                engine.initial_capital,
                &ohlcv_bars,
            );
            
            // Print comprehensive metrics
            info!("\n📊 Backtest Results:");
            info!("════════════════════════════════════════════════════════════");
            info!("  Sharpe Ratio:        {:.4}", perf_metrics.sharpe_ratio);
            info!("  Win Rate:            {:.2}%", perf_metrics.win_rate);
            info!("  Max Drawdown:        {:.2}%", perf_metrics.max_drawdown_pct);
            info!("  Total Return:        {:.2}%", perf_metrics.total_return_pct);
            info!("  Total Trades:        {}", perf_metrics.total_trades);
            info!("════════════════════════════════════════════════════════════");
            info!("  Sortino Ratio:       {:.4}", perf_metrics.sortino_ratio);
            info!("  Calmar Ratio:        {:.4}", perf_metrics.calmar_ratio);
            info!("  Omega Ratio:         {:.4}", perf_metrics.omega_ratio);
            info!("════════════════════════════════════════════════════════════");
            info!("  Value at Risk (95%): {:.4}", perf_metrics.var_95);
            info!("  CVaR (95%):          {:.4}", perf_metrics.cvar_95);
            info!("════════════════════════════════════════════════════════════");
            info!("  Beta:                {:.4}", perf_metrics.beta);
            info!("  Alpha:               {:.4}", perf_metrics.alpha);
            info!("  Information Ratio:   {:.4}", perf_metrics.information_ratio);
            info!("════════════════════════════════════════════════════════════");
            info!("  Final Equity:        ${:.2}", perf_metrics.final_equity);
            info!("  Max Equity:          ${:.2}", perf_metrics.max_equity);
            info!("  Avg Trade PnL:       ${:.2}", perf_metrics.avg_trade_pnl);
            info!("════════════════════════════════════════════════════════════");
            
            info!("\n✅ Backtest evaluation complete!");
        }

    // Save final model
    let final_model_path = output_path.join(format!("dqn_final_epoch{}.safetensors", opts.epochs));
    info!("\n💾 Saving final model to: {}", final_model_path.display());

    // Get final model state
    let final_checkpoint_data = trainer
        .serialize_model()
        .await
        .context("Failed to serialize final model")?;

    std::fs::write(&final_model_path, &final_checkpoint_data)
        .context("Failed to save final model")?;

    info!(
        "✅ Final model saved: {} ({} bytes)",
        final_model_path.display(),
        final_checkpoint_data.len()
    );

    info!("\n🎉 DQN training complete!");
    info!("📁 Model files saved to: {}", opts.output_dir);

    Ok(())
}