from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional, List

@dataclass
class ModelConfig:
    # Core architecture
    vocab_size: int
    dim: int
    n_layers: int
    n_heads: int
    n_kv_heads: int
    ff_mult: float
    max_seq_len: int = 1024          # INCREASED: reasoning needs 1024+

    # Regularisation
    dropout: float = 0.0

    # Positional encoding
    rope_theta: float = 10000.0

    # Normalisation
    rms_norm_eps: float = 1e-5
    use_qk_norm: bool = False
    use_post_norm: bool = False

    # Embeddings
    tie_embeddings: bool = True

    # Attention
    use_flash_if_available: bool = True
    sliding_window_attention: bool = False
    window_size: int = 2048
    residual_scale: float = 1.0

    # MoE (Mixture of Experts) — future expansion
    use_moe: bool = False
    moe_every: int = 2
    n_experts: int = 8
    moe_top_k: int = 2
    aux_loss_alpha: float = 1e-2

    # MLA (Multi-head Latent Attention) — future expansion
    use_mla: bool = False
    mla_latent_dim: int = 512

    # ── SELF-LEARNING & EVOLUTION ─────────────────────────────
    # MOTHER CORE is a self-evolving sovereign AI system
    # These flags enable continual learning capabilities

    # Continual learning — model can continue learning after deployment
    continual_learning: bool = True

    # Self-correction — model checks and corrects its own outputs
    self_correction: bool = True

    # Knowledge consolidation — periodic replay of core facts
    knowledge_consolidation: bool = True

    # Adaptive learning rate — adjusts per-layer based on gradient norms
    adaptive_lr: bool = True

    # Memory augmentation — external memory for long-term knowledge
    use_external_memory: bool = False
    memory_size: int = 10000

    # Identity — sovereign AI metadata
    model_name: str = "MOTHER CORE"
    model_version: str = "v11"
    architect: str = "Christopher Kenna"
    organisation: str = "MediaStream AI Limited"
    sovereign: bool = True
    sovereign_nation: str = "United Kingdom"