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

"""
Pre-training sanity check for AXIOM-1.

MUST be run before every training phase. No exceptions (Rule 1).
All 5 checks must pass. If any fails, diagnose before training.

Checks:
  1. Forward pass: no NaN/Inf in logits
  2. Initial loss ≈ log(132000) ≈ 11.79 (within ±3.0)
  3. Backward pass: finite gradients, grad_norm < 100
  4. Parameter count: 60M–110M
  5. Loop step embedding std < 0.02
  6. Loop gradient flow: gradients reach loop_step_embed
"""

import math

import torch
import torch.nn.functional as F

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


def run_sanity_check(
    model: AXIOM1,
    cfg: AXIOM1Config,
    device: str = 'cuda'
) -> bool:
    """
    Must pass ALL checks before training begins.
    If any check fails: fix the problem before proceeding.

    Args:
        model: AXIOM1 model instance
        cfg: AXIOM1Config instance
        device: Device to run checks on ('cuda' or 'cpu')

    Returns:
        True if all checks pass, raises AssertionError otherwise.
    """
    print("=" * 50)
    print(f"AXIOM-1 Sanity Check | {cfg.model_id}")
    print("=" * 50)

    model = model.to(device).train()
    ids = torch.randint(0, cfg.vocab_size, (2, 64), device=device)

    # CHECK 1: Forward pass — no NaN/Inf
    logits = model(ids)
    assert not torch.isnan(logits).any(), "FAIL: NaN in logits"
    assert not torch.isinf(logits).any(), "FAIL: Inf in logits"
    print("[PASS] Forward pass — no NaN/Inf")

    # CHECK 2: Initial loss near log(vocab_size)
    loss = F.cross_entropy(
        logits[:, :-1].reshape(-1, cfg.vocab_size),
        ids[:, 1:].reshape(-1)
    )
    expected_loss = math.log(cfg.vocab_size)   # ≈ 11.79 for vocab=132000
    assert abs(loss.item() - expected_loss) < 3.0, \
        f"FAIL: Loss {loss.item():.2f} too far from {expected_loss:.2f}"
    print(f"[PASS] Initial loss: {loss.item():.3f} (expected ≈ {expected_loss:.3f})")

    # CHECK 3: Backward — all gradients finite
    loss.backward()
    bad_grads = [(n, p) for n, p in model.named_parameters()
                 if p.grad is not None and (torch.isnan(p.grad).any()
                                            or torch.isinf(p.grad).any())]
    assert len(bad_grads) == 0, f"FAIL: Bad gradients in {[n for n, _ in bad_grads]}"
    grad_norm = sum(p.grad.norm()**2 for p in model.parameters()
                    if p.grad is not None) ** 0.5
    assert grad_norm.item() < 100.0, f"FAIL: Grad norm {grad_norm:.1f} too large"
    print(f"[PASS] Backward — grad_norm={grad_norm.item():.3f}")

    # CHECK 4: Parameter count
    # With 132K vocab, embedding table alone is 132K×512 = 67.6M params.
    # LoopLLM has ~85M unique trainable params (1 shared block + large embeddings).
    # Accept 60M–110M to account for weight tying deduplication.
    n_params = model.num_params()
    assert 60e6 < n_params < 110e6, f"FAIL: Param count {n_params/1e6:.1f}M out of range"
    print(f"[PASS] Parameters: {n_params/1e6:.1f}M unique (effective ~130M via 12-loop reuse)")

    # CHECK 5: Loop step embeddings — small std
    lse_std = model.loop_step_embed.weight.std().item()
    assert lse_std < 0.02, f"FAIL: LSE std {lse_std:.4f} too large"
    print(f"[PASS] Loop step embed std: {lse_std:.5f}")

    # CHECK 6: Loop gradient flow — verify gradients reach all loop iterations
    model.zero_grad()
    ids_6 = torch.randint(0, cfg.vocab_size, (1, 16), device=device)
    _, intermediates_6 = model(ids_6, return_intermediates=True)
    # Backprop from the first intermediate (loop step 0)
    intermediates_6[0].sum().backward()
    lse_grad = model.loop_step_embed.weight.grad
    assert lse_grad is not None, "FAIL: No gradient on loop_step_embed"
    has_grad = lse_grad.abs().sum() > 0
    assert has_grad, "FAIL: loop_step_embed gradient is all zeros — loop is disconnected"
    print(f"[PASS] Loop gradient flow — loop_step_embed grad norm: {lse_grad.norm().item():.6f}")

    model.zero_grad()
    print("=" * 50)
    print("All 6 checks PASSED. Safe to begin training.")
    print("=" * 50)
    return True


if __name__ == "__main__":
    # Standalone execution for quick validation
    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg = AXIOM1Config()
    model = AXIOM1(cfg)
    run_sanity_check(model, cfg, device=device)
