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

"""
AXIOM-1 Full Model — EGen Core LoopLLM Base.

Single shared block applied num_loops times to produce output.
Loop Step Embeddings differentiate each iteration.
Weight tying between token_embed and lm_head is mandatory.

Initialization is critical:
  - Output projections (o_proj, down) scaled by 1/√(2 × num_loops)
  - Loop step embeddings initialized at std=0.005 (small — must not disrupt h)
  - RMSNorm weights initialized to ones
  - All biases initialized to zeros
"""

import math

import torch
import torch.nn as nn

from axiom1.model.components import RMSNorm
from axiom1.model.block import LoopBlock
from axiom1.model.config import AXIOM1Config


class AXIOM1(nn.Module):
    """
    AXIOM-1: EGen Core LoopLLM Base Model.
    EGen Labs | ErebusTN

    Single shared block applied num_loops times to produce output.
    Loop Step Embeddings differentiate each iteration.
    Weight tying between token_embed and lm_head.

    Architecture:
        input_ids → token_embed × √d_model → pre_loop_norm
        → [loop_step_embed + shared_block] × num_loops
        → final_norm → lm_head (weight-tied) → logits
    """

    def __init__(self, cfg: AXIOM1Config) -> None:
        super().__init__()
        self.cfg = cfg

        # --- Core modules ---
        self.token_embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.loop_step_embed = nn.Embedding(cfg.max_loops, cfg.d_model)
        self.pre_loop_norm = RMSNorm(cfg.d_model)
        self.block = LoopBlock(cfg)
        self.final_norm = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)

        # --- Weight tying (Rule 12: mandatory) ---
        self.lm_head.weight = self.token_embed.weight

        # --- Initialize ---
        self._init_weights()

    def _init_weights(self) -> None:
        """
        Initialization is critical for LoopLLM stability.

        Output projections (o_proj, down) must be scaled by 1/sqrt(2 × num_loops)
        to prevent activation magnitude explosion across iterations.

        Rules:
          token_embed:       normal(std=0.02)
          loop_step_embed:   normal(std=0.005) — small, must not disrupt h
          o_proj, down:      normal(std=0.02 / sqrt(2 × 12)) — SCALED DOWN
          all other weights: normal(std=0.02)
          RMSNorm weights:   ones
          all biases:        zeros
        """
        std = self.cfg.init_std
        N = self.cfg.num_loops

        nn.init.normal_(self.token_embed.weight, std=std)
        nn.init.normal_(self.loop_step_embed.weight, std=0.005)  # small — don't disrupt h

        for name, p in self.block.named_parameters():
            if p.ndim >= 2 and 'weight' in name:
                if any(tag in name for tag in ['o_proj', 'down']):
                    # Output projections: scale down to prevent compounding
                    nn.init.normal_(p, std=std / math.sqrt(2 * N))
                else:
                    nn.init.normal_(p, std=std)
            elif 'bias' in name:
                nn.init.zeros_(p)
            elif 'weight' in name and p.ndim == 1:
                # RMSNorm weights
                nn.init.ones_(p)

    def forward(
        self,
        input_ids: torch.Tensor,
        mask: torch.Tensor | None = None,
        kv_cache: object | None = None,
        return_intermediates: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
        """
        Full forward pass through the AXIOM-1 model.

        Args:
            input_ids: Token IDs (batch, seq_len)
            mask: Optional causal mask (bool, True=masked)
            kv_cache: PerStepKVCache instance for inference, or None
            return_intermediates: If True, return list of hidden states per loop step

        Returns:
            logits: (batch, seq_len, vocab_size)
            [intermediates]: Optional list of (batch, seq_len, d_model) per loop step
        """
        # Embed + scale
        h = self.token_embed(input_ids) * (self.cfg.d_model ** 0.5)
        h = self.pre_loop_norm(h)

        intermediates: list[torch.Tensor] = []
        loop_contributions: list[float] = []

        for step in range(self.cfg.num_loops):
            prev_h = h.detach().clone()

            # Add loop step signal (differentiates each iteration)
            lse = self.loop_step_embed(
                torch.tensor(step, device=h.device)
            ).unsqueeze(0).unsqueeze(0)
            h = h + lse

            # Apply shared block
            h = self.block(h, mask=mask, kv_cache=kv_cache, step_idx=step)

            # Mid-loop re-normalization: prevent hidden state magnitude drift
            # Applied every mid_loop_renorm_interval steps (default: every 4)
            if (step + 1) % self.cfg.mid_loop_renorm_interval == 0 and step < self.cfg.num_loops - 1:
                h = self.pre_loop_norm(h)

            # Track loop contributions (detached — monitoring only, no grad impact)
            with torch.no_grad():
                delta = (h - prev_h).norm().item()
                loop_contributions.append(delta)

            if return_intermediates:
                intermediates.append(h)

        # Store latest contributions for external access (e.g., logging)
        self._last_loop_contributions = loop_contributions

        # Final output
        logits = self.lm_head(self.final_norm(h))

        if return_intermediates:
            return logits, intermediates
        return logits

    def num_params(self) -> int:
        """
        Return unique parameter count.

        PyTorch's parameters() already deduplicates tied weights
        (lm_head.weight IS token_embed.weight — same tensor), so
        total already excludes the duplicate. No further subtraction needed.

        Note: AXIOM-1 has ~18M unique params. The effective parameter
        count across 12 loop iterations is much higher (~130M equivalent
        depth), but the actual trainable parameters are ~18M due to
        weight sharing — this is the core efficiency of LoopLLM.
        """
        return sum(p.numel() for p in self.parameters())

    def __repr__(self) -> str:
        return (
            f"AXIOM-1 | EGen Labs\n"
            f"  d_model={self.cfg.d_model}, loops={self.cfg.num_loops}, "
            f"vocab={self.cfg.vocab_size}\n"
            f"  GQA: {self.cfg.n_q_heads}q/{self.cfg.n_kv_heads}kv, "
            f"SwiGLU d_ff={self.cfg.d_ff}\n"
            f"  Params: {self.num_params()/1e6:.1f}M"
        )
