//! Generic Traits for Hyperparameter Optimization
//!
//! This module provides a trait-based architecture for hyperparameter optimization
//! that works across all Foxhunt ML models (MAMBA-2, DQN, PPO, TFT). The design
//! follows these principles:
//!
//! - **Model-agnostic**: Models implement `HyperparameterOptimizable` trait
//! - **Type-safe**: Parameter spaces are strongly typed with compile-time validation
//! - **Composable**: Traits can be combined for complex optimization workflows
//! - **Efficient**: Zero-cost abstractions with no runtime overhead
//!
//! ## Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────┐
//! │                   HyperparameterOptimizable                 │
//! │  ┌─────────────┐    ┌──────────────┐   ┌─────────────┐    │
//! │  │   MAMBA-2   │    │     DQN      │   │     PPO     │    │
//! │  └─────────────┘    └──────────────┘   └─────────────┘    │
//! └─────────────────────────────────────────────────────────────┘
//!                              ↓
//!                    ┌──────────────────┐
//!                    │  ArgminOptimizer │
//!                    │  (Generic impl)  │
//!                    └──────────────────┘
//!                              ↓
//!                    ┌──────────────────┐
//!                    │ OptimizationResult│
//!                    │  (Trial history)  │
//!                    └──────────────────┘
//! ```
//!
//! ## Usage Example
//!
//! ```rust,no_run
//! use ml::hyperopt::{ArgminOptimizer, HyperparameterOptimizable};
//! use ml::hyperopt::adapters::mamba2::{Mamba2Trainer, Mamba2Params};
//!
//! # async fn example() -> anyhow::Result<()> {
//! // Create model trainer
//! let trainer = Mamba2Trainer::new(
//!     "test_data/ES_FUT_180d.parquet",
//!     50, // epochs
//! )?;
//!
//! // Run optimization
//! let optimizer = ArgminOptimizer::new(30, 5);
//! let result = optimizer.optimize(trainer)?;
//!
//! println!("Best params: {:?}", result.best_params);
//! println!("Best loss: {:.6}", result.best_objective);
//! # Ok(())
//! # }
//! ```

use serde::{Deserialize, Serialize};
use std::fmt::Debug;
use tracing::info;

use crate::MLError;

/// Hyperopt objective mode for two-phase optimization.
///
/// In two-phase optimization, Phase A uses `EpisodeReward` for fast convergence,
/// then Phase B switches to `Sharpe` for financial quality refinement.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Default)]
pub enum ObjectiveMode {
    /// Phase A: Optimize episode reward (fast convergence signal)
    EpisodeReward,
    /// Phase B: Optimize Sharpe ratio (financial quality)
    #[default]
    Sharpe,
}


/// Strategy computed by HardwareBudget for optimal GPU utilization during hyperopt.
#[derive(Debug, Clone)]
pub struct HyperoptStrategy {
    /// Max number of trials to evaluate concurrently on the GPU.
    pub max_concurrent_trials: usize,
    /// VRAM-aware (lower, upper) batch_size bounds for PSO search.
    pub batch_size_bounds: (f64, f64),
    /// Estimated VRAM per trial in MB (for logging/monitoring).
    pub per_trial_memory_mb: f64,
    /// Total VRAM reserved for all concurrent trials (for watchdog).
    pub total_reserved_mb: f64,
}

/// Hardware budget passed to dynamic bounds computation.
///
/// Contains detected GPU memory so that `ParameterSpace` implementations
/// can scale batch_size upper bounds to match available VRAM, instead of
/// using hardcoded limits tuned for a specific GPU (e.g., RTX 3050 Ti 4GB).
#[derive(Debug, Clone)]
pub struct HardwareBudget {
    /// Available GPU memory in MB (0 = CPU-only, use conservative defaults)
    pub gpu_memory_mb: usize,
    /// GPU device name (e.g. "NVIDIA L40S"), empty for CPU-only
    pub gpu_name: String,
}

impl HardwareBudget {
    /// Detect hardware budget from the current system.
    ///
    /// Uses `nvidia-smi` via the existing `detect_gpu_memory()` utility.
    /// Uses **free** VRAM (not total) so that display server, compositor,
    /// and other GPU consumers are already subtracted.
    /// Falls back to `cpu_only()` if detection fails (no GPU, driver error, etc.).
    pub fn detect() -> Self {
        if let Ok((total_mb, free_mb, name)) = ml_core::memory_optimization::auto_batch_size::detect_gpu_memory() {
            // Use free VRAM — on desktop GPUs the display server can consume
            // 300-1500 MB of the total. Using total_mb would overestimate
            // available budget and cause OOM on small GPUs (e.g. RTX 3050 Ti 4GB).
            let gpu_memory_mb = free_mb as usize;
            info!(
                "HardwareBudget: detected {} — {}MB total, {}MB free (using free as budget)",
                name, total_mb as usize, gpu_memory_mb
            );
            Self { gpu_memory_mb, gpu_name: name }
        } else {
            info!("HardwareBudget: no GPU detected, using CPU-only defaults");
            Self::cpu_only()
        }
    }

    /// Fallback for CPU-only or detection failure.
    pub fn cpu_only() -> Self {
        Self { gpu_memory_mb: 0, gpu_name: String::new() }
    }

    /// Create with a specific memory value (for testing).
    pub fn with_memory_mb(gpu_memory_mb: usize) -> Self {
        Self { gpu_memory_mb, gpu_name: String::new() }
    }

    /// Compute the max batch_size for a given model memory profile.
    ///
    /// Formula: `(available_vram * safety_factor - model_overhead) / per_sample_mb`
    /// Result is clamped to `[min_batch, max_batch]`.
    ///
    /// Returns `None` if GPU is not available (budget is 0).
    pub fn max_batch_size(
        &self,
        model_overhead_mb: f64,
        per_sample_mb: f64,
        min_batch: f64,
        max_batch: f64,
    ) -> Option<f64> {
        if self.gpu_memory_mb == 0 {
            return None;
        }
        // Adaptive safety factor: larger GPUs can safely use more VRAM
        let safety_factor = match self.gpu_memory_mb {
            0..=8192 => 0.70,
            8193..=24576 => 0.75,
            24577..=49152 => 0.80,
            _ => 0.85,
        };
        let available = self.gpu_memory_mb as f64 * safety_factor - model_overhead_mb;
        if available <= 0.0 || per_sample_mb <= 0.0 {
            return Some(min_batch);
        }
        Some((available / per_sample_mb).clamp(min_batch, max_batch))
    }

    /// Compute optimal hyperopt strategy based on model size vs available VRAM.
    ///
    /// # Arguments
    /// * `model_overhead_mb` - Fixed per-trial VRAM: model weights + optimizer + gradients + activations
    /// * `per_sample_mb` - VRAM per sample in a batch
    /// * `min_batch` - Minimum batch_size floor
    /// * `max_batch` - Maximum batch_size ceiling
    pub fn plan_hyperopt(
        &self,
        model_overhead_mb: f64,
        per_sample_mb: f64,
        min_batch: f64,
        max_batch: f64,
    ) -> HyperoptStrategy {
        if self.gpu_memory_mb == 0 {
            return HyperoptStrategy {
                max_concurrent_trials: 1,
                batch_size_bounds: (min_batch, max_batch),
                per_trial_memory_mb: model_overhead_mb,
                total_reserved_mb: model_overhead_mb,
            };
        }

        let safety_margin = 0.20;
        let available_mb = self.gpu_memory_mb as f64 * (1.0 - safety_margin);

        let default_batch_mb = min_batch * per_sample_mb;
        let per_trial_mb = model_overhead_mb + default_batch_mb;

        let max_concurrent = if per_trial_mb <= 0.0 {
            1
        } else {
            let raw = (available_mb / per_trial_mb).floor() as usize;
            raw.clamp(1, 128) // Hardware cap: avoid thread explosion
        };

        let per_trial_budget = available_mb / max_concurrent as f64;
        let batch_budget = per_trial_budget - model_overhead_mb;
        let batch_upper = if batch_budget > 0.0 && per_sample_mb > 0.0 {
            (batch_budget / per_sample_mb).clamp(min_batch, max_batch)
        } else {
            min_batch
        };

        let total_reserved = per_trial_mb * max_concurrent as f64;

        HyperoptStrategy {
            max_concurrent_trials: max_concurrent,
            batch_size_bounds: (min_batch, batch_upper),
            per_trial_memory_mb: per_trial_mb,
            total_reserved_mb: total_reserved,
        }
    }

    /// Whether the detected GPU supports BF16 tensor cores (Ampere+ architecture).
    pub fn supports_bf16(&self) -> bool {
        let name = self.gpu_name.to_uppercase();
        name.contains("A100") || name.contains("A10G") || name.contains("A30")
            || name.contains("A40") || name.contains("L40") || name.contains("L4 ")
            || name.contains("H100") || name.contains("H200")
            || name.contains("RTX 30") || name.contains("RTX 40") || name.contains("RTX 50")
            || name.contains("ADA") || name.contains("HOPPER") || name.contains("BLACKWELL")
    }

    /// Whether the detected GPU supports FP16 tensor cores (Volta+ architecture).
    pub fn supports_fp16(&self) -> bool {
        self.supports_bf16() || {
            let name = self.gpu_name.to_uppercase();
            name.contains("V100") || name.contains("T4") || name.contains("RTX 20")
                || name.contains("TURING")
        }
    }

    /// Compute max hidden_dim_base that fits in VRAM for the actual DQN architecture.
    ///
    /// Accounts for: C51/QR distributional atoms, dueling streams, NoisyNet (2x params),
    /// main+target networks, Adam optimizer, gradients, activations, and GPU PER buffer.
    pub fn max_hidden_dim_base(
        &self,
        concurrent_trials: usize,
        batch_size: usize,
        state_dim: usize,
        num_actions: usize,
    ) -> usize {
        self.max_hidden_dim_base_full(
            concurrent_trials,
            batch_size,
            state_dim,
            num_actions,
            1,     // num_atoms (vanilla DQN)
            false, // noisy_nets
            false, // dueling
            0,     // replay_buffer_capacity
        )
    }

    /// Compute max hidden_dim_base with full architecture parameters.
    ///
    /// This is the accurate estimator that accounts for the real DQN architecture:
    /// C51/QR atoms, dueling value+advantage streams, NoisyNet mu+sigma doubling,
    /// GPU PER replay buffer pre-allocation, BF16 dtype on CUDA, and regime-conditional
    /// head multiplier (3 independent DQN heads for Trending/Ranging/Volatile).
    #[allow(clippy::too_many_arguments)]
    pub fn max_hidden_dim_base_full(
        &self,
        concurrent_trials: usize,
        batch_size: usize,
        state_dim: usize,
        num_actions: usize,
        num_atoms: usize,
        noisy_nets: bool,
        dueling: bool,
        replay_buffer_capacity: usize,
    ) -> usize {
        if self.gpu_memory_mb == 0 {
            return 256; // CPU fallback
        }
        let safety_margin = 0.20;
        let available_per_trial = (self.gpu_memory_mb as f64 * (1.0 - safety_margin))
            / concurrent_trials.max(1) as f64;

        // Binary search for largest base that fits in budget
        let mut lo: usize = 256;
        let mut hi: usize = 8192;
        while lo < hi {
            let mid = (lo + hi).div_ceil(2);
            let vram = Self::estimate_trial_vram_mb_full(
                mid,
                batch_size,
                state_dim,
                num_actions,
                num_atoms,
                noisy_nets,
                dueling,
                replay_buffer_capacity,
            );
            if vram <= available_per_trial {
                lo = mid;
            } else {
                hi = mid - 1;
            }
        }
        // Round down to nearest 256 (all values are already multiples of 8 for tensor cores)
        (lo / 256) * 256
    }

    /// Check whether a trial with the given parameters fits in GPU VRAM.
    ///
    /// Returns `Ok(estimated_mb)` if it fits, `Err(message)` if it exceeds budget.
    ///
    /// Uses a two-gate approach:
    /// 1. **Static estimate**: architecture-based VRAM estimate with 3× fragmentation
    ///    multiplier must be under 80% of initially detected free VRAM.
    /// 2. **Live measurement**: re-queries nvidia-smi for CURRENT free VRAM (accounts
    ///    for CUDA context already initialized, previous trial leftovers, etc.)
    ///    and checks the estimate against 80% of that.
    ///
    /// Both gates must pass. This prevents stale initial measurements from allowing
    /// trials that won't fit after CUDA context initialization.
    #[allow(clippy::too_many_arguments)]
    pub fn trial_fits_vram(
        &self,
        hidden_dim_base: usize,
        batch_size: usize,
        state_dim: usize,
        num_actions: usize,
        num_atoms: usize,
        noisy_nets: bool,
        dueling: bool,
        replay_buffer_capacity: usize,
    ) -> Result<f64, String> {
        if self.gpu_memory_mb == 0 {
            return Ok(0.0); // CPU mode, no VRAM constraint
        }
        let estimated = Self::estimate_trial_vram_mb_full(
            hidden_dim_base,
            batch_size,
            state_dim,
            num_actions,
            num_atoms,
            noisy_nets,
            dueling,
            replay_buffer_capacity,
        );

        // Gate 1: static budget from initial detection
        let static_budget = self.gpu_memory_mb as f64 * 0.80;
        if estimated > static_budget {
            return Err(format!(
                "Trial VRAM estimate {estimated:.0} MB exceeds static budget {static_budget:.0} MB \
                 (GPU: {} MB initial free). Params: hidden={hidden_dim_base} atoms={num_atoms} \
                 batch={batch_size} buffer={replay_buffer_capacity}",
                self.gpu_memory_mb
            ));
        }

        // Gate 2: live nvidia-smi measurement for current free VRAM
        // By the time this runs, CUDA context may already be initialized (eating ~300-500 MB),
        // so live free is typically less than the initial detection.
        if let Ok((_total, live_free, _name)) =
            ml_core::memory_optimization::auto_batch_size::detect_gpu_memory()
        {
            let live_budget = live_free * 0.80;
            if estimated > live_budget {
                return Err(format!(
                    "Trial VRAM estimate {estimated:.0} MB exceeds live budget {live_budget:.0} MB \
                     (GPU: {live_free:.0} MB currently free). Params: hidden={hidden_dim_base} \
                     atoms={num_atoms} batch={batch_size} buffer={replay_buffer_capacity}",
                ));
            }
        }

        Ok(estimated)
    }

    /// Estimate total VRAM for one trial (vanilla DQN, no distributional/dueling/noisy).
    ///
    /// For production DQN with C51/Dueling/NoisyNet, use `estimate_trial_vram_mb_full`.
    pub fn estimate_trial_vram_mb(
        hidden_dim_base: usize,
        batch_size: usize,
        state_dim: usize,
        num_actions: usize,
    ) -> f64 {
        Self::estimate_trial_vram_mb_full(
            hidden_dim_base,
            batch_size,
            state_dim,
            num_actions,
            1,     // vanilla DQN
            false, // no noisy nets
            false, // no dueling
            0,     // no GPU PER buffer
        )
    }

    /// Accurate VRAM estimate for the real DQN architecture.
    ///
    /// Model weights and replay buffer states use BF16 (2 bytes) on CUDA.
    /// Optimizer moments and gradients are stored in F32 (4 bytes) for
    /// numerical stability (AdamW moment accumulation in BF16 causes
    /// catastrophic underflow).  Loss computation is fully F32.
    ///
    /// The final multiplier (3×) accounts for severe memory fragmentation
    /// caused by the cudarc + Candle dual CUDA allocator architecture.
    /// nvidia-smi may report GB of "free" memory during OOM because the
    /// driver sees total free, but neither allocator can use the other's
    /// freed blocks.  Empirically validated: a trial estimating 750 MB
    /// OOM'd on a 4 GB GPU with 2.9 GB "free" per nvidia-smi.
    ///
    /// Accounts for:
    /// - Constant-width hidden layers: `[base, base]` (NOT pyramid `[b, b/2, b/4]`)
    /// - C51/QR distributional: output = num_actions × num_atoms (not just num_actions)
    /// - Dueling: separate value stream (1 × num_atoms) + advantage stream (num_actions × num_atoms)
    /// - NoisyNet: doubles parameters per layer (mu + sigma) + noise buffers
    /// - GPU PER replay buffer: pre-allocated [capacity, state_dim] × 2 (states + next_states)
    /// - Regime-conditional: 3 independent heads (Trending/Ranging/Volatile), each with
    ///   its own network weights, optimizer, gradients, and replay buffer
    /// - GPU experience collector: dynamic output buffers (n_episodes × timesteps)
    /// - cudarc weight copies: F32 duplicates of all model weights for CUDA kernels
    /// - C51 distributional loss intermediates: projected distributions, softmax, cross-entropy
    /// - Training data: Candle BF16 + cudarc F32 copies of market features/targets
    /// - PER sampling intermediates: cumsum, index_select on GPU buffer
    /// - CUDA context: driver, cuBLAS workspace, NVRTC modules, memory allocator overhead
    /// - Backward pass autograd graph: Candle retains ALL intermediate tensors from forward
    ///   through loss through backward — 20× activation memory for C51 distributional
    pub fn estimate_trial_vram_mb_full(
        hidden_dim_base: usize,
        batch_size: usize,
        state_dim: usize,
        num_actions: usize,
        num_atoms: usize,
        noisy_nets: bool,
        dueling: bool,
        replay_buffer_capacity: usize,
    ) -> f64 {
        // Actual network: constant-width [base, base] (see trainer.rs vec![b, b])
        let b = hidden_dim_base;
        let output_dim = num_actions * num_atoms.max(1);

        // Backbone: input → base → base (2 constant-width layers)
        // + bias terms: b + b
        let mut param_count = state_dim * b + b   // layer 1: weights + bias
            + b * b + b;                           // layer 2: weights + bias

        if dueling {
            // Dueling: value stream (b → value_hidden → num_atoms)
            //        + advantage stream (b → adv_hidden → output_dim)
            // value_hidden and adv_hidden default to 128 (dueling_hidden_dim)
            // but scale with hidden_dim_base in hyperopt
            let vh = (b / 4).max(128);
            param_count += b * vh + vh             // value hidden layer
                + vh * num_atoms.max(1) + num_atoms.max(1) // value output
                + b * vh + vh                       // advantage hidden layer
                + vh * output_dim + output_dim;     // advantage output
        } else {
            // Single head: b → output_dim
            param_count += b * output_dim + output_dim;
        }

        // NoisyNet: mu + sigma (2x trainable params) + noise buffers (0.5x, non-trainable)
        let noisy_multiplier: f64 = if noisy_nets { 2.5 } else { 1.0 };
        let effective_params = param_count as f64 * noisy_multiplier;
        // Trainable params only (for optimizer/gradient accounting)
        let trainable_multiplier: f64 = if noisy_nets { 2.0 } else { 1.0 };
        let trainable_params = param_count as f64 * trainable_multiplier;

        // BF16 = 2 bytes, F32 = 4 bytes
        let bf16_bytes: f64 = 2.0;
        let f32_bytes: f64 = 4.0;

        // === Per-head costs (× 3 for regime-conditional) ===

        // Weights: main + target networks (BF16 in Candle via training_dtype)
        let weights_mb = effective_params * bf16_bytes * 2.0 / 1e6;
        // Adam optimizer: 2 moments per trainable param.
        // Stored in F32 for numerical stability — BF16 moment accumulation
        // causes catastrophic underflow in the running variance (second moment).
        let optimizer_mb = trainable_params * f32_bytes * 2.0 / 1e6;
        // Gradients: F32 in Candle's autograd (backward() always produces F32
        // GradStore, even when model tensors are BF16 — casts happen in optimizer step)
        let gradients_mb = trainable_params * f32_bytes / 1e6;

        // GPU PER replay buffer: pre-allocated ring buffer tensors
        let replay_mb = if replay_buffer_capacity > 0 {
            let cap = replay_buffer_capacity as f64;
            let sd = state_dim as f64;
            cap * sd * bf16_bytes * 2.0 / 1e6 // states + next_states (BF16)
                + cap * f32_bytes * 4.0 / 1e6 // actions(U32)/rewards/dones/priorities (F32)
        } else {
            0.0
        };

        // PER sampling intermediates: cumsum, pow, index_select create temporary tensors
        // on top of the replay buffer. Worst case ~3× single column of the buffer.
        let per_sampling_mb = if replay_buffer_capacity > 0 {
            replay_buffer_capacity as f64 * f32_bytes * 3.0 / 1e6
        } else {
            0.0
        };

        // Regime-conditional: 3 independent heads
        let num_heads: f64 = 3.0;
        let per_head_mb = weights_mb + optimizer_mb + gradients_mb + replay_mb + per_sampling_mb;

        // === Shared costs (not multiplied by heads) ===

        // Activation memory: 2 hidden layers × batch_size × BF16 × 2 (forward+backward)
        let total_hidden = (b * 2) as f64;
        let activation_mb = total_hidden * batch_size as f64 * bf16_bytes * 2.0 / 1e6;

        // Backward pass autograd graph: Candle retains ALL intermediate tensors
        // from forward pass through loss computation until backward() completes.
        // For C51 distributional with 100 atoms, the Tz projection creates ~20
        // intermediate tensors ([batch, atoms] or [batch, actions, atoms]),
        // softmax creates ~5 more, cross-entropy ~3 more, plus dueling stream
        // intermediates. Empirically measured at 20× forward activation memory.
        let backward_graph_mb = activation_mb * 20.0;

        // Batch tensors: states + next_states (BF16) + actions/rewards/dones/weights (F32 scalars)
        let batch_mb = batch_size as f64 * state_dim as f64 * bf16_bytes * 2.0 / 1e6
            + batch_size as f64 * f32_bytes * 4.0 / 1e6;

        // C51 distributional loss intermediates (per training step):
        // - projected distribution: [batch, atoms] × F32
        // - target distribution: [batch, atoms] × F32
        // - online logits: [batch, actions, atoms] × F32
        // - softmax + cross-entropy intermediates: ~3× above (Tz clip, floor, ceil,
        //   offset lower, offset upper, weighted scatter, log_softmax, kl_div)
        let c51_loss_mb = if num_atoms > 1 {
            let atoms = num_atoms as f64;
            let bs = batch_size as f64;
            let na = num_actions as f64;
            (bs * atoms * 2.0 + bs * na * atoms * 5.0) * f32_bytes / 1e6
        } else {
            0.0
        };

        // GPU-resident features: experience collector + cudarc weight copies + data staging.
        // These are ONLY active when GPU PER is enabled (replay_buffer_capacity > 0 as proxy).
        // On small GPUs (≤8 GB), GPU PER and experience collector are disabled to avoid
        // cudarc/Candle dual-allocator fragmentation.
        let gpu_features_active = replay_buffer_capacity > 0;

        let collector_mb = if gpu_features_active {
            // GPU experience collector: 128 episodes × 500 timesteps = 64K output slots
            let collector_output: f64 = 64_000.0;
            let collector_per_episode: f64 = 128.0;
            collector_output * state_dim as f64 * f32_bytes / 1e6
                + collector_output * f32_bytes * 5.0 / 1e6
                + collector_per_episode * 114.0 * f32_bytes / 1e6
        } else {
            0.0
        };

        // cudarc F32 weight copies: only when GPU experience collector is active
        let cudarc_weights_mb = if gpu_features_active {
            effective_params * f32_bytes * 2.0 / 1e6
        } else {
            0.0
        };

        // Training data on GPU (Candle BF16 always present for training)
        let training_bars = if gpu_features_active {
            (replay_buffer_capacity as f64).max(50_000.0)
        } else {
            50_000.0 // minimum training data estimate
        };
        let data_candle_mb = training_bars * (42.0 + 4.0) * bf16_bytes / 1e6;
        // cudarc F32 data copies: only when GPU experience collector is active
        let data_cudarc_mb = if gpu_features_active {
            training_bars * (42.0 + 4.0) * f32_bytes / 1e6
        } else {
            0.0
        };

        // GPU buffer pool: only when GPU experience collector is active
        let buffer_pool_mb = if gpu_features_active {
            100_000.0 * 42.0 * f32_bytes / 1e6
        } else {
            0.0
        };

        // CUDA context + cuBLAS workspace. Reduced when no NVRTC/cudarc overhead.
        let cuda_overhead: f64 = if gpu_features_active { 500.0 } else { 350.0 };

        let subtotal = per_head_mb * num_heads
            + activation_mb + backward_graph_mb + batch_mb
            + c51_loss_mb
            + collector_mb + cudarc_weights_mb
            + data_candle_mb + data_cudarc_mb
            + buffer_pool_mb
            + cuda_overhead;

        // Fragmentation multiplier:
        // - 3× when dual allocators active (cudarc + Candle cannot reuse each other's freed blocks)
        // - 1.5× when only Candle allocator (still some fragmentation from autograd temporaries)
        let fragmentation = if gpu_features_active { 3.0 } else { 1.5 };
        subtotal * fragmentation
    }
}

/// Trait for models that support hyperparameter optimization
///
/// This trait defines the interface for models that can be optimized using
/// Bayesian optimization. Models must:
/// 1. Define a parameter space type implementing `ParameterSpace`
/// 2. Define a metrics type for training results
/// 3. Implement training with arbitrary parameters
/// 4. Extract a scalar objective from metrics (lower is better)
///
/// ## Type Parameters
///
/// - `Params`: Parameter space type (must implement `ParameterSpace`)
/// - `Metrics`: Training metrics type (must be Clone + Debug)
///
/// ## Example Implementation
///
/// ```rust,no_run
/// use ml::hyperopt::{HyperparameterOptimizable, ParameterSpace};
/// use ml::MLError;
///
/// #[derive(Debug, Clone)]
/// struct MyParams { learning_rate: f64 }
///
/// impl ParameterSpace for MyParams {
///     fn continuous_bounds() -> Vec<(f64, f64)> {
///         vec![(1e-5_f64.ln(), 1e-2_f64.ln())]
///     }
///     fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
///         Ok(Self { learning_rate: x[0].exp() })
///     }
///     fn to_continuous(&self) -> Vec<f64> {
///         vec![self.learning_rate.ln()]
///     }
///     fn param_names() -> Vec<&'static str> {
///         vec!["learning_rate"]
///     }
/// }
///
/// #[derive(Debug, Clone)]
/// struct MyMetrics { val_loss: f64 }
///
/// struct MyModel;
///
/// impl HyperparameterOptimizable for MyModel {
///     type Params = MyParams;
///     type Metrics = MyMetrics;
///
///     fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError> {
///         // Train model with params...
///         Ok(MyMetrics { val_loss: 0.5 })
///     }
///
///     fn extract_objective(metrics: &Self::Metrics) -> f64 {
///         metrics.val_loss
///     }
/// }
/// ```
pub trait HyperparameterOptimizable {
    /// Parameter space type (must implement `ParameterSpace`)
    type Params: ParameterSpace + Clone + Debug;

    /// Metrics type returned by training (must be Clone + Debug for result tracking)
    type Metrics: Clone + Debug;

    /// Train model with given parameters and return metrics
    ///
    /// This method should:
    /// 1. Configure the model with the provided parameters
    /// 2. Run the full training loop
    /// 3. Return comprehensive metrics including validation loss
    ///
    /// # Arguments
    ///
    /// * `params` - Hyperparameters to use for training
    ///
    /// # Returns
    ///
    /// Training metrics including validation loss for optimization
    ///
    /// # Errors
    ///
    /// Returns `MLError` if training fails due to:
    /// - Invalid parameter combinations
    /// - Out of memory conditions
    /// - Numerical instability (NaN/Inf)
    /// - CUDA errors
    fn train_with_params(&mut self, params: Self::Params) -> Result<Self::Metrics, MLError>;

    /// Extract scalar objective value from metrics (lower is better)
    ///
    /// This method defines what metric to optimize. Common choices:
    /// - Validation loss (minimize)
    /// - Negative accuracy (minimize negative = maximize accuracy)
    /// - Custom scoring function (e.g., Sharpe ratio)
    ///
    /// # Arguments
    ///
    /// * `metrics` - Training metrics from `train_with_params`
    ///
    /// # Returns
    ///
    /// Scalar objective value where **lower is better**. The optimizer
    /// will try to minimize this value across trials.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use ml::hyperopt::HyperparameterOptimizable;
    /// # struct MyModel;
    /// # #[derive(Clone, Debug)] struct MyParams;
    /// # #[derive(Clone, Debug)] struct MyMetrics { val_loss: f64, accuracy: f64 }
    /// # impl HyperparameterOptimizable for MyModel {
    /// #     type Params = MyParams;
    /// #     type Metrics = MyMetrics;
    /// #     fn train_with_params(&mut self, _: Self::Params) -> Result<Self::Metrics, ml::MLError> { unimplemented!() }
    ///     fn extract_objective(metrics: &Self::Metrics) -> f64 {
    ///         // Minimize validation loss
    ///         metrics.val_loss
    ///
    ///         // Or maximize accuracy (negate to minimize)
    ///         // -metrics.accuracy
    ///     }
    /// # }
    /// ```
    fn extract_objective(metrics: &Self::Metrics) -> f64;

    /// Extract per-trial metrics as opaque JSON for reporting.
    ///
    /// Override this to attach domain-specific metrics (Sharpe, drawdown, etc.)
    /// to each `TrialResult`.  The default returns `None` (no extra metrics).
    fn extract_metrics(_metrics: &Self::Metrics) -> Option<serde_json::Value> {
        None
    }
}

/// Trait for parameter space definition
///
/// This trait defines how to convert between parameter representations:
/// - **Continuous space** `[0, 1]^n`: Used by the optimizer
/// - **Original space**: Human-readable parameter values (e.g., learning_rate=0.001)
///
/// The continuous space allows the optimizer to:
/// - Sample uniformly across the space using Latin Hypercube
/// - Build Gaussian Process surrogates
/// - Perform efficient gradient-based acquisition optimization
///
/// ## Log-Scale Parameters
///
/// For parameters spanning multiple orders of magnitude (learning rates, weight decay),
/// use log-scale transformations in `from_continuous`/`to_continuous`:
///
/// ```rust
/// # use ml::hyperopt::ParameterSpace;
/// # use ml::MLError;
/// #[derive(Clone, Debug)]
/// struct Params { learning_rate: f64 }
///
/// impl ParameterSpace for Params {
///     fn continuous_bounds() -> Vec<(f64, f64)> {
///         vec![(1e-5_f64.ln(), 1e-2_f64.ln())] // Log-scale bounds
///     }
///
///     fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
///         Ok(Self { learning_rate: x[0].exp() }) // exp() to recover original scale
///     }
///
///     fn to_continuous(&self) -> Vec<f64> {
///         vec![self.learning_rate.ln()] // ln() for continuous space
///     }
///
///     fn param_names() -> Vec<&'static str> {
///         vec!["learning_rate"]
///     }
/// }
/// ```
pub trait ParameterSpace: Sized {
    /// Return bounds for continuous parameters: `[(min, max), ...]`
    ///
    /// Each tuple defines the valid range for one continuous parameter.
    /// These bounds are used by the optimizer to:
    /// 1. Generate initial samples via Latin Hypercube Sampling
    /// 2. Constrain Gaussian Process predictions
    /// 3. Limit acquisition function optimization
    ///
    /// # Log-Scale Parameters
    ///
    /// For parameters like learning rates (1e-5 to 1e-2), use log-scale bounds:
    /// ```rust
    /// vec![
    ///     (1e-5_f64.ln(), 1e-2_f64.ln()), // learning_rate (log scale)
    ///     (16.0, 256.0),                  // batch_size (linear scale)
    /// ]
    /// ```
    ///
    /// # Returns
    ///
    /// Vector of `(min, max)` tuples, one per parameter. Length must match
    /// the dimensionality returned by `to_continuous()`.
    ///
    /// # Panics
    ///
    /// Implementations should never panic. Invalid bounds should return
    /// default values or use type-level guarantees.
    fn continuous_bounds() -> Vec<(f64, f64)>;

    /// Convert from continuous vector `[0, 1]^n` to parameters
    ///
    /// This method reconstructs original parameter values from the continuous
    /// representation used by the optimizer. It should:
    /// 1. Denormalize continuous values to actual ranges
    /// 2. Apply inverse transformations (e.g., `exp()` for log-scale params)
    /// 3. Convert to discrete types where needed (e.g., rounding for batch_size)
    ///
    /// # Arguments
    ///
    /// * `x` - Continuous parameter vector from the optimizer
    ///
    /// # Returns
    ///
    /// Reconstructed parameters, or `MLError` if conversion fails
    ///
    /// # Errors
    ///
    /// Returns `MLError::ConfigError` if:
    /// - Vector length doesn't match expected dimensionality
    /// - Values are outside valid bounds (should be impossible with proper bounds)
    /// - Conversion produces invalid parameter combinations
    ///
    /// # Example
    ///
    /// ```rust
    /// # use ml::hyperopt::ParameterSpace;
    /// # use ml::MLError;
    /// # #[derive(Clone, Debug)]
    /// # struct Params { learning_rate: f64, batch_size: usize }
    /// # impl ParameterSpace for Params {
    /// #     fn continuous_bounds() -> Vec<(f64, f64)> { vec![] }
    ///     fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
    ///         if x.len() != 2 {
    ///             return Err(MLError::ConfigError(
    ///                 format!("Expected 2 params, got {}", x.len())
    ///             ));
    ///         }
    ///
    ///         Ok(Self {
    ///             learning_rate: x[0].exp(), // Log-scale
    ///             batch_size: x[1].round() as usize, // Discrete
    ///         })
    ///     }
    /// #     fn to_continuous(&self) -> Vec<f64> { vec![] }
    /// #     fn param_names() -> Vec<&'static str> { vec![] }
    /// # }
    /// ```
    fn from_continuous(x: &[f64]) -> Result<Self, MLError>;

    /// Convert parameters to continuous vector for the optimizer
    ///
    /// This method should be the exact inverse of `from_continuous()`.
    /// It converts parameters to the continuous representation used by the optimizer.
    ///
    /// # Returns
    ///
    /// Continuous parameter vector. Length must match `continuous_bounds()`.
    ///
    /// # Invariants
    ///
    /// The following must hold for all valid parameters:
    /// ```rust,ignore
    /// let params = MyParams { ... };
    /// let continuous = params.to_continuous();
    /// let recovered = MyParams::from_continuous(&continuous)?;
    /// assert_eq!(params, recovered); // Approximate equality for floats
    /// ```
    fn to_continuous(&self) -> Vec<f64>;

    /// Get parameter names for display and logging
    ///
    /// Returns human-readable names for each parameter dimension.
    /// Used for:
    /// - Trial result logging
    /// - Convergence plot labels
    /// - Final result reporting
    ///
    /// # Returns
    ///
    /// Vector of parameter names. Length must match `continuous_bounds()`.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use ml::hyperopt::ParameterSpace;
    /// # use ml::MLError;
    /// # #[derive(Clone, Debug)]
    /// # struct Params { learning_rate: f64, batch_size: usize }
    /// # impl ParameterSpace for Params {
    /// #     fn continuous_bounds() -> Vec<(f64, f64)> { vec![] }
    /// #     fn from_continuous(_: &[f64]) -> Result<Self, MLError> { unimplemented!() }
    /// #     fn to_continuous(&self) -> Vec<f64> { vec![] }
    ///     fn param_names() -> Vec<&'static str> {
    ///         vec!["learning_rate", "batch_size"]
    ///     }
    /// # }
    /// ```
    fn param_names() -> Vec<&'static str>;

    /// Return hardware-aware bounds for continuous parameters.
    ///
    /// Override this to scale batch_size (or other memory-sensitive params)
    /// based on detected GPU VRAM. The default implementation delegates to
    /// `continuous_bounds()`, preserving existing static behavior.
    ///
    /// Adapters that override this should call `Self::continuous_bounds()`
    /// first, then replace the batch_size bound using `budget.max_batch_size()`.
    fn continuous_bounds_for(_budget: &HardwareBudget) -> Vec<(f64, f64)> {
        Self::continuous_bounds()
    }

    /// Estimated fixed VRAM overhead per model instance in MB (weights + optimizer + gradients + activations).
    /// Override for accurate VRAM-based concurrency planning.
    fn model_overhead_mb() -> f64 {
        500.0
    }

    /// Estimated VRAM per batch sample in MB.
    fn per_sample_mb() -> f64 {
        0.02
    }
}

/// Result of a single optimization trial
///
/// Contains parameter values, objective score, and timing information
/// for one training run during optimization.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrialResult<P> {
    /// Trial number (1-indexed)
    pub trial_num: usize,
    /// Parameters used for this trial
    pub params: P,
    /// Objective value achieved (lower is better)
    pub objective: f64,
    /// Wall-clock time for training (seconds)
    pub duration_secs: f64,
    /// Optional domain-specific metrics (Sharpe, drawdown, etc.)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metrics: Option<serde_json::Value>,
}

/// Complete optimization result with best parameters and trial history
///
/// Contains all information from a hyperparameter optimization run:
/// - Best parameters found
/// - Best objective value
/// - Full trial history for analysis
/// - Convergence plot data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizationResult<P> {
    /// Best parameters found during optimization
    pub best_params: P,
    /// Best objective value (lower is better)
    pub best_objective: f64,
    /// All trial results in execution order
    pub all_trials: Vec<TrialResult<P>>,
    /// Convergence plot data: `(trial_num, best_objective_so_far)`
    ///
    /// This tracks the best objective seen up to each trial, showing
    /// convergence behavior over time. Can be used to:
    /// - Plot optimization progress
    /// - Detect early convergence
    /// - Identify exploration vs exploitation phases
    pub convergence_plot_data: Vec<(usize, f64)>,
}

impl<P: Clone> OptimizationResult<P> {
    /// Create new optimization result from trial history
    ///
    /// Automatically computes best parameters and convergence data.
    ///
    /// # Arguments
    ///
    /// * `trials` - Complete trial history
    ///
    /// # Panics
    ///
    /// Panics if `trials` is empty (should never happen in practice)
    #[allow(clippy::indexing_slicing)] // assert! guarantees non-empty; min_by always returns Some
    pub fn from_trials(trials: Vec<TrialResult<P>>) -> Self {
        assert!(!trials.is_empty(), "Cannot create result from empty trials");

        // Find best trial (minimum objective)
        // assert! guarantees non-empty, so min_by always returns Some
        let best_trial = trials
            .iter()
            .min_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap_or(std::cmp::Ordering::Equal))
            .unwrap_or(&trials[0]);

        // Build convergence plot data
        let mut best_so_far = f64::INFINITY;
        let convergence_plot_data: Vec<_> = trials
            .iter()
            .map(|trial| {
                if trial.objective < best_so_far {
                    best_so_far = trial.objective;
                }
                (trial.trial_num, best_so_far)
            })
            .collect();

        Self {
            best_params: best_trial.params.clone(),
            best_objective: best_trial.objective,
            all_trials: trials,
            convergence_plot_data,
        }
    }

    /// Get the trial with the best objective.
    ///
    /// Returns `None` if `all_trials` is empty (should not happen for well-constructed results).
    pub fn best_trial(&self) -> Option<&TrialResult<P>> {
        self.all_trials
            .iter()
            .min_by(|a, b| a.objective.partial_cmp(&b.objective).unwrap_or(std::cmp::Ordering::Equal))
    }

    /// Get improvement from first to best trial.
    ///
    /// Returns `0.0` if there are no trials.
    pub fn total_improvement(&self) -> f64 {
        match self.all_trials.first() {
            Some(first) => self.best_objective - first.objective,
            None => 0.0,
        }
    }

    /// Get improvement percentage.
    ///
    /// Returns `0.0` if there are no trials or the first trial objective is near zero.
    pub fn improvement_percentage(&self) -> f64 {
        match self.all_trials.first() {
            Some(first) if first.objective.abs() >= 1e-10 => {
                ((first.objective - self.best_objective) / first.objective.abs()) * 100.0
            }
            _ => 0.0,
        }
    }
}

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

    #[derive(Debug, Clone, PartialEq)]
    struct TestParams {
        learning_rate: f64,
        batch_size: usize,
    }

    impl ParameterSpace for TestParams {
        fn continuous_bounds() -> Vec<(f64, f64)> {
            vec![
                (1e-5_f64.ln(), 1e-2_f64.ln()), // lr (log scale)
                (16.0, 256.0),                  // batch_size
            ]
        }

        fn from_continuous(x: &[f64]) -> Result<Self, MLError> {
            if x.len() != 2 {
                return Err(MLError::ConfigError(format!("Expected 2 params, got {}", x.len())));
            }

            Ok(Self {
                learning_rate: x[0].exp(),
                batch_size: x[1].round() as usize,
            })
        }

        fn to_continuous(&self) -> Vec<f64> {
            vec![self.learning_rate.ln(), self.batch_size as f64]
        }

        fn param_names() -> Vec<&'static str> {
            vec!["learning_rate", "batch_size"]
        }
    }

    #[test]
    fn test_parameter_space_roundtrip() {
        let params = TestParams {
            learning_rate: 0.001,
            batch_size: 64,
        };

        let continuous = params.to_continuous();
        let recovered = TestParams::from_continuous(&continuous).unwrap();

        assert!((recovered.learning_rate - params.learning_rate).abs() < 1e-10);
        assert_eq!(recovered.batch_size, params.batch_size);
    }

    #[test]
    fn test_optimization_result_convergence() {
        let trials = vec![
            TrialResult {
                trial_num: 1,
                params: TestParams {
                    learning_rate: 0.01,
                    batch_size: 32,
                },
                objective: 1.5,
                duration_secs: 10.0,
                metrics: None,
            },
            TrialResult {
                trial_num: 2,
                params: TestParams {
                    learning_rate: 0.001,
                    batch_size: 64,
                },
                objective: 0.8,
                duration_secs: 12.0,
                metrics: None,
            },
            TrialResult {
                trial_num: 3,
                params: TestParams {
                    learning_rate: 0.005,
                    batch_size: 128,
                },
                objective: 0.5,
                duration_secs: 15.0,
                metrics: None,
            },
        ];

        let result = OptimizationResult::from_trials(trials);

        assert_eq!(result.best_objective, 0.5);
        assert_eq!(result.best_params.batch_size, 128);
        assert_eq!(result.convergence_plot_data.len(), 3);
        assert_eq!(result.convergence_plot_data[0].1, 1.5);
        assert_eq!(result.convergence_plot_data[1].1, 0.8);
        assert_eq!(result.convergence_plot_data[2].1, 0.5);
    }

    #[test]
    fn test_hardware_budget_cpu_only() {
        let budget = HardwareBudget::cpu_only();
        assert_eq!(budget.gpu_memory_mb, 0);
        assert!(budget.max_batch_size(100.0, 0.05, 16.0, 4096.0).is_none());
    }

    #[test]
    fn test_hardware_budget_with_memory() {
        let budget = HardwareBudget::with_memory_mb(4096); // 4GB like RTX 3050 Ti
        assert_eq!(budget.gpu_memory_mb, 4096);

        // 4096 * 0.70 = 2867.2; 2867.2 - 150 overhead = 2717.2; 2717.2 / 0.05 = 54344
        // clamped to max_batch=4096
        let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
        assert!((max - 4096.0).abs() < 0.01);
    }

    #[test]
    fn test_hardware_budget_h100_wide_range() {
        let budget = HardwareBudget::with_memory_mb(81920); // 80GB H100
        // 81920 * 0.85 = 69632; 69632 - 150 = 69482; 69482 / 0.05 = 1_389_640
        // clamped to 4096
        let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
        assert!((max - 4096.0).abs() < 0.01);
    }

    #[test]
    fn test_hardware_budget_small_gpu() {
        let budget = HardwareBudget::with_memory_mb(2048); // 2GB GPU
        // 2048 * 0.70 = 1433.6; 1433.6 - 150 = 1283.6; 1283.6 / 0.05 = 25672
        // clamped to 4096
        let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
        assert!((max - 4096.0).abs() < 0.01);

        // With much heavier model overhead
        // 1433.6 - 1500 = -66.4 → available <= 0 → returns min_batch
        let max2 = budget.max_batch_size(1500.0, 0.05, 16.0, 4096.0).unwrap();
        assert!((max2 - 16.0).abs() < 0.01);
    }

    #[test]
    fn test_hardware_budget_clamps_to_min() {
        let budget = HardwareBudget::with_memory_mb(256); // Tiny GPU
        // 256 * 0.70 = 179.2; 179.2 - 200 = -20.8 → available <= 0 → returns min_batch
        let max = budget.max_batch_size(200.0, 0.05, 16.0, 4096.0).unwrap();
        assert!((max - 16.0).abs() < 0.01);
    }

    #[test]
    fn test_continuous_bounds_for_default_delegates() {
        // Default impl should return same as continuous_bounds()
        let budget = HardwareBudget::with_memory_mb(81920);
        let static_bounds = TestParams::continuous_bounds();
        let dynamic_bounds = TestParams::continuous_bounds_for(&budget);
        assert_eq!(static_bounds, dynamic_bounds);
    }

    #[test]
    fn test_plan_hyperopt_tiny_model_high_concurrency() {
        let budget = HardwareBudget { gpu_memory_mb: 24576, gpu_name: String::new() }; // L4 24GB
        // L4: 24576 * 0.80 = 19660MB available, per_trial = 300 + 64*0.0002 = ~300MB
        // raw = 19660 / 300 = ~65, clamped to [1, 128]
        let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0);
        assert!(strategy.max_concurrent_trials >= 10,
            "L4 should support 10+ concurrent DQN trials, got {}", strategy.max_concurrent_trials);
        assert!(strategy.max_concurrent_trials <= 128,
            "concurrency should be capped at hardware limit 128, got {}", strategy.max_concurrent_trials);
    }

    #[test]
    fn test_plan_hyperopt_large_model_low_concurrency() {
        let budget = HardwareBudget { gpu_memory_mb: 24576, gpu_name: String::new() }; // L4 24GB
        let strategy = budget.plan_hyperopt(2000.0, 0.05, 64.0, 4096.0);
        assert!(strategy.max_concurrent_trials >= 1,
            "should always allow at least 1 trial");
        assert!(strategy.max_concurrent_trials <= 9,
            "L4 should not fit 10+ TFT trials, got {}", strategy.max_concurrent_trials);
    }

    #[test]
    fn test_plan_hyperopt_cpu_only_sequential() {
        let budget = HardwareBudget::cpu_only();
        let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0);
        assert_eq!(strategy.max_concurrent_trials, 1,
            "CPU-only should fall back to sequential");
    }

    #[test]
    fn test_plan_hyperopt_h100_maxes_out() {
        let budget = HardwareBudget { gpu_memory_mb: 81920, gpu_name: String::new() }; // H100 80GB
        // H100: 81920 * 0.80 = 65536MB available, per_trial ~300MB
        // raw = 65536 / 300 = ~218, clamped to 128
        let strategy = budget.plan_hyperopt(300.0, 0.0002, 64.0, 4096.0);
        assert_eq!(strategy.max_concurrent_trials, 128,
            "H100 should max out at hardware cap 128 for DQN");
    }
}
