//! Canonical model type definitions for the Foxhunt trading system.
//!
//! This module provides the unified `ModelType` enum used across all crates
//! (ml, model_loader, services) to identify ML model architectures.

use serde::{Deserialize, Serialize};
use std::fmt;

/// Canonical model type enum used throughout the Foxhunt system.
///
/// This is the single source of truth for model type identification.
/// All crates should import this type rather than defining their own.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModelType {
    /// Compact Deep Q-Network
    CompactDQN,
    /// Distilled micro network for ultra-low latency
    DistilledMicroNet,
    /// Standard Deep Q-Network
    DQN,
    /// Rainbow DQN with all enhancements
    RainbowDQN,
    /// MAMBA/Mamba2 state space model
    MAMBA,
    /// Temporal Fusion Transformer
    TFT,
    /// Temporal Graph Neural Network
    TGGN,
    /// Liquid Neural Network
    LNN,
    /// Temporal Limit Order Book transformer
    TLOB,
    /// Proximal Policy Optimization
    PPO,
    /// Transformer for sequence modeling
    Transformer,
    /// Mamba state space model (alias for MAMBA)
    Mamba,
    /// Liquid time constant networks (alias for LNN)
    LiquidNet,
    /// Temporal Graph Neural Network (alias for TGGN)
    TGNN,
    /// Ensemble of multiple models
    Ensemble,
    /// Kolmogorov-Arnold Network
    KAN,
    /// Extended LSTM (sLSTM + mLSTM)
    XLSTM,
    /// Diffusion model (DDPM/DDIM) for price path generation
    Diffusion,
}

impl ModelType {
    /// Get file extension for model type (used for checkpoint filenames).
    pub fn file_extension(&self) -> &'static str {
        match self {
            Self::DQN => "dqn",
            Self::MAMBA | Self::Mamba => "mamba",
            Self::TFT => "tft",
            Self::TGGN | Self::TGNN => "tggn",
            Self::LNN | Self::LiquidNet => "lnn",
            Self::CompactDQN => "compact_dqn",
            Self::DistilledMicroNet => "distilled",
            Self::RainbowDQN => "rainbow_dqn",
            Self::TLOB => "tlob",
            Self::PPO => "ppo",
            Self::Transformer => "transformer",
            Self::Ensemble => "ensemble",
            Self::KAN => "kan",
            Self::XLSTM => "xlstm",
            Self::Diffusion => "diffusion",
        }
    }

    /// String representation of the model type (lowercase, unique per variant family).
    pub const fn as_str(&self) -> &'static str {
        match self {
            Self::DQN | Self::CompactDQN | Self::RainbowDQN => "dqn",
            Self::MAMBA | Self::Mamba => "mamba",
            Self::TFT => "tft",
            Self::TGGN | Self::TGNN => "tggn",
            Self::LNN | Self::LiquidNet => "liquid",
            Self::TLOB => "tlob",
            Self::PPO => "ppo",
            Self::DistilledMicroNet => "distilled",
            Self::Transformer => "transformer",
            Self::Ensemble => "ensemble",
            Self::KAN => "kan",
            Self::XLSTM => "xlstm",
            Self::Diffusion => "diffusion",
        }
    }

    /// S3 storage prefix for model files (used by model_loader).
    pub const fn s3_prefix(&self) -> &'static str {
        match self {
            Self::TLOB => "tlob_transformer",
            Self::DQN | Self::CompactDQN | Self::RainbowDQN => "dqn",
            Self::MAMBA | Self::Mamba => "mamba2",
            Self::TFT => "tft",
            Self::PPO => "ppo",
            Self::LNN | Self::LiquidNet => "liquid",
            Self::Ensemble => "ensemble",
            Self::TGGN | Self::TGNN => "tggn",
            Self::DistilledMicroNet => "distilled",
            Self::Transformer => "transformer",
            Self::KAN => "kan",
            Self::XLSTM => "xlstm",
            Self::Diffusion => "diffusion",
        }
    }

    /// Convert to database string representation (for PostgreSQL storage).
    pub fn to_db_string(&self) -> &'static str {
        match self {
            Self::DQN => "DQN",
            Self::PPO => "PPO",
            Self::MAMBA | Self::Mamba => "MAMBA-2",
            Self::TFT => "TFT",
            Self::CompactDQN => "COMPACT_DQN",
            Self::DistilledMicroNet => "DISTILLED_MICRO_NET",
            Self::RainbowDQN => "RAINBOW_DQN",
            Self::TGGN | Self::TGNN => "TGGN",
            Self::LNN | Self::LiquidNet => "LNN",
            Self::TLOB => "TLOB",
            Self::Transformer => "TRANSFORMER",
            Self::Ensemble => "ENSEMBLE",
            Self::KAN => "KAN",
            Self::XLSTM => "XLSTM",
            Self::Diffusion => "DIFFUSION",
        }
    }

    /// Get model weight for progress aggregation.
    pub fn weight(&self) -> f64 {
        0.25 // Equal weight for all models
    }

    /// Parse model type from string (case-insensitive).
    #[allow(clippy::should_implement_trait)]
    pub fn from_str(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "dqn" => Some(Self::DQN),
            "mamba" | "mamba2" | "mamba-2" => Some(Self::MAMBA),
            "tft" => Some(Self::TFT),
            "tggn" | "tgnn" => Some(Self::TGGN),
            "lnn" | "liquidnet" | "liquid" => Some(Self::LNN),
            "compact_dqn" | "compactdqn" => Some(Self::CompactDQN),
            "distilled" | "distilledmicronet" => Some(Self::DistilledMicroNet),
            "rainbow_dqn" | "rainbowdqn" => Some(Self::RainbowDQN),
            "tlob" | "tlob_transformer" => Some(Self::TLOB),
            "ppo" => Some(Self::PPO),
            "transformer" => Some(Self::Transformer),
            "ensemble" => Some(Self::Ensemble),
            "kan" => Some(Self::KAN),
            "xlstm" | "x-lstm" | "extended_lstm" => Some(Self::XLSTM),
            "diffusion" | "ddpm" | "ddim" => Some(Self::Diffusion),
            _ => None,
        }
    }
}

impl fmt::Display for ModelType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.as_str())
    }
}

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

    #[test]
    fn test_file_extension() {
        assert_eq!(ModelType::DQN.file_extension(), "dqn");
        assert_eq!(ModelType::PPO.file_extension(), "ppo");
        assert_eq!(ModelType::TFT.file_extension(), "tft");
        assert_eq!(ModelType::MAMBA.file_extension(), "mamba");
        assert_eq!(ModelType::Mamba.file_extension(), "mamba");
        assert_eq!(ModelType::Ensemble.file_extension(), "ensemble");
    }

    #[test]
    fn test_as_str() {
        assert_eq!(ModelType::TLOB.as_str(), "tlob");
        assert_eq!(ModelType::DQN.as_str(), "dqn");
        assert_eq!(ModelType::MAMBA.as_str(), "mamba");
        assert_eq!(ModelType::TFT.as_str(), "tft");
        assert_eq!(ModelType::PPO.as_str(), "ppo");
        assert_eq!(ModelType::LNN.as_str(), "liquid");
        assert_eq!(ModelType::Ensemble.as_str(), "ensemble");
    }

    #[test]
    fn test_s3_prefix() {
        assert_eq!(ModelType::TLOB.s3_prefix(), "tlob_transformer");
        assert_eq!(ModelType::DQN.s3_prefix(), "dqn");
        assert_eq!(ModelType::MAMBA.s3_prefix(), "mamba2");
        assert_eq!(ModelType::TFT.s3_prefix(), "tft");
    }

    #[test]
    fn test_to_db_string() {
        assert_eq!(ModelType::DQN.to_db_string(), "DQN");
        assert_eq!(ModelType::PPO.to_db_string(), "PPO");
        assert_eq!(ModelType::MAMBA.to_db_string(), "MAMBA-2");
        assert_eq!(ModelType::TFT.to_db_string(), "TFT");
    }

    #[test]
    fn test_weight() {
        assert_eq!(ModelType::DQN.weight(), 0.25);
        assert_eq!(ModelType::PPO.weight(), 0.25);
    }

    #[test]
    fn test_display() {
        assert_eq!(format!("{}", ModelType::DQN), "dqn");
        assert_eq!(format!("{}", ModelType::PPO), "ppo");
        assert_eq!(format!("{}", ModelType::MAMBA), "mamba");
        assert_eq!(format!("{}", ModelType::TLOB), "tlob");
    }

    #[test]
    fn test_from_str() {
        assert_eq!(ModelType::from_str("dqn"), Some(ModelType::DQN));
        assert_eq!(ModelType::from_str("DQN"), Some(ModelType::DQN));
        assert_eq!(ModelType::from_str("mamba2"), Some(ModelType::MAMBA));
        assert_eq!(ModelType::from_str("ppo"), Some(ModelType::PPO));
        assert_eq!(ModelType::from_str("unknown"), None);
    }

    #[test]
    fn test_serde_roundtrip() {
        let model = ModelType::DQN;
        let json = serde_json::to_string(&model).unwrap_or_default();
        let deserialized: ModelType =
            serde_json::from_str(&json).unwrap_or(ModelType::DQN);
        assert_eq!(model, deserialized);
    }
}
