//! Machine learning configuration

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct MLConfig {
    pub model_config: ModelArchitectureConfig,
    pub training_config: TrainingConfig,
    pub simulation_config: SimulationConfig,
}

/// Configuration for market data simulation and stress testing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationConfig {
    /// Initial market state with configurable symbol prices
    pub initial_market_state: MarketState,
    /// Simulation parameters
    pub parameters: SimulationParameters,
    /// Test symbol configuration for generic testing
    pub test_symbols: TestSymbolConfig,
}

/// Initial market state configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MarketState {
    /// Symbol-specific initial prices and configuration
    pub symbols: HashMap<String, SymbolConfig>,
    /// Default configuration for unlisted symbols
    pub default_symbol: SymbolConfig,
}

/// Configuration for individual symbols
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SymbolConfig {
    /// Initial price for the symbol
    pub initial_price: f64,
    /// Base volatility for the symbol
    pub volatility: f64,
    /// Base trading volume
    pub base_volume: f64,
    /// Minimum spread in basis points
    pub min_spread_bps: f64,
    /// Maximum spread in basis points
    pub max_spread_bps: f64,
    /// Market capitalization tier (affects behavior)
    pub market_cap_tier: MarketCapTier,
}

/// Market capitalization tiers for different symbol behaviors
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MarketCapTier {
    /// Large cap stocks (>$10B)
    LargeCap,
    /// Mid cap stocks ($2B-$10B)
    MidCap,
    /// Small cap stocks (<$2B)
    SmallCap,
    /// Generic test symbol
    Test,
}

/// Simulation parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SimulationParameters {
    /// Update rate in Hz
    pub update_rate_hz: u32,
    /// Base market volatility
    pub base_volatility: f64,
    /// Market trend direction (-1.0 to 1.0)
    pub trend: f64,
    /// Enable realistic market microstructure
    pub enable_microstructure: bool,
    /// Enable correlated movements between symbols
    pub enable_correlation: bool,
}

/// Test symbol configuration for generic testing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestSymbolConfig {
    /// Prefix for test symbols (e.g., "TEST")
    pub symbol_prefix: String,
    /// Number of test symbols to generate
    pub count: usize,
    /// Price range for test symbols
    pub price_range: (f64, f64),
    /// Volume range for test symbols
    pub volume_range: (f64, f64),
}

/// Default simulation configuration
impl Default for SimulationConfig {
    fn default() -> Self {
        let mut symbols = HashMap::new();

        // Production-ready major symbols with realistic configurations
        symbols.insert(
            "AAPL".to_owned(),
            SymbolConfig {
                initial_price: 150.0,
                volatility: 0.25,
                base_volume: 50_000_000.0,
                min_spread_bps: 1.0,
                max_spread_bps: 5.0,
                market_cap_tier: MarketCapTier::LargeCap,
            },
        );

        symbols.insert(
            "MSFT".to_owned(),
            SymbolConfig {
                initial_price: 300.0,
                volatility: 0.22,
                base_volume: 30_000_000.0,
                min_spread_bps: 1.0,
                max_spread_bps: 5.0,
                market_cap_tier: MarketCapTier::LargeCap,
            },
        );

        symbols.insert(
            "GOOGL".to_owned(),
            SymbolConfig {
                initial_price: 2500.0,
                volatility: 0.28,
                base_volume: 20_000_000.0,
                min_spread_bps: 2.0,
                max_spread_bps: 8.0,
                market_cap_tier: MarketCapTier::LargeCap,
            },
        );

        symbols.insert(
            "TSLA".to_owned(),
            SymbolConfig {
                initial_price: 800.0,
                volatility: 0.45,
                base_volume: 80_000_000.0,
                min_spread_bps: 2.0,
                max_spread_bps: 10.0,
                market_cap_tier: MarketCapTier::LargeCap,
            },
        );

        symbols.insert(
            "AMZN".to_owned(),
            SymbolConfig {
                initial_price: 3200.0,
                volatility: 0.30,
                base_volume: 25_000_000.0,
                min_spread_bps: 2.0,
                max_spread_bps: 8.0,
                market_cap_tier: MarketCapTier::LargeCap,
            },
        );

        symbols.insert(
            "NVDA".to_owned(),
            SymbolConfig {
                initial_price: 500.0,
                volatility: 0.40,
                base_volume: 40_000_000.0,
                min_spread_bps: 2.0,
                max_spread_bps: 8.0,
                market_cap_tier: MarketCapTier::LargeCap,
            },
        );

        Self {
            initial_market_state: MarketState {
                symbols,
                default_symbol: SymbolConfig {
                    initial_price: 100.0,
                    volatility: 0.30,
                    base_volume: 1_000_000.0,
                    min_spread_bps: 5.0,
                    max_spread_bps: 20.0,
                    market_cap_tier: MarketCapTier::Test,
                },
            },
            parameters: SimulationParameters {
                update_rate_hz: 1000,
                base_volatility: 0.02,
                trend: 0.0,
                enable_microstructure: true,
                enable_correlation: false,
            },
            test_symbols: TestSymbolConfig {
                symbol_prefix: "TEST".to_owned(),
                count: 10,
                price_range: (50.0, 500.0),
                volume_range: (100_000.0, 10_000_000.0),
            },
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelArchitectureConfig {
    pub model_type: String,
    pub hidden_dims: Vec<usize>,
    pub dropout_rate: f64,
    pub activation: String,
}

impl Default for ModelArchitectureConfig {
    fn default() -> Self {
        Self {
            model_type: "transformer".to_owned(),
            hidden_dims: vec![256, 128, 64],
            dropout_rate: 0.1,
            activation: "relu".to_owned(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
    pub batch_size: usize,
    pub learning_rate: f64,
    pub epochs: u32,
    pub early_stopping_patience: u32,
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            batch_size: 32,
            learning_rate: 0.001,
            epochs: 100,
            early_stopping_patience: 10,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Mamba2Config {
    pub d_model: usize,
    pub d_state: usize,
    pub d_conv: usize,
    pub expand: usize,
    pub dt_rank: Option<usize>,
    pub dt_min: f64,
    pub dt_max: f64,
    pub dt_init: String,
    pub dt_scale: f64,
    pub dt_init_floor: f64,
    pub conv_bias: bool,
    pub bias: bool,
    pub use_fast_path: bool,
    pub layer_idx: Option<usize>,
    pub device: Option<String>,
    pub dtype: Option<String>,
    pub d_head: usize,
    pub num_heads: usize,
    pub num_layers: usize,
    pub target_latency_us: u64,
    pub hardware_aware: bool,
    pub use_ssd: bool,
    pub use_selective_state: bool,
    pub max_seq_len: usize,
    pub batch_size: usize,
    pub seq_len: usize,
    pub dropout: f64,
}

impl Default for Mamba2Config {
    fn default() -> Self {
        Self {
            d_model: 768,
            d_state: 128,
            d_conv: 4,
            expand: 2,
            dt_rank: None, // Auto-calculated as ceil(d_model / 16)
            dt_min: 0.001,
            dt_max: 0.1,
            dt_init: "random".to_owned(),
            dt_scale: 1.0,
            dt_init_floor: 1e-4,
            conv_bias: true,
            bias: false,
            use_fast_path: true,
            layer_idx: None,
            device: None,
            dtype: None,
            d_head: 32,
            num_heads: 8,
            num_layers: 4,
            target_latency_us: 3,
            hardware_aware: true,
            use_ssd: true,
            use_selective_state: true,
            max_seq_len: 1024,
            batch_size: 1,
            seq_len: 256,
            dropout: 0.0,
        }
    }
}
