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

"""
Tests for LoopBlock and AXIOM1 model:
  - Forward pass shapes
  - Parameter counts
  - Weight initialization (output projection scaling, LSE std)
  - Weight tying
  - Various batch/seq sizes
  - NaN/Inf checks
"""

import math
import torch
import pytest

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.model.block import LoopBlock


class TestLoopBlock:
    """Tests for the shared LoopBlock."""

    def test_output_shape(self, cfg: AXIOM1Config) -> None:
        block = LoopBlock(cfg)
        x = torch.randn(2, 16, cfg.d_model)
        out = block(x)
        assert out.shape == x.shape, \
            f"LoopBlock output shape {out.shape} != input {x.shape}"

    def test_no_nan(self, cfg: AXIOM1Config) -> None:
        block = LoopBlock(cfg)
        x = torch.randn(2, 16, cfg.d_model)
        out = block(x)
        assert not torch.isnan(out).any(), "LoopBlock output contains NaN"
        assert not torch.isinf(out).any(), "LoopBlock output contains Inf"

    def test_pre_ln_placement(self, cfg: AXIOM1Config) -> None:
        """Rule 4: Pre-LN only — norms before sub-layers."""
        block = LoopBlock(cfg)
        assert hasattr(block, 'ln1'), "LoopBlock missing ln1 (pre-attention norm)"
        assert hasattr(block, 'ln2'), "LoopBlock missing ln2 (pre-FFN norm)"

    def test_no_bias_in_block(self, cfg: AXIOM1Config) -> None:
        """No bias in any Linear layer inside LoopBlock."""
        block = LoopBlock(cfg)
        for name, module in block.named_modules():
            if isinstance(module, torch.nn.Linear):
                assert module.bias is None, \
                    f"LoopBlock Linear '{name}' has bias — must be bias=False"


class TestAXIOM1Model:
    """Tests for the full AXIOM1 model."""

    def test_output_shape(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 32))
        logits = small_model(ids)
        assert logits.shape == (2, 32, small_cfg.vocab_size), \
            f"Expected (2, 32, {small_cfg.vocab_size}), got {logits.shape}"

    def test_various_batch_seq(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        """Test different batch sizes and sequence lengths."""
        for batch, seq in [(1, 8), (4, 16), (2, 64)]:
            ids = torch.randint(0, small_cfg.vocab_size, (batch, seq))
            logits = small_model(ids)
            assert logits.shape == (batch, seq, small_cfg.vocab_size), \
                f"Shape mismatch for batch={batch}, seq={seq}: {logits.shape}"

    def test_no_nan_inf(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 32))
        logits = small_model(ids)
        assert not torch.isnan(logits).any(), "Model output contains NaN"
        assert not torch.isinf(logits).any(), "Model output contains Inf"

    def test_param_count_range(self, model: AXIOM1) -> None:
        """Unique params should be ~18M for LoopLLM."""
        n = model.num_params()
        assert 15e6 < n < 25e6, \
            f"Expected 15M-25M unique params, got {n/1e6:.1f}M"

    def test_weight_tying(self, model: AXIOM1) -> None:
        """Rule 12: lm_head.weight must be token_embed.weight."""
        assert model.lm_head.weight is model.token_embed.weight, \
            "Weight tying failed: lm_head.weight is not token_embed.weight"

    def test_single_block_instance(self, model: AXIOM1) -> None:
        """Rule 11: Only one LoopBlock instance, reused N times."""
        blocks = [m for m in model.modules() if isinstance(m, LoopBlock)]
        assert len(blocks) == 1, \
            f"Expected 1 LoopBlock instance, found {len(blocks)}"

    def test_return_intermediates(self, small_model: AXIOM1, small_cfg: AXIOM1Config) -> None:
        ids = torch.randint(0, small_cfg.vocab_size, (2, 16))
        logits, intermediates = small_model(ids, return_intermediates=True)
        assert len(intermediates) == small_cfg.num_loops, \
            f"Expected {small_cfg.num_loops} intermediates, got {len(intermediates)}"
        for i, h in enumerate(intermediates):
            assert h.shape == (2, 16, small_cfg.d_model), \
                f"Intermediate {i} shape mismatch: {h.shape}"


class TestInitialization:
    """Tests for weight initialization rules."""

    def test_output_proj_scaling(self, cfg: AXIOM1Config) -> None:
        """Rule 5: o_proj and down must be scaled by 1/sqrt(2*N)."""
        model = AXIOM1(cfg)
        expected_std = cfg.init_std / math.sqrt(2 * cfg.num_loops)

        o_proj_std = model.block.attn.o_proj.weight.std().item()
        down_std = model.block.ffn.down.weight.std().item()

        # Allow ±50% tolerance for random initialization
        assert abs(o_proj_std - expected_std) < expected_std * 0.5, \
            f"o_proj std={o_proj_std:.4f}, expected ~{expected_std:.4f}"
        assert abs(down_std - expected_std) < expected_std * 0.5, \
            f"down std={down_std:.4f}, expected ~{expected_std:.4f}"

    def test_loop_step_embed_small_std(self, model: AXIOM1) -> None:
        """Rule 6: loop_step_embed init std <= 0.01."""
        lse_std = model.loop_step_embed.weight.std().item()
        assert lse_std < 0.02, \
            f"Loop step embed std={lse_std:.4f}, should be < 0.02"

    def test_token_embed_std(self, model: AXIOM1) -> None:
        """token_embed should be initialized at std=0.02."""
        std = model.token_embed.weight.std().item()
        assert abs(std - 0.02) < 0.01, \
            f"token_embed std={std:.4f}, expected ~0.02"

    def test_rmsnorm_weights_ones(self, model: AXIOM1) -> None:
        """RMSNorm weights should be initialized to ones."""
        assert torch.allclose(model.pre_loop_norm.weight, torch.ones_like(model.pre_loop_norm.weight)), \
            "pre_loop_norm weight should be ones"
        assert torch.allclose(model.final_norm.weight, torch.ones_like(model.final_norm.weight)), \
            "final_norm weight should be ones"
