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

"""
End-to-end training test.

Runs a minimal Phase 1 training loop (4 loops, a few steps)
to verify nothing crashes under realistic conditions.
"""

import torch
import pytest

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.training.optimizer import build_optimizer, get_lr
from axiom1.training.loss import compute_ce_loss


class TestEndToEnd:
    """Minimal end-to-end training loop."""

    def test_phase1_mini_loop(self) -> None:
        """Run 5 training steps with Phase 1 config (4 loops)."""
        cfg = AXIOM1Config()
        cfg.num_loops = 4
        cfg.max_seq_len = 64

        model = AXIOM1(cfg)
        optimizer = build_optimizer(model, cfg)

        model.train()
        initial_loss = None

        for step in range(5):
            ids = torch.randint(0, cfg.vocab_size, (2, 64))

            logits = model(ids)
            loss = compute_ce_loss(logits, ids)

            if initial_loss is None:
                initial_loss = loss.item()

            loss.backward()

            # Rule 7: gradient clipping
            grad_norm = torch.nn.utils.clip_grad_norm_(
                model.parameters(), cfg.grad_clip
            )

            # LR schedule
            lr = get_lr(step, warmup_steps=10, total_steps=100,
                        peak_lr=3e-4, min_lr=3e-5)
            for pg in optimizer.param_groups:
                pg['lr'] = lr

            optimizer.step()
            optimizer.zero_grad()

            assert torch.isfinite(loss), f"Loss is not finite at step {step}"
            assert torch.isfinite(grad_norm), f"Grad norm not finite at step {step}"

        final_loss = loss.item()
        assert torch.isfinite(torch.tensor(final_loss)), "Final loss not finite"
        print(f"  E2E: initial_loss={initial_loss:.2f}, "
              f"final_loss={final_loss:.2f}, steps=5")

    def test_phase1_with_consistency(self) -> None:
        """Test a step with consistency loss enabled."""
        from axiom1.training.loss import compute_consistency_loss

        cfg = AXIOM1Config()
        cfg.num_loops = 4
        model = AXIOM1(cfg)
        optimizer = build_optimizer(model, cfg)
        model.train()

        ids = torch.randint(0, cfg.vocab_size, (2, 32))
        logits, intermediates = model(ids, return_intermediates=True)

        ce_loss = compute_ce_loss(logits, ids)
        cons_loss = compute_consistency_loss(
            model, intermediates, ids, start_loop=2
        )
        total_loss = ce_loss + 0.25 * cons_loss

        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()

        assert torch.isfinite(total_loss), \
            f"Total loss (CE + consistency) not finite: {total_loss.item()}"
