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

"""
LoopBlock — THE core unit of AXIOM-1.

One instance exists. It is applied 12 times (num_loops) per forward pass.
The same weights are used at every iteration — this is the fundamental
principle of the LoopLLM architecture.

Pre-LN placement is mandatory: normalize BEFORE each sub-layer.
DropPath is the ONLY regularization — no standard dropout inside the loop.
"""

import torch
import torch.nn as nn

from axiom1.model.components import RMSNorm, GroupedQueryAttention, SwiGLU_FFN, DropPath


class LoopBlock(nn.Module):
    """
    THE core unit of AXIOM-1. One instance. Applied 12 times.

    Architecture per iteration:
        h = h + DropPath(Attn(RMSNorm(h)))   ← Pre-LN attention
        h = h + DropPath(FFN(RMSNorm(h)))    ← Pre-LN FFN

    Pre-LN placement is mandatory — see architecture spec.
    No bias in any Linear layer — prevents activation drift across loops.
    """

    def __init__(self, cfg: object) -> None:
        super().__init__()
        self.ln1 = RMSNorm(cfg.d_model)
        self.attn = GroupedQueryAttention(cfg.d_model, cfg.n_q_heads, cfg.n_kv_heads)
        self.ln2 = RMSNorm(cfg.d_model)
        self.ffn = SwiGLU_FFN(cfg.d_model, cfg.d_ff)
        self.dp_attn = DropPath(cfg.drop_path)
        self.dp_ffn = DropPath(cfg.drop_path)

    def forward(
        self,
        h: torch.Tensor,
        mask: torch.Tensor | None = None,
        kv_cache: object | None = None,
        step_idx: int | None = None
    ) -> torch.Tensor:
        """
        Apply one iteration of the shared block.

        Args:
            h: Hidden state (batch, seq_len, d_model)
            mask: Optional causal mask (bool, True=masked)
            kv_cache: PerStepKVCache instance for inference, or None
            step_idx: Current loop iteration index (for KV cache)

        Returns:
            Updated hidden state (batch, seq_len, d_model)
        """
        # Attention sub-layer (Pre-LN)
        h = h + self.dp_attn(
            self.attn(self.ln1(h), mask=mask, kv_cache=kv_cache, step_idx=step_idx)
        )
        # FFN sub-layer (Pre-LN)
        h = h + self.dp_ffn(
            self.ffn(self.ln2(h))
        )
        return h
