"""
Evaluation Metrics for Q-TensorFormer (`src/metrics.py`).

Provides:
  - `evaluate_model`: Computes token-level cross-entropy loss, perplexity, and top-1 accuracy
    with proper ignore_index handling for `<pad>` tokens (`pad_token_id=0`).
"""

import math
import time
from typing import Any, Dict, Optional

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


@torch.no_grad()
def evaluate_model(
    model: nn.Module,
    dataloader: DataLoader,
    device: str = "cpu",
    pad_token_id: int = 0,
    preset: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Evaluate a causal language model on a DataLoader yielding `(x, y)` or `(x, y, mask)`.
    """
    was_training = model.training
    model.eval()
    model.to(device)

    total_nll = 0.0
    total_tokens = 0
    correct_tokens = 0
    total_batches = 0
    t0 = time.perf_counter()

    for batch in dataloader:
        if len(batch) == 3:
            x, y, attn_mask = batch
            attn_mask = attn_mask.to(device)
        else:
            x, y = batch
            attn_mask = (x != pad_token_id).long().to(device)

        x = x.to(device)
        y = y.to(device)

        if preset is not None and hasattr(model, "set_preset"):
            out = model(x, attention_mask=attn_mask, preset=preset)
        else:
            try:
                out = model(x, attention_mask=attn_mask)
            except TypeError:
                out = model(x)

        logits = out[0] if isinstance(out, tuple) else (out.logits if hasattr(out, "logits") else out)
        vocab_size = logits.size(-1)

        valid_mask = y != pad_token_id
        n_valid = int(valid_mask.sum().item())
        if n_valid == 0:
            continue

        loss = F.cross_entropy(
            logits.reshape(-1, vocab_size),
            y.reshape(-1),
            ignore_index=pad_token_id,
            reduction="sum",
        )
        preds = logits.argmax(dim=-1)
        correct = int(((preds == y) & valid_mask).sum().item())

        total_nll += float(loss.item())
        total_tokens += n_valid
        correct_tokens += correct
        total_batches += 1

    elapsed_s = max(1e-6, time.perf_counter() - t0)
    if was_training:
        model.train()

    avg_loss = total_nll / max(1, total_tokens)
    ppl = math.exp(min(avg_loss, 20.0))
    acc = correct_tokens / max(1, total_tokens)
    total_params = sum(p.numel() for p in model.parameters())

    return {
        "loss": round(avg_loss, 6),
        "test_loss": round(avg_loss, 6),
        "perplexity": round(ppl, 4),
        "test_ppl": round(ppl, 4),
        "accuracy": round(acc, 4),
        "total_params": total_params,
        "total_tokens": total_tokens,
        "total_batches": total_batches,
        "eval_time_s": round(elapsed_s, 4),
        "tokens_per_sec": round(total_tokens / elapsed_s, 2),
        "classification": "MEASURED",
    }