from dataclasses import dataclass, field
from typing import List

@dataclass
class ModelConfig:
    # ── Architecture ──────────────────────────────────────────────────────
    vocab_size: int = 5010      # 5000 Words + 4 Specials + 6 Agentic (CoT & Tools)
    block_size: int = 2048      # Context window (increased for deeper reasoning)
    n_embd: int = 512           # Embedding dimension (384→512 for higher capacity)
    n_head: int = 16            # Query heads (12→16)
    n_kv_heads: int = 4         # Key/Value heads (GQA 4:1 for KV cache savings)
    n_layer: int = 16           # Transformer layers (12→16 for deeper representations)
    
    # ── MoE ───────────────────────────────────────────────────────────────
    n_experts: int = 8          # Total number of routed experts
    n_shared_experts: int = 1   # Always-active shared experts (DeepSeek pattern)
    top_k: int = 2              # Experts activated per token
    moe_intermediate_size: int = 768  # Expert MLP hidden size (512→768)
    expert_capacity_factor: float = 1.25  # Capacity buffer to prevent dropping
    
    # ── Training ──────────────────────────────────────────────────────────
    batch_size: int = 8         # Micro batch (reduced for 8GB VRAM)
    learning_rate: float = 2e-4 # Slightly lower for stability
    weight_decay: float = 0.01
    dropout: float = 0.05       # Lower dropout for better generalization
    label_smoothing: float = 0.1
    max_grad_norm: float = 1.0
    warmup_steps: int = 100
    
    # ── Gradient Accumulation ─────────────────────────────────────────────
    gradient_accumulation_steps: int = 16  # Effective batch = 8*16 = 128
    
    # ── RTX 4060 Optimizations ────────────────────────────────────────────
    device: str = 'cuda'
    use_checkpointing: bool = True   # Gradient checkpointing (cuts VRAM ~2x)
    pin_memory: bool = True          # DMA pinned memory for async transfer
    expert_offloading: bool = True   # Offload idle MoE experts to CPU
    use_compile: bool = True         # torch.compile for CUDA graphs
    use_amp: bool = True             # Automatic mixed precision (bf16/fp16)
    
    # ── MLA (Multi-head Latent Attention) ─────────────────────────────────
    use_mla: bool = True             # Compress KV through low-rank bottleneck
    mla_latent_dim: int = 128        # Bottleneck dimension for KV compression
    
    # ── EMA (Exponential Moving Average) ──────────────────────────────────
    use_ema: bool = True
    ema_decay: float = 0.999
    
    # ── Early Stopping ────────────────────────────────────────────────────
    patience: int = 5
    min_delta: float = 1e-4