"""Small deterministic before/after evaluator for causal language models."""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Iterable, Mapping, Optional

import torch


@dataclass(frozen=True)
class LossMetrics:
    nll: float
    perplexity: float
    predicted_tokens: int


@torch.inference_mode()
def evaluate_causal_lm(
    model,
    batches: Iterable[Mapping[str, torch.Tensor]],
    *,
    device: Optional[torch.device | str] = None,
) -> LossMetrics:
    """Evaluate an HF-style causal LM on already-tokenized frozen batches."""
    total_nll = 0.0
    total_tokens = 0
    was_training = model.training
    model.eval()
    try:
        for original in batches:
            batch = {
                key: value.to(device) if device is not None else value
                for key, value in original.items()
            }
            labels = batch.get("labels")
            if labels is None:
                labels = batch["input_ids"].clone()
                batch["labels"] = labels
            predicted = labels[..., 1:]
            token_count = int((predicted != -100).sum().item())
            if token_count == 0:
                continue
            output = model(**batch)
            loss = output.loss if hasattr(output, "loss") else output[0]
            total_nll += float(loss.detach().float().item()) * token_count
            total_tokens += token_count
    finally:
        model.train(was_training)
    if total_tokens == 0:
        raise ValueError("evaluation batches contain no predicted tokens")
    mean_nll = total_nll / total_tokens
    return LossMetrics(mean_nll, math.exp(mean_nll), total_tokens)