"""MANIFOLD configuration models using Pydantic v2."""

from pydantic import BaseModel
from typing import Literal, Optional


class ModelConfig(BaseModel):
    """MANIFOLD-Lite model configuration"""
    input_dim: int = 64
    embed_dim: int = 256
    sequence_length: int = 128
    
    # IHE
    ihe_layers: int = 4
    ihe_heads: int = 8
    ihe_ff_dim: int = 1024
    ihe_dropout: float = 0.1
    
    # MDM
    mdm_hidden: int = 512
    mdm_steps: int = 4
    
    # MPL
    latent_dim: int = 64
    mpl_hidden: int = 256
    kl_weight: float = 0.001
    
    # CCA
    num_cf_probes: int = 16
    cca_heads: int = 8
    
    # HSE
    manifold_dim: int = 32
    num_skill_levels: int = 7
    
    # TIV
    num_domains: int = 4
    adversarial_lambda: float = 0.1
    
    # Verdict
    num_classes: int = 3
    evidence_scale: float = 10.0
    dropout: float = 0.1


class TrainingConfig(BaseModel):
    """Training configuration"""
    batch_size: int = 32
    effective_batch_size: int = 128
    learning_rate: float = 3e-4
    min_learning_rate: float = 1e-6
    weight_decay: float = 0.01
    warmup_ratio: float = 0.1
    max_epochs: int = 50
    gradient_clip: float = 1.0
    use_amp: bool = True
    amp_dtype: Literal["float16", "bfloat16"] = "float16"
    gradient_checkpointing: bool = True
    save_every_n_epochs: int = 5
    
    loss_weights: dict = {
        "classification": 1.0,
        "reconstruction": 0.1,
        "kl_divergence": 0.001,
        "physics_violation": 0.5,
        "invariance": 0.1,
    }


class DataConfig(BaseModel):
    """Data generation configuration"""
    num_legit_players: int = 70000
    num_cheaters: int = 30000
    engagements_per_session: int = 200
    num_features: int = 64
    trajectory_length: int = 128
    seed: Optional[int] = None
    
    cheater_distribution: dict = {
        "blatant_rage": 0.10,
        "obvious": 0.15,
        "closet_moderate": 0.30,
        "closet_subtle": 0.30,
        "wallhack_only": 0.15,
    }
    
    rank_distribution: dict = {
        "silver": 0.20,
        "gold_nova": 0.25,
        "master_guardian": 0.25,
        "legendary_eagle": 0.15,
        "supreme_global": 0.10,
        "pro": 0.05,
    }