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

"""
Validation perplexity tracker for AXIOM-1.

Computes CE loss and perplexity on a held-out validation set.
Always uses bfloat16 autocast (never float16).
"""

import math

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

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


@torch.no_grad()
def evaluate(
    model: AXIOM1,
    val_dataloader: DataLoader,
    cfg: AXIOM1Config,
    max_batches: int = 50,
    device: str = 'cuda'
) -> float:
    """
    Compute average validation loss over a subset of the validation set.

    Args:
        model: AXIOM1 model instance (will be set to eval mode)
        val_dataloader: DataLoader yielding dicts with 'input_ids'
        cfg: AXIOM1Config instance
        max_batches: Maximum number of batches to evaluate
        device: Device to evaluate on

    Returns:
        Average cross-entropy loss (float). Use math.exp(loss) for perplexity.
    """
    model.eval()
    total_loss = 0.0
    total_tokens = 0

    for i, batch in enumerate(val_dataloader):
        if i >= max_batches:
            break

        input_ids = batch['input_ids'].to(device)

        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
            logits = model(input_ids)

        shift_logits = logits[:, :-1, :].contiguous()
        shift_labels = input_ids[:, 1:].contiguous()

        loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            reduction='sum'
        )

        n_tokens = shift_labels.numel()
        total_loss += loss.item()
        total_tokens += n_tokens

    avg_loss = total_loss / max(total_tokens, 1)
    return avg_loss


def compute_perplexity(loss: float) -> float:
    """Convert average CE loss to perplexity."""
    return math.exp(min(loss, 100.0))  # cap to prevent overflow


def log_validation(
    model: AXIOM1,
    val_dataloader: DataLoader,
    cfg: AXIOM1Config,
    step: int,
    device: str = 'cuda'
) -> tuple[float, float]:
    """
    Run validation and print results in AXIOM-1 format.

    Args:
        model: AXIOM1 model instance
        val_dataloader: DataLoader yielding dicts with 'input_ids'
        cfg: AXIOM1Config instance
        step: Current training step
        device: Device to evaluate on

    Returns:
        Tuple of (val_loss, val_perplexity)
    """
    val_loss = evaluate(model, val_dataloader, cfg, device=device)
    val_ppl = compute_perplexity(val_loss)
    print(f"[AXIOM-1] VAL step={step} | loss={val_loss:.4f} | ppl={val_ppl:.1f}")
    model.train()
    return val_loss, val_ppl
