"""Evaluation metrics for ablation studies."""

from __future__ import annotations

import math
from typing import Sequence

import torch
import torch.nn.functional as F
from sklearn.metrics import f1_score as sklearn_f1


def perplexity(logits: torch.Tensor, labels: torch.Tensor) -> float:
    """Compute perplexity from causal-LM logits and label token IDs.

    Args:
        logits: (batch, seq_len, vocab_size) — raw model output.
        labels: (batch, seq_len) — ground-truth token IDs (use -100 for padding).

    Returns:
        Scalar perplexity (lower is better).
    """
    if logits.dim() != 3:
        raise ValueError(f"logits must be 3D (batch, sequence, vocabulary), got {logits.dim()}D")
    if labels.dim() != 2:
        raise ValueError(f"labels must be 2D (batch, sequence), got {labels.dim()}D")
    if labels.dtype not in {
        torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8,
    }:
        raise ValueError("labels must contain integer token IDs")
    if logits.shape[:2] != labels.shape:
        raise ValueError(
            f"logits batch/sequence shape {tuple(logits.shape[:2])} does not match "
            f"labels shape {tuple(labels.shape)}",
        )
    if logits.shape[1] < 2:
        raise ValueError("perplexity requires at least two sequence positions")
    if not torch.isfinite(logits).all():
        raise ValueError("logits must contain only finite values")

    # Shift so that tokens < n predict n
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    valid_labels = shift_labels[shift_labels != -100]
    if valid_labels.numel() == 0:
        raise ValueError("perplexity is unavailable because every target label is ignored")
    if (valid_labels < 0).any() or (valid_labels >= logits.shape[-1]).any():
        raise ValueError("labels contain token IDs outside the logits vocabulary")

    loss = F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=-100,
        reduction="mean",
    )
    return math.exp(loss.item())


def accuracy(predictions: Sequence[int], references: Sequence[int]) -> float:
    """Simple accuracy."""
    if len(predictions) != len(references):
        raise ValueError(
            f"predictions and references must have equal length, got "
            f"{len(predictions)} and {len(references)}",
        )
    if len(predictions) == 0:
        return 0.0
    correct = sum(p == r for p, r in zip(predictions, references))
    return correct / len(predictions)


def f1_score_metric(
    predictions: Sequence[int],
    references: Sequence[int],
    average: str = "macro",
) -> float:
    """F1 score wrapper around sklearn."""
    if len(predictions) != len(references):
        raise ValueError(
            f"predictions and references must have equal length, got "
            f"{len(predictions)} and {len(references)}",
        )
    if len(predictions) == 0:
        return 0.0
    return float(sklearn_f1(references, predictions, average=average, zero_division=0))