//! Per-model memory estimation and dynamic batch size resolution.
//!
//! Each ML model provides a `ModelMemoryEstimate` describing its memory footprint.
//! `resolve_batch_size()` combines this with `GpuCapabilities` to compute the
//! optimal batch size for the detected hardware.

use serde::{Deserialize, Serialize};

/// Memory footprint estimate for an ML model.
///
/// Each model type provides one of these so the batch size resolver
/// can compute hardware-appropriate limits.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMemoryEstimate {
    /// Model name (for logging)
    pub name: &'static str,
    /// Estimated parameter count
    pub param_count: usize,
    /// Activation memory multiplier (1.0 for simple MLPs, 2.0+ for attention)
    pub activation_multiplier: f64,
    /// Whether the model supports gradient checkpointing
    pub supports_checkpointing: bool,
    /// Default sequence length for this model
    pub default_seq_len: usize,
    /// Default feature dimension for this model
    pub default_feature_dim: usize,
}

impl ModelMemoryEstimate {
    /// Estimated model size in MB (FP32, parameters only).
    pub fn estimated_size_mb(&self) -> f64 {
        (self.param_count as f64 * 4.0) / (1024.0 * 1024.0)
    }
}

/// Pre-defined memory estimates for all 10 Foxhunt models.
pub mod estimates {
    use super::ModelMemoryEstimate;

    pub const DQN: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "DQN",
        param_count: 12_000_000, // worst-case: [54->4096->2048->1024->45]
        activation_multiplier: 1.5,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 54,
    };

    pub const PPO: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "PPO",
        param_count: 20_000_000, // actor + critic with large hidden dims
        activation_multiplier: 1.5,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 54,
    };

    pub const TFT: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "TFT",
        param_count: 50_000_000, // hidden_size=2048, multi-head attention
        activation_multiplier: 3.0, // attention workspace is large
        supports_checkpointing: true,
        default_seq_len: 60,
        default_feature_dim: 225,
    };

    pub const MAMBA2: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "Mamba2",
        param_count: 5_000_000, // state-space layers with large hidden
        activation_multiplier: 2.0,
        supports_checkpointing: false,
        default_seq_len: 60,
        default_feature_dim: 64,
    };

    pub const TGGN: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "TGGN",
        param_count: 5_000_000, // hidden_dim up to 1024
        activation_multiplier: 2.0,
        supports_checkpointing: false,
        default_seq_len: 32,
        default_feature_dim: 64,
    };

    pub const TLOB: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "TLOB",
        param_count: 10_000_000, // d_model up to 1024
        activation_multiplier: 2.5,
        supports_checkpointing: false,
        default_seq_len: 128,
        default_feature_dim: 43,
    };

    pub const LIQUID: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "Liquid",
        param_count: 10_000_000, // hidden up to 2048
        activation_multiplier: 1.5,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 54,
    };

    pub const KAN: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "KAN",
        param_count: 2_000_000, // hidden_width up to 256
        activation_multiplier: 1.5,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 54,
    };

    pub const XLSTM: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "xLSTM",
        param_count: 20_000_000, // hidden_dim up to 2048
        activation_multiplier: 2.0,
        supports_checkpointing: false,
        default_seq_len: 60,
        default_feature_dim: 64,
    };

    pub const DIFFUSION: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "Diffusion",
        param_count: 20_000_000, // hidden_dim up to 2048
        activation_multiplier: 2.5,
        supports_checkpointing: true,
        default_seq_len: 32,
        default_feature_dim: 64,
    };
}

// detect_vram_mb, vram_scaled_hidden_dims, vram_scaled_base_dim,
// resolve_hidden_dim_base — ALL REMOVED.
// Replaced by GpuProfile::load() from gpu::profile (TOML config files).
// Replaced by GpuProfile::load() from gpu::profile (TOML config files).

/// Compute total parameter count for a fully-connected network given layer dimensions.
///
/// Each layer has `in * out` weights + `out` biases.
/// Example: `[54, 256, 128, 45]` → 54×256+256 + 256×128+128 + 128×45+45
pub fn network_param_count(layer_dims: &[usize]) -> usize {
    layer_dims
        .windows(2)
        .map(|w| w[0] * w[1] + w[1])
        .sum()
}

// NOTE: `resolve_batch_size()` lives in `ml::batch_size_resolver` because it depends
// on `memory_optimization::auto_batch_size` which remains in the `ml` crate.
// It is re-exported as `ml::gpu::memory_profile::resolve_batch_size` for backward compat.

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

    #[test]
    fn test_dqn_estimate_size() {
        let size = estimates::DQN.estimated_size_mb();
        // 12M params * 4B / 1MB = ~45.8 MB
        assert!(size > 40.0 && size < 50.0, "DQN should be ~45.8 MB, got {}", size);
    }

    #[test]
    fn test_tft_estimate_size() {
        let size = estimates::TFT.estimated_size_mb();
        // 50M params * 4B / 1MB = ~190.7 MB
        assert!(size > 180.0 && size < 200.0, "TFT should be ~190.7 MB, got {}", size);
    }

    #[test]
    fn test_network_param_count() {
        // [54, 256, 128, 64, 45] → 54*256+256 + 256*128+128 + 128*64+64 + 64*45+45
        // = 14080 + 32896 + 8256 + 2925 = 58157
        assert_eq!(network_param_count(&[54, 256, 128, 64, 45]), 58157);
        // Single layer: [10, 5] → 10*5 + 5 = 55
        assert_eq!(network_param_count(&[10, 5]), 55);
        // Edge: single dim → 0 params
        assert_eq!(network_param_count(&[128]), 0);
    }

    #[test]
    fn test_all_models_have_valid_estimates() {
        let all = [
            &estimates::DQN, &estimates::PPO, &estimates::TFT,
            &estimates::MAMBA2, &estimates::TGGN, &estimates::TLOB,
            &estimates::LIQUID, &estimates::KAN, &estimates::XLSTM,
            &estimates::DIFFUSION,
        ];
        for est in all {
            assert!(!est.name.is_empty());
            assert!(est.param_count > 0);
            assert!(est.activation_multiplier >= 1.0);
            assert!(est.default_seq_len > 0);
            assert!(est.default_feature_dim > 0);
            let size = est.estimated_size_mb();
            assert!(size > 0.0, "{} has zero estimated size", est.name);
        }
    }

}
