//! Auto Batch Size Tuning for GPU Training
//!
//! Automatically detects available GPU memory and calculates optimal batch size
//! to prevent OOM errors while maximizing GPU utilization.
//!
//! ## Memory Budget Calculation
//!
//! Total GPU memory is allocated as follows:
//! - Model parameters: ~M bytes (model-specific)
//! - Forward pass activations: ~M bytes (same as model size)
//! - Backward pass gradients: ~M bytes (same as model size)
//! - Optimizer states (Adam): ~2M bytes (momentum + variance)
//! - Batch data: batch_size × sequence_length × feature_dim × 4 bytes
//! - Safety margin: 20% of total memory
//!
//! ## Usage
//!
//! ```rust,ignore
//! use ml::memory_optimization::auto_batch_size::{AutoBatchSizer, BatchSizeConfig};
//!
//! // Create auto batch sizer
//! let sizer = AutoBatchSizer::new()?;
//!
//! // Configure model memory requirements
//! let config = BatchSizeConfig {
//!     model_memory_mb: 125.0,  // TFT model size
//!     sequence_length: 60,     // Lookback window
//!     feature_dim: 225,        // 225 features (Wave C + Wave D)
//!     gradient_checkpointing: false,
//!     optimizer_type: OptimizerType::Adam,
//!     safety_margin: 0.20,     // 20% safety margin
//! };
//!
//! // Get optimal batch size
//! let batch_size = sizer.calculate_optimal_batch_size(&config)?;
//! println!("Optimal batch size: {}", batch_size);
//! ```

use serde::{Deserialize, Serialize};
use std::process::Command;
use tracing::{debug, info, warn};

use crate::{MLError, MLResult};

/// Optimizer type for memory calculation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum OptimizerType {
    /// SGD with momentum (1x model memory for momentum)
    SGD,
    /// Adam optimizer (2x model memory for momentum + variance)
    Adam,
    /// AdamW optimizer (2x model memory for momentum + variance)
    AdamW,
}

impl OptimizerType {
    /// Get memory multiplier for optimizer states
    pub fn memory_multiplier(&self) -> f64 {
        match self {
            OptimizerType::SGD => 1.0,   // Only momentum
            OptimizerType::Adam => 2.0,  // Momentum + variance
            OptimizerType::AdamW => 2.0, // Momentum + variance
        }
    }
}

/// Model precision for memory calculation
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelPrecision {
    /// FP32 (32-bit floating point) - 4 bytes per parameter
    FP32,
    /// BF16 (bfloat16) - 2 bytes per parameter (2x smaller than FP32)
    BF16,
}

impl ModelPrecision {
    /// Get bytes per parameter for this precision
    pub fn bytes_per_param(&self) -> f64 {
        match self {
            ModelPrecision::FP32 => 4.0,
            ModelPrecision::BF16 => 2.0,
        }
    }

    /// Get memory multiplier relative to BF16 (1x for BF16, 2x for FP32)
    pub fn memory_multiplier(&self) -> f64 {
        match self {
            ModelPrecision::FP32 => 2.0,
            ModelPrecision::BF16 => 1.0,
        }
    }
}

/// Batch size configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchSizeConfig {
    /// Base model memory in MB (scaled by precision multiplier for actual usage).
    /// For TFT-225: ~250MB (BF16) → ~500MB (FP32).
    pub model_memory_mb: f64,

    /// Model precision (FP32 or BF16)
    pub model_precision: ModelPrecision,

    /// Sequence length (lookback window)
    pub sequence_length: usize,

    /// Feature dimension (number of input features)
    pub feature_dim: usize,

    /// Enable gradient checkpointing (reduces activation memory by ~35%)
    pub gradient_checkpointing: bool,

    /// Optimizer type (affects memory overhead)
    pub optimizer_type: OptimizerType,

    /// Safety margin (0.0-1.0, recommended: 0.20 for 20%)
    pub safety_margin: f64,

    /// Minimum batch size (default: 1)
    pub min_batch_size: usize,

    /// Maximum batch size (default: 256)
    pub max_batch_size: usize,
}

impl Default for BatchSizeConfig {
    fn default() -> Self {
        Self {
            model_memory_mb: 125.0,
            model_precision: ModelPrecision::BF16,
            sequence_length: 60,
            feature_dim: 225,
            gradient_checkpointing: false,
            optimizer_type: OptimizerType::Adam,
            safety_margin: 0.20,
            min_batch_size: 1,
            max_batch_size: 256,
        }
    }
}

/// Auto batch size tuner
#[derive(Debug)]
pub struct AutoBatchSizer {
    /// Total GPU memory in MB
    total_memory_mb: f64,

    /// Free GPU memory in MB
    free_memory_mb: f64,

    /// GPU device name
    device_name: String,
}

impl AutoBatchSizer {
    /// Create new auto batch sizer with GPU memory detection
    pub fn new() -> MLResult<Self> {
        let (total_mb, free_mb, name) = detect_gpu_memory()?;

        info!(
            "GPU detected: {} (Total: {:.1} MB, Free: {:.1} MB)",
            name, total_mb, free_mb
        );

        Ok(Self {
            total_memory_mb: total_mb,
            free_memory_mb: free_mb,
            device_name: name,
        })
    }

    /// Create auto batch sizer with manual memory specification (for testing)
    pub fn with_manual_memory(total_mb: f64, free_mb: f64, device_name: String) -> Self {
        Self {
            total_memory_mb: total_mb,
            free_memory_mb: free_mb,
            device_name,
        }
    }

    /// Calculate optimal batch size based on available GPU memory
    pub fn calculate_optimal_batch_size(&self, config: &BatchSizeConfig) -> MLResult<usize> {
        // Calculate memory budget with precision-aware safety margin
        // Safety margins account for:
        // - FP32: 25% margin (CUDA allocator overhead + minor fragmentation)
        //         Batch overhead (250MB) already includes activation buffers
        //         Gradient checkpointing (35% discount) already reduces activation memory
        // - BF16: 20% margin (half-precision models have more predictable memory)
        let precision_safety_margin: f64 = match config.model_precision {
            ModelPrecision::FP32 => 0.25, // 25% safety margin (overhead already in batch_overhead_mb)
            ModelPrecision::BF16 => 0.20, // 20% safety margin for BF16 models
        };

        // Apply whichever safety margin is more conservative (larger)
        let effective_safety_margin = precision_safety_margin.max(config.safety_margin);
        let usable_memory_mb = self.free_memory_mb * (1.0 - effective_safety_margin);

        debug!(
            "Memory budget calculation: Free={:.1}MB, Safety margin={:.1}% (precision: {:.1}%, config: {:.1}%), Usable={:.1}MB",
            self.free_memory_mb,
            effective_safety_margin * 100.0,
            precision_safety_margin * 100.0,
            config.safety_margin * 100.0,
            usable_memory_mb
        );

        // Scale base memory by precision multiplier
        let model_mb = config.model_memory_mb * config.model_precision.memory_multiplier();

        // Calculate fixed memory overhead (model + optimizer states)
        let optimizer_mb = model_mb * config.optimizer_type.memory_multiplier();
        let gradient_mb = model_mb; // Gradients = model size

        // Gradient checkpointing reduces activation memory by 30-40% in practice
        // (not 50% as theoretical - some layers still need full activations)
        let activation_multiplier = if config.gradient_checkpointing {
            0.65 // Gradient checkpointing reduces activation memory by ~35%
        } else {
            1.0
        };
        let activation_mb = model_mb * activation_multiplier;

        let fixed_overhead_mb = model_mb + optimizer_mb + gradient_mb + activation_mb;

        // Add batch-level overhead (calculated before error checking)
        // This accounts for intermediate buffers that don't scale linearly with batch size
        // Batch overhead scales with GPU VRAM (attention workspaces, cuDNN scratch)
        // Base: ~1.5% of total VRAM for FP32, with per-precision multipliers
        let base_overhead_pct = match config.model_precision {
            ModelPrecision::FP32 => 0.015,
            ModelPrecision::BF16 => 0.005,
        };
        let batch_overhead_mb = (self.total_memory_mb * base_overhead_pct).max(50.0);

        debug!(
            "Fixed overhead: Model={:.1}MB, Optimizer={:.1}MB, Gradients={:.1}MB, Activations={:.1}MB, Total={:.1}MB, Batch overhead={:.1}MB",
            model_mb, optimizer_mb, gradient_mb, activation_mb, fixed_overhead_mb, batch_overhead_mb
        );

        // Check if we have enough memory for minimum batch size
        let total_overhead_mb = fixed_overhead_mb + batch_overhead_mb;
        if usable_memory_mb < total_overhead_mb {
            return Err(MLError::ConfigError(format!(
                    "Insufficient GPU memory: {:.1}MB available, {:.1}MB required (model: {:.1}MB + batch overhead: {:.1}MB). Consider enabling gradient_checkpointing or using a larger GPU.",
                    usable_memory_mb, total_overhead_mb, fixed_overhead_mb, batch_overhead_mb
                )));
        }

        // Calculate available memory for batch data (after fixed overhead + batch overhead)
        let available_for_batches_mb = usable_memory_mb - fixed_overhead_mb - batch_overhead_mb;

        // Calculate memory per sample (sequence_length × feature_dim × bytes_per_param)
        // Account for:
        // - Input data: sequence_length × feature_dim × bytes_per_param
        // - Target data: ~20% of input size
        let bytes_per_param = config.model_precision.bytes_per_param();
        let bytes_per_sample =
            (config.sequence_length * config.feature_dim) as f64 * bytes_per_param;
        let bytes_per_target = bytes_per_sample * 0.2;
        let total_bytes_per_sample = bytes_per_sample + bytes_per_target;
        let mb_per_sample = total_bytes_per_sample / (1024.0 * 1024.0);

        debug!(
            "Memory per sample: {:.3}MB (input: {:.1} bytes, target: {:.1} bytes), Batch overhead: {:.1}MB",
            mb_per_sample, bytes_per_sample, bytes_per_target, batch_overhead_mb
        );

        // Calculate maximum batch size
        let max_batch_size_from_memory = if available_for_batches_mb <= 0.0 {
            0
        } else {
            (available_for_batches_mb / mb_per_sample).floor() as usize
        };

        // Clamp to configured min/max
        let optimal_batch_size = max_batch_size_from_memory
            .max(config.min_batch_size)
            .min(config.max_batch_size);

        // Round down to nearest power of 2 for GPU alignment
        let batch_size = if optimal_batch_size == 0 {
            0
        } else {
            1_usize << (usize::BITS - 1 - optimal_batch_size.leading_zeros())
        };
        let final_batch_size = batch_size
            .max(config.min_batch_size)
            .min(config.max_batch_size);

        info!(
            "Optimal batch size calculated: {} (memory-based: {}, rounded: {}, final: {})",
            final_batch_size, max_batch_size_from_memory, batch_size, final_batch_size
        );

        // Estimate actual memory usage
        let estimated_usage_mb = fixed_overhead_mb + (final_batch_size as f64 * mb_per_sample);
        let memory_utilization = (estimated_usage_mb / usable_memory_mb) * 100.0;

        info!(
            "Estimated memory usage: {:.1}MB / {:.1}MB ({:.1}% utilization)",
            estimated_usage_mb, usable_memory_mb, memory_utilization
        );

        // Warn if utilization is very low
        if memory_utilization < 50.0 {
            info!(
                "Low GPU memory utilization ({:.1}%). Consider increasing max_batch_size or model size.",
                memory_utilization
            );
        }

        Ok(final_batch_size)
    }

    /// Calculate the maximum safe batch size that fits in GPU VRAM.
    ///
    /// Unlike `calculate_optimal_batch_size()` which rounds to power-of-2 and
    /// clamps to `config.max_batch_size`, this returns the raw VRAM-derived
    /// ceiling. Use this when an external optimizer (PSO, hyperopt) controls
    /// the batch size and you only need a safety cap.
    pub fn max_safe_batch_size(&self, config: &BatchSizeConfig) -> usize {
        let effective_safety_margin = match config.model_precision {
            ModelPrecision::FP32 => 0.25_f64,
            ModelPrecision::BF16 => 0.20,
        }
        .max(config.safety_margin);
        let usable_memory_mb = self.free_memory_mb * (1.0 - effective_safety_margin);

        let model_mb = config.model_memory_mb * config.model_precision.memory_multiplier();

        let optimizer_mb = model_mb * config.optimizer_type.memory_multiplier();
        let gradient_mb = model_mb;
        let activation_multiplier = if config.gradient_checkpointing {
            0.65
        } else {
            1.0
        };
        let activation_mb = model_mb * activation_multiplier;
        let fixed_overhead_mb = model_mb + optimizer_mb + gradient_mb + activation_mb;
        // Batch overhead scales with GPU VRAM (attention workspaces, cuDNN scratch)
        let base_overhead_pct = match config.model_precision {
            ModelPrecision::FP32 => 0.015,
            ModelPrecision::BF16 => 0.005,
        };
        let batch_overhead_mb = (self.total_memory_mb * base_overhead_pct).max(50.0);

        let available = usable_memory_mb - fixed_overhead_mb - batch_overhead_mb;
        if available <= 0.0 {
            return config.min_batch_size;
        }

        let bytes_per_param = config.model_precision.bytes_per_param();
        let bytes_per_sample =
            (config.sequence_length * config.feature_dim) as f64 * bytes_per_param;
        let total_bytes_per_sample = bytes_per_sample * 1.2; // input + 20% target
        let mb_per_sample = total_bytes_per_sample / (1024.0 * 1024.0);

        if mb_per_sample <= 0.0 {
            return config.min_batch_size;
        }

        let max_from_memory = (available / mb_per_sample).floor() as usize;
        max_from_memory.max(config.min_batch_size)
    }

    /// Get GPU memory info
    pub fn memory_info(&self) -> GpuMemoryInfo {
        GpuMemoryInfo {
            device_name: self.device_name.clone(),
            total_memory_mb: self.total_memory_mb,
            free_memory_mb: self.free_memory_mb,
            used_memory_mb: self.total_memory_mb - self.free_memory_mb,
        }
    }

    /// Reduce batch size for OOM recovery (exponential backoff)
    ///
    /// # Arguments
    /// * `current_batch_size` - Current batch size that caused OOM
    ///
    /// # Returns
    /// * New batch size (halved, minimum 1)
    ///
    /// # Example
    /// ```rust,ignore
    /// let new_batch_size = AutoBatchSizer::reduce_batch_size(32); // Returns 16
    /// let min_batch_size = AutoBatchSizer::reduce_batch_size(1);  // Returns 1
    /// ```
    pub fn reduce_batch_size(current_batch_size: usize) -> usize {
        (current_batch_size / 2).max(1)
    }

    /// Check if batch size is below minimum viable threshold
    ///
    /// # Arguments
    /// * `batch_size` - Batch size to check
    ///
    /// # Returns
    /// * true if batch size is too small for efficient training
    ///
    /// # Threshold
    /// * Below 4: Training is inefficient (GPU underutilized)
    pub fn is_batch_size_too_small(batch_size: usize) -> bool {
        batch_size < 4
    }
}

/// GPU memory information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuMemoryInfo {
    pub device_name: String,
    pub total_memory_mb: f64,
    pub free_memory_mb: f64,
    pub used_memory_mb: f64,
}

/// Bundled output from [`GpuHardwareInfo::optimal_replay_config`].
///
/// Ensures callers always receive both the capacity **and** the PER memory
/// budget together, preventing the temporal-coupling bug where the auto-sizer
/// proposes a capacity that the PER pre-flight check rejects.
#[derive(Debug, Clone, Copy)]
pub struct OptimalReplayConfig {
    /// VRAM-optimal replay buffer capacity (clamped 100K..10M).
    pub capacity: usize,
    /// Maximum PER buffer allocation in bytes (20% of total VRAM, min 1 GB).
    pub per_max_buffer_bytes: usize,
}

/// Full GPU hardware info for dynamic scaling decisions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuHardwareInfo {
    pub device_name: String,
    pub total_memory_mb: f64,
    pub free_memory_mb: f64,
    /// Number of Streaming Multiprocessors. 0 if detection failed.
    pub sm_count: usize,
}

impl GpuHardwareInfo {
    /// Compute the optimal number of parallel episodes for the GPU
    /// experience collector kernel.
    ///
    /// Each CUDA thread runs one complete episode. To saturate the GPU we
    /// want enough threads to keep all SMs busy with multiple warps.
    /// With high register pressure per thread, target ~2 warps/SM (64 threads).
    ///
    /// The result is capped by VRAM: each episode needs ~110 KB of output
    /// buffers (500 timesteps × (48 state + 5 scalar) × 4 bytes).
    ///
    /// Falls back to a VRAM-proportional heuristic when SM count is unknown.
    pub fn optimal_n_episodes(&self, state_dim: usize, timesteps: usize) -> usize {
        const WARPS_PER_SM: usize = 2;  // conservative for high-register kernels
        const THREADS_PER_WARP: usize = 32;
        const MIN_EPISODES: usize = 128;
        const MAX_EPISODES: usize = 4096;

        // Per-episode output buffer: states + actions + rewards + done + target_q + td_error
        // Plus per-episode state: portfolio(6) + barrier(10) + diversity(30+2) + rng(1) ≈ 50 floats
        let per_episode_bytes = timesteps * (state_dim + 5) * 4 + 50 * 4;

        // VRAM budget for experience collector output buffers.
        // `free_memory_mb` is measured at runtime (after model/optimizer are loaded).
        // Reserve headroom for:
        //   ~200 MB backward graph (scales with batch size during training step)
        //   ~150 MB candle block allocator fragmentation + temporary tensors
        //   ~50 MB training data upload (market features buffer, re-uploaded per epoch)
        // Then use 50% of remaining for collector output buffers.
        let headroom_mb = 400.0_f64;
        let available_mb = (self.free_memory_mb - headroom_mb).max(64.0);
        let vram_budget_bytes = (available_mb * 1024.0 * 1024.0 * 0.50) as usize;
        let vram_cap = if per_episode_bytes > 0 {
            vram_budget_bytes / per_episode_bytes
        } else {
            MAX_EPISODES
        };

        // SM-based target
        let sm_target = if self.sm_count > 0 {
            self.sm_count * WARPS_PER_SM * THREADS_PER_WARP
        } else {
            // Fallback: VRAM-proportional heuristic (128 episodes per GB)
            let vram_gb = self.total_memory_mb / 1024.0;
            (vram_gb * 128.0) as usize
        };

        let optimal = sm_target.min(vram_cap).clamp(MIN_EPISODES, MAX_EPISODES);
        // Round up to nearest multiple of 256 for clean block scheduling.
        // The VRAM budget cap above already scales dynamically with free memory,
        // so no additional hard-cap is needed for small GPUs.
        let aligned = ((optimal + 255) / 256) * 256;
        aligned.min(MAX_EPISODES)
    }

    /// Calculate VRAM-optimal replay buffer capacity and PER memory budget.
    ///
    /// Allocates `vram_fraction` of free GPU memory to the replay buffer.
    /// Each transition stores: 2 × state_dim × f32 (state + next_state)
    ///   + 4 scalars (action, reward, done, priority) = (2 * state_dim + 4) * 4 bytes.
    ///
    /// Enforces guardrails:
    ///   - MIN_REPLAY_CAPACITY = 100_000 (below this, sample efficiency degrades)
    ///   - MAX_REPLAY_CAPACITY = 10_000_000 (above this, priority staleness dominates)
    ///   - PER memory budget cap: capacity must fit within `per_max_buffer_bytes()` so
    ///     the GPU PER pre-flight check never rejects the auto-sizer's output.
    ///
    /// Returns an [`OptimalReplayConfig`] bundling the capacity with the PER
    /// memory budget so callers cannot forget to pass the budget downstream.
    pub fn optimal_replay_config(&self, state_dim: usize, vram_fraction: f64) -> OptimalReplayConfig {
        const MIN_REPLAY_CAPACITY: usize = 100_000;
        const MAX_REPLAY_CAPACITY: usize = 10_000_000;

        let bytes_per_transition = (2 * state_dim + 4) * 4; // f32 each
        let budget_bytes = (self.free_memory_mb * vram_fraction * 1024.0 * 1024.0) as usize;
        let raw_capacity = budget_bytes / bytes_per_transition.max(1);

        // Cap by PER buffer VRAM budget (proportional to total VRAM).
        // Uses the same bytes-per-transition formula as GpuReplayBuffer's pre-flight
        // check, so the auto-sizer never proposes a capacity the PER would reject.
        let per_max = self.per_max_buffer_bytes();
        let per_cap = per_max / bytes_per_transition.max(1);

        let capacity = raw_capacity.min(per_cap).clamp(MIN_REPLAY_CAPACITY, MAX_REPLAY_CAPACITY);
        OptimalReplayConfig { capacity, per_max_buffer_bytes: per_max }
    }

    /// Maximum PER replay buffer allocation (bytes) proportional to total VRAM.
    ///
    /// Formula: 20% of total GPU memory (min 1 GB).
    /// No upper cap — scales with available hardware.
    pub fn per_max_buffer_bytes(&self) -> usize {
        let max_mb = (self.total_memory_mb * 0.20).max(1024.0);
        (max_mb * 1024.0 * 1024.0) as usize
    }
}

/// Detect available GPU memory using nvidia-smi
///
/// Returns (total_mb, free_mb, device_name)
pub fn detect_gpu_memory() -> MLResult<(f64, f64, String)> {
    let info = detect_gpu_hardware()?;
    Ok((info.total_memory_mb, info.free_memory_mb, info.device_name))
}

/// Detect full GPU hardware properties (memory + SM count) via nvidia-smi.
///
/// SM count enables dynamic scaling of CUDA kernel parallelism.
/// Falls back gracefully: SM count = 0 when detection fails.
pub fn detect_gpu_hardware() -> MLResult<GpuHardwareInfo> {
    match Command::new("nvidia-smi")
        .args([
            "--query-gpu=memory.total,memory.free,name",
            "--format=csv,noheader,nounits",
        ])
        .output()
    {
        Ok(output) if output.status.success() => {
            let stdout = String::from_utf8_lossy(&output.stdout);
            let parts: Vec<&str> = stdout.trim().split(',').collect();

            if parts.len() >= 3 {
                let total_mb =
                    parts[0]
                        .trim()
                        .parse::<f64>()
                        .map_err(|e| MLError::ConfigError(format!("Failed to parse total memory: {}", e)))?;
                let free_mb = parts[1]
                    .trim()
                    .parse::<f64>()
                    .map_err(|e| MLError::ConfigError(format!("Failed to parse free memory: {}", e)))?;
                let device_name = parts[2].trim().to_string();

                // Resolve SM count from device name (nvidia-smi doesn't expose it directly)
                let sm_count = detect_sm_count_from_name(&device_name);

                return Ok(GpuHardwareInfo {
                    device_name,
                    total_memory_mb: total_mb,
                    free_memory_mb: free_mb,
                    sm_count,
                });
            }

            Err(MLError::ConfigError(format!("Unexpected nvidia-smi output format: {}", stdout)))
        },
        Ok(output) => {
            let stderr = String::from_utf8_lossy(&output.stderr);
            Err(MLError::ConfigError(format!("nvidia-smi failed: {}", stderr)))
        },
        Err(e) => {
            warn!("nvidia-smi not available ({}), using CPU fallback", e);
            Ok(GpuHardwareInfo {
                device_name: "CPU".to_owned(),
                total_memory_mb: 0.0,
                free_memory_mb: 0.0,
                sm_count: 0,
            })
        },
    }
}

/// Detect SM count by looking up the GPU device name.
///
/// nvidia-smi's `--query-gpu` doesn't expose SM count directly,
/// so we resolve it from the device name returned by the memory query.
fn detect_sm_count_from_name(device_name: &str) -> usize {
    let count = sm_count_from_device_name(device_name);
    if count > 0 {
        debug!("SM count for '{}': {}", device_name, count);
    } else {
        debug!("Unknown GPU '{}', SM count = 0 (will use VRAM heuristic)", device_name);
    }
    count
}

/// Look up SM count from known GPU device names.
///
/// This avoids requiring CUDA runtime or deviceQuery binary.
/// Returns 0 for unrecognized GPUs (caller falls back to VRAM heuristic).
pub fn sm_count_from_device_name(name: &str) -> usize {
    let name_upper = name.to_uppercase();
    // Data centre GPUs
    if name_upper.contains("H100") { return 132; }
    if name_upper.contains("H200") { return 132; }
    if name_upper.contains("L40S") { return 142; }
    if name_upper.contains("L40 ") || name_upper.ends_with("L40") { return 142; }
    if name_upper.contains("A100") { return 108; }
    if name_upper.contains("A10G") { return 80; }
    if name_upper.contains("A10 ") || name_upper.ends_with("A10") { return 72; }
    if name_upper.contains("V100") { return 80; }
    if name_upper.contains("T4") { return 40; }
    // Consumer GPUs
    if name_upper.contains("4090") { return 128; }
    if name_upper.contains("4080") { return 76; }
    if name_upper.contains("4070") { return 46; }
    if name_upper.contains("3090") { return 82; }
    if name_upper.contains("3080") { return 68; }
    if name_upper.contains("3070") { return 46; }
    if name_upper.contains("3060") { return 28; }
    if name_upper.contains("3050") { return 20; }
    0 // unknown — caller uses VRAM heuristic
}

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

    #[test]
    fn test_optimizer_memory_multiplier() {
        assert_eq!(OptimizerType::SGD.memory_multiplier(), 1.0);
        assert_eq!(OptimizerType::Adam.memory_multiplier(), 2.0);
        assert_eq!(OptimizerType::AdamW.memory_multiplier(), 2.0);
    }

    #[test]
    fn test_model_precision_memory_multiplier() {
        assert_eq!(ModelPrecision::BF16.memory_multiplier(), 1.0);
        assert_eq!(ModelPrecision::FP32.memory_multiplier(), 2.0);
        assert_eq!(ModelPrecision::BF16.bytes_per_param(), 2.0);
        assert_eq!(ModelPrecision::FP32.bytes_per_param(), 4.0);
    }

    #[test]
    fn test_batch_size_config_default() {
        let config = BatchSizeConfig::default();
        assert_eq!(config.model_memory_mb, 125.0);
        assert_eq!(config.model_memory_mb, 125.0);
        assert_eq!(config.model_precision, ModelPrecision::BF16);
        assert_eq!(config.sequence_length, 60);
        assert_eq!(config.feature_dim, 225);
        assert!(!config.gradient_checkpointing);
        assert_eq!(config.optimizer_type, OptimizerType::Adam);
        assert_eq!(config.safety_margin, 0.20);
    }

    #[test]
    fn test_auto_batch_sizer_rtx_3050_ti() {
        // RTX 3050 Ti: 4GB total, ~3.7GB free
        let sizer = AutoBatchSizer::with_manual_memory(4096.0, 3700.0, "RTX 3050 Ti".to_owned());

        let config = BatchSizeConfig {
            model_memory_mb: 125.0, // TFT model (BF16)
            model_precision: ModelPrecision::BF16,

            sequence_length: 60,
            feature_dim: 225,
            gradient_checkpointing: false,
            optimizer_type: OptimizerType::Adam,
            safety_margin: 0.20,
            min_batch_size: 1,
            max_batch_size: 256,
        };

        let batch_size = sizer.calculate_optimal_batch_size(&config).unwrap();

        // Expected calculation with BF16:
        // Usable: 3700 * (1 - 0.20) = 2960 MB
        // Model scaled: 125 * 1.0 (BF16 multiplier) = 125 MB
        // Fixed: 125 (model) + 250 (optimizer) + 125 (gradients) + 125 (activations) = 625 MB
        // Batch overhead (BF16): max(4096 * 0.005, 50) = 50 MB
        // Available: 2960 - 625 - 50 = 2285 MB
        // Bytes per sample: 60 * 225 * 2 (BF16) * 1.2 = 32,400 bytes = 0.0309 MB
        // Max batch: 2285 / 0.0309 ≈ 73,948
        // Clamped to max_batch_size: 256
        assert!(
            batch_size >= 128 && batch_size <= 256,
            "BF16 batch_size should be 128-256 on RTX 3050 Ti, got {}",
            batch_size
        );
    }

    #[test]
    fn test_auto_batch_sizer_t4() {
        // T4: 16GB total, ~15GB free
        let sizer = AutoBatchSizer::with_manual_memory(16384.0, 15000.0, "Tesla T4".to_owned());

        let config = BatchSizeConfig {
            model_memory_mb: 125.0,
            model_precision: ModelPrecision::BF16,

            sequence_length: 60,
            feature_dim: 225,
            gradient_checkpointing: false,
            optimizer_type: OptimizerType::Adam,
            safety_margin: 0.20,
            min_batch_size: 1,
            max_batch_size: 256,
        };

        let batch_size = sizer.calculate_optimal_batch_size(&config).unwrap();

        // With 15GB free, should calculate large batch size but clamp to max_batch_size
        // The actual calculation will produce a very large number (>150K samples)
        // which clamps to 256 (max_batch_size) then floors to 256 (already power-of-2)
        assert_eq!(batch_size, 256);
    }

    #[test]
    fn test_gradient_checkpointing_increases_batch_size() {
        let sizer = AutoBatchSizer::with_manual_memory(4096.0, 3700.0, "RTX 3050 Ti".to_owned());

        let config_no_checkpointing = BatchSizeConfig {
            model_memory_mb: 125.0,
            gradient_checkpointing: false,
            ..Default::default()
        };

        let config_with_checkpointing = BatchSizeConfig {
            model_memory_mb: 125.0,
            gradient_checkpointing: true,
            ..Default::default()
        };

        let batch_size_no_cp = sizer
            .calculate_optimal_batch_size(&config_no_checkpointing)
            .unwrap();
        let batch_size_with_cp = sizer
            .calculate_optimal_batch_size(&config_with_checkpointing)
            .unwrap();

        // Gradient checkpointing should allow larger batch size (or at least same)
        assert!(batch_size_with_cp >= batch_size_no_cp);
    }

    #[test]
    fn test_insufficient_memory_error() {
        // Very small GPU memory
        let sizer = AutoBatchSizer::with_manual_memory(512.0, 400.0, "Small GPU".to_owned());

        let config = BatchSizeConfig {
            model_memory_mb: 500.0, // Model too large
            ..Default::default()
        };

        let result = sizer.calculate_optimal_batch_size(&config);
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("Insufficient GPU memory"));
    }

    #[test]
    fn test_memory_info() {
        let sizer = AutoBatchSizer::with_manual_memory(4096.0, 3700.0, "RTX 3050 Ti".to_owned());
        let info = sizer.memory_info();

        assert_eq!(info.device_name, "RTX 3050 Ti");
        assert_eq!(info.total_memory_mb, 4096.0);
        assert_eq!(info.free_memory_mb, 3700.0);
        assert_eq!(info.used_memory_mb, 396.0);
    }

    #[test]
    fn test_sgd_uses_less_memory_than_adam() {
        let sizer = AutoBatchSizer::with_manual_memory(2048.0, 1800.0, "Small GPU".to_owned());

        let config_sgd = BatchSizeConfig {
            model_memory_mb: 100.0,
            optimizer_type: OptimizerType::SGD,
            ..Default::default()
        };

        let config_adam = BatchSizeConfig {
            model_memory_mb: 100.0,
            optimizer_type: OptimizerType::Adam,
            ..Default::default()
        };

        let batch_size_sgd = sizer.calculate_optimal_batch_size(&config_sgd).unwrap();
        let batch_size_adam = sizer.calculate_optimal_batch_size(&config_adam).unwrap();

        // SGD uses less memory, so should allow larger batch size
        assert!(batch_size_sgd >= batch_size_adam);
    }

    #[test]
    fn test_fp32_vs_int8_rtx_3050_ti() {
        // RTX 3050 Ti: 4GB total, ~3.7GB free (realistic with minimal background)
        // FP32 TFT-225 is TOO LARGE for 4GB GPU with new safety margins
        // Use smaller model (100MB base instead of 125MB) to demonstrate comparison
        let sizer = AutoBatchSizer::with_manual_memory(4096.0, 3700.0, "RTX 3050 Ti".to_owned());

        // FP32 smaller TFT model (80MB base) with gradient checkpointing
        // This is a SMALL model for testing comparison, NOT production TFT-225
        let config_fp32 = BatchSizeConfig {
            model_memory_mb: 80.0, // Legacy field (backward compat)
            model_precision: ModelPrecision::FP32,

            sequence_length: 60,
            feature_dim: 150,             // Fewer features for testing
            gradient_checkpointing: true, // Enable to fit in memory
            optimizer_type: OptimizerType::Adam,
            safety_margin: 0.20, // 20% margin (overridden by FP32 precision margin: 50%)
            min_batch_size: 1,
            max_batch_size: 64,
        };

        // BF16 TFT-225 model config (standard production size)
        let config_bf16 = BatchSizeConfig {
            model_memory_mb: 125.0,
            model_precision: ModelPrecision::BF16,

            sequence_length: 60,
            feature_dim: 225,
            gradient_checkpointing: false,
            optimizer_type: OptimizerType::Adam,
            safety_margin: 0.20,
            min_batch_size: 1,
            max_batch_size: 256,
        };

        let batch_size_fp32 = sizer.calculate_optimal_batch_size(&config_fp32).unwrap();
        let batch_size_bf16 = sizer.calculate_optimal_batch_size(&config_bf16).unwrap();

        info!(
            batch_size_fp32,
            batch_size_bf16,
            "DEBUG: FP32 vs BF16 batch sizes"
        );

        // FP32 should get smaller batch size due to:
        // - 2x larger model (160MB vs 125MB with BF16 multiplier)
        // - 25% safety margin vs 20%
        assert!(
            batch_size_fp32 <= batch_size_bf16,
            "FP32 batch_size ({}) should be <= BF16 batch_size ({})",
            batch_size_fp32,
            batch_size_bf16
        );

        // Verify FP32 gets reasonable batch size (1-64)
        assert!(
            batch_size_fp32 >= 1 && batch_size_fp32 <= 64,
            "FP32 batch_size should be 1-64 on 4GB GPU with small model, got {}",
            batch_size_fp32
        );

        // Verify BF16 gets larger batch size (32+)
        assert!(
            batch_size_bf16 >= 32,
            "BF16 batch_size should be >=32 on 4GB GPU, got {}",
            batch_size_bf16
        );
    }

    #[test]
    fn test_fp32_requires_larger_gpu() {
        // Very small GPU: 1GB total, ~900MB free
        let sizer = AutoBatchSizer::with_manual_memory(1024.0, 900.0, "Small GPU".to_owned());

        // FP32 TFT-225 model config
        // FP32 125MB × 2.0 multiplier = 250MB model
        // Fixed overhead: 250 (model) + 500 (Adam) + 250 (grad) + 250 (act) = 1250MB
        // Usable: 900 × 0.75 = 675MB — NOT enough for 1250MB fixed overhead
        let config_fp32 = BatchSizeConfig {
            model_memory_mb: 125.0,
            model_precision: ModelPrecision::FP32,

            sequence_length: 60,
            feature_dim: 225,
            gradient_checkpointing: false,
            optimizer_type: OptimizerType::Adam,
            safety_margin: 0.20,
            min_batch_size: 1,
            max_batch_size: 256,
        };

        let result = sizer.calculate_optimal_batch_size(&config_fp32);

        assert!(
            result.is_err(),
            "FP32 model should fail on small GPU, but got: {:?}",
            result
        );
        if let Err(e) = result {
            assert!(
                e.to_string().contains("Insufficient GPU memory"),
                "Error should mention insufficient memory, got: {}",
                e
            );
        }
    }

    #[test]
    fn test_bf16_works_on_small_gpu() {
        // Small GPU: 2GB total, ~1.8GB free
        let sizer = AutoBatchSizer::with_manual_memory(2048.0, 1800.0, "Small GPU".to_owned());

        // BF16 TFT-225 model config
        let config_bf16 = BatchSizeConfig {
            model_memory_mb: 125.0,
            model_precision: ModelPrecision::BF16,

            sequence_length: 60,
            feature_dim: 225,
            gradient_checkpointing: false,
            optimizer_type: OptimizerType::Adam,
            safety_margin: 0.20,
            min_batch_size: 1,
            max_batch_size: 256,
        };

        let batch_size = sizer.calculate_optimal_batch_size(&config_bf16).unwrap();

        // BF16 model (125MB * 1.0 multiplier = 125MB, overhead ~500MB) fits comfortably
        assert!(
            batch_size >= 16,
            "BF16 model should work on small GPU with reasonable batch size, got {}",
            batch_size
        );
    }

    #[test]
    fn test_reduce_batch_size() {
        // Test exponential backoff: 64 → 32 → 16 → 8 → 4 → 2 → 1
        assert_eq!(AutoBatchSizer::reduce_batch_size(64), 32);
        assert_eq!(AutoBatchSizer::reduce_batch_size(32), 16);
        assert_eq!(AutoBatchSizer::reduce_batch_size(16), 8);
        assert_eq!(AutoBatchSizer::reduce_batch_size(8), 4);
        assert_eq!(AutoBatchSizer::reduce_batch_size(4), 2);
        assert_eq!(AutoBatchSizer::reduce_batch_size(2), 1);
        assert_eq!(AutoBatchSizer::reduce_batch_size(1), 1); // Minimum is 1
    }

    #[test]
    fn test_is_batch_size_too_small() {
        // Below 4 is too small (GPU underutilized)
        assert!(AutoBatchSizer::is_batch_size_too_small(1));
        assert!(AutoBatchSizer::is_batch_size_too_small(2));
        assert!(AutoBatchSizer::is_batch_size_too_small(3));

        // 4 and above is acceptable
        assert!(!AutoBatchSizer::is_batch_size_too_small(4));
        assert!(!AutoBatchSizer::is_batch_size_too_small(8));
        assert!(!AutoBatchSizer::is_batch_size_too_small(32));
    }

    #[test]
    fn test_oom_recovery_simulation() {
        // Simulate OOM recovery with exponential backoff
        let mut current_batch_size = 64;
        let mut oom_retry_count = 0;
        const MAX_OOM_RETRIES: usize = 3;

        let mut batch_sizes = vec![current_batch_size];

        // Simulate 3 OOM events
        while oom_retry_count < MAX_OOM_RETRIES {
            oom_retry_count += 1;
            current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);
            batch_sizes.push(current_batch_size);

            // Check if batch size is too small (abort condition)
            if AutoBatchSizer::is_batch_size_too_small(current_batch_size) {
                break;
            }
        }

        // Verify batch size reduction sequence: 64 → 32 → 16 → 8
        assert_eq!(batch_sizes, vec![64, 32, 16, 8]);

        // Verify final batch size is still viable (>=4)
        assert!(!AutoBatchSizer::is_batch_size_too_small(current_batch_size));
    }

    #[test]
    fn test_sm_count_from_device_name() {
        assert_eq!(sm_count_from_device_name("NVIDIA H100 80GB HBM3"), 132);
        assert_eq!(sm_count_from_device_name("NVIDIA L40S"), 142);
        assert_eq!(sm_count_from_device_name("NVIDIA A100-SXM4-80GB"), 108);
        assert_eq!(sm_count_from_device_name("Tesla T4"), 40);
        assert_eq!(sm_count_from_device_name("NVIDIA GeForce RTX 4090"), 128);
        assert_eq!(sm_count_from_device_name("NVIDIA GeForce RTX 3050 Ti Laptop GPU"), 20);
        assert_eq!(sm_count_from_device_name("Unknown GPU Model"), 0);
    }

    #[test]
    fn test_optimal_n_episodes_h100() {
        let hw = GpuHardwareInfo {
            device_name: "NVIDIA H100".to_owned(),
            total_memory_mb: 81920.0,
            free_memory_mb: 78000.0,
            sm_count: 132,
        };
        let n = hw.optimal_n_episodes(48, 500);
        // 132 SMs × 2 warps × 32 threads = 8448, rounded to 8192 (MAX)
        assert!(n >= 4096, "H100 should use at least 4096 episodes, got {}", n);
        assert!(n <= 8192, "H100 should cap at 8192 episodes, got {}", n);
        assert_eq!(n % 256, 0, "Should be 256-aligned, got {}", n);
    }

    #[test]
    fn test_optimal_n_episodes_rtx3050() {
        let hw = GpuHardwareInfo {
            device_name: "RTX 3050 Ti".to_owned(),
            total_memory_mb: 4096.0,
            free_memory_mb: 3500.0,
            sm_count: 20,
        };
        let n = hw.optimal_n_episodes(48, 500);
        // 20 SMs × 2 × 32 = 1280, rounded up to 1280 → aligned to 1536
        assert!(n >= 128, "RTX 3050 should use at least 128 episodes, got {}", n);
        assert!(n <= 2048, "RTX 3050 should stay ≤2048, got {}", n);
        assert_eq!(n % 256, 0, "Should be 256-aligned, got {}", n);
    }

    #[test]
    fn test_optimal_n_episodes_unknown_gpu_vram_heuristic() {
        let hw = GpuHardwareInfo {
            device_name: "Unknown GPU".to_owned(),
            total_memory_mb: 24576.0, // 24 GB
            free_memory_mb: 22000.0,
            sm_count: 0, // unknown
        };
        let n = hw.optimal_n_episodes(48, 500);
        // Fallback: 24 GB × 128 = 3072, aligned to 3072
        assert!(n >= 1024, "24GB unknown GPU should use at least 1024, got {}", n);
        assert_eq!(n % 256, 0, "Should be 256-aligned, got {}", n);
    }

    #[test]
    fn test_optimal_n_episodes_vram_capped() {
        let hw = GpuHardwareInfo {
            device_name: "NVIDIA H100".to_owned(),
            total_memory_mb: 81920.0,
            free_memory_mb: 500.0, // almost no free VRAM
            sm_count: 132,
        };
        let n = hw.optimal_n_episodes(48, 500);
        // Free VRAM = 500 MB, 15% = 75 MB budget
        // Per episode: 500 × (48+5) × 4 = 106 KB → ~707 episodes → capped, aligned
        assert!(n >= 128, "Should hit minimum 128 even with low VRAM, got {}", n);
        assert!(n <= 1024, "Should be VRAM-capped with only 500MB free, got {}", n);
    }

    #[test]
    fn test_optimal_replay_capacity_h100() {
        let hw = GpuHardwareInfo {
            device_name: "NVIDIA H100".to_owned(),
            total_memory_mb: 81920.0,
            free_memory_mb: 72_000.0,
            sm_count: 132,
        };
        let cfg = hw.optimal_replay_config(56, 0.70);
        assert!(cfg.capacity >= 500_000, "H100 should get at least 500K capacity, got {}", cfg.capacity);
        assert!(cfg.capacity <= 10_000_000, "Should be capped at max, got {}", cfg.capacity);
        assert!(cfg.per_max_buffer_bytes > 0, "PER budget must be positive");
    }

    #[test]
    fn test_optimal_replay_capacity_small_gpu() {
        let hw = GpuHardwareInfo {
            device_name: "RTX 3050 Ti".to_owned(),
            total_memory_mb: 4096.0,
            free_memory_mb: 2_000.0,
            sm_count: 20,
        };
        let cfg = hw.optimal_replay_config(48, 0.70);
        assert!(cfg.capacity >= 100_000, "Small GPU should get at least min capacity");
        // 20% of 4096 MB = ~819 MB, but min is 1 GB
        assert_eq!(cfg.per_max_buffer_bytes, 1024 * 1024 * 1024, "Small GPU PER budget should be min 1 GB");
    }
}
