# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/model/config.py
# ─────────────────────────────────────────────────────────

"""
AXIOM-1 Configuration — Canonical hyperparameter dataclass.

All values are frozen per the AXIOM-1 specification.
Do not change any default without explicit approval from ErebusTN.
"""

from dataclasses import dataclass


@dataclass
class AXIOM1Config:
    """
    Complete configuration for the AXIOM-1 LoopLLM model.

    This dataclass holds every hyperparameter, training setting, and inference
    default for AXIOM-1. All values match the canonical specification in
    basis/AXIOM1_Context_File.md exactly.

    Architecture: 1 shared LoopBlock × 12 iterations, ~130M params after tying.
    """

    # ── Identity ──────────────────────────────
    model_name:   str = "AXIOM-1"
    model_id:     str = "EGen-Core/AXIOM-1"
    lab:          str = "EGen Labs"
    author:       str = "ErebusTN"
    version:      str = "1.0.1"

    # ── Architecture ──────────────────────────
    d_model:      int   = 512
    n_q_heads:    int   = 8
    n_kv_heads:   int   = 2
    num_loops:    int   = 12
    max_loops:    int   = 16
    vocab_size:   int   = 132000
    max_seq_len:  int   = 2048
    drop_path:    float = 0.10
    init_std:     float = 0.02
    rope_base:    int   = 10000
    weight_tying: bool  = True

    # Adaptive Depth (Phase 3 only — False during Phase 1 & 2)
    use_adaptive_depth: bool  = False
    halt_threshold:     float = 0.99
    ponder_cost_tau:    float = 0.01

    # Special tokens (LLaMA-compatible)
    pad_token_id: int = 0
    bos_token_id: int = 1
    eos_token_id: int = 2

    # ── Derived properties ─────────────────────
    @property
    def head_dim(self) -> int:
        """Head dimension: d_model / n_q_heads = 64."""
        return self.d_model // self.n_q_heads

    @property
    def d_ff(self) -> int:
        """FFN intermediate size: int(d_model × 8/3) = 1365 for SwiGLU."""
        return int(self.d_model * 8 / 3)

    # ── Training — Phase 1 (Warm Start) ───────
    phase1_loops:        int   = 4
    phase1_seq_len:      int   = 512
    phase1_peak_lr:      float = 3e-4
    phase1_min_lr:       float = 3e-5
    phase1_warmup_steps: int   = 1000
    phase1_total_steps:  int   = 50_000
    phase1_batch_size:   int   = 32
    phase1_grad_accum:   int   = 4

    # ── Training — Phase 2 (Full Depth) ───────
    phase2_loops:        int   = 12
    phase2_seq_len:      int   = 1024
    phase2_peak_lr:      float = 1e-4
    phase2_min_lr:       float = 1e-5
    phase2_warmup_steps: int   = 500
    phase2_total_steps:  int   = 300_000
    phase2_batch_size:   int   = 32
    phase2_grad_accum:   int   = 8

    # ── Training — Phase 3 (ADC, Optional) ────
    phase3_loops:        int   = 16
    phase3_seq_len:      int   = 2048
    phase3_peak_lr:      float = 3e-5
    phase3_min_lr:       float = 3e-6
    phase3_warmup_steps: int   = 200
    phase3_total_steps:  int   = 50_000
    phase3_batch_size:   int   = 16
    phase3_grad_accum:   int   = 8

    # ── Shared Training Settings ──────────────
    weight_decay:           float = 0.1
    grad_clip:              float = 1.0
    optimizer_beta1:        float = 0.9
    optimizer_beta2:        float = 0.95    # NOT 0.999 — recurrent models need lower β2
    optimizer_eps:          float = 1e-8
    mixed_precision:        str   = "bfloat16"   # NEVER float16
    consistency_weight:     float = 0.25
    consistency_start_loop: int   = 6            # num_loops // 2
    trunc_bptt_prob:        float = 0.05

    # ── Stability ─────────────────────────────────
    mid_loop_renorm_interval: int   = 4       # Re-normalize hidden state every N loops
    consistency_warmup_steps: int   = 5000    # Steps to ramp consistency weight 0→full
    loop_collapse_threshold:  float = 0.999   # cos_sim above this = collapse risk

    # ── Inference ─────────────────────────────
    inference_dtype:    str   = "bfloat16"
    use_kv_cache:       bool  = True
    use_cpu_offload:    bool  = False
    temperature:        float = 0.8
    top_p:              float = 0.9
    top_k:              int   = 50
    max_new_tokens:     int   = 512
    repetition_penalty: float = 1.1
