"""Reusable single-label classification metrics."""

from __future__ import annotations

import torch


_EPSILON = 1e-10


def _normalize_classification_tensors(output: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Validate and normalize classification logits and targets.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor broadcastable to shape ``[N]``.

    Returns:
        Tuple of normalized ``(output, target)`` tensors.
    """
    if output.ndim != 2:
        raise ValueError(f"Expected `output` with shape [N, C], got {tuple(output.shape)}")
    target = target.reshape(-1).to(device=output.device, dtype=torch.long)
    if target.numel() != output.size(0):
        raise ValueError("`target` size must match output batch dimension")
    return output, target


def calculate_metrics(output: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Compute one-vs-rest TP/TN/FP/FN vectors for each class.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Tuple of tensors ``(tp, tn, fp, fn)`` for every class.
    """
    output, target = _normalize_classification_tensors(output, target)
    pred = torch.argmax(output, dim=1)
    num_classes = output.size(1)

    class_ids = torch.arange(num_classes, device=pred.device).view(1, -1)
    pred_pos = pred.view(-1, 1).eq(class_ids)
    target_pos = target.view(-1, 1).eq(class_ids)

    tp = (pred_pos & target_pos).sum(dim=0).to(dtype=torch.float32, device="cpu")
    fp = (pred_pos & ~target_pos).sum(dim=0).to(dtype=torch.float32, device="cpu")
    fn = (~pred_pos & target_pos).sum(dim=0).to(dtype=torch.float32, device="cpu")
    tn = (~pred_pos & ~target_pos).sum(dim=0).to(dtype=torch.float32, device="cpu")

    return tp, tn, fp, fn


def accuracy_per_class(output: torch.Tensor, target: torch.Tensor) -> dict[int, float]:
    """Compute per-class one-vs-rest accuracy values.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Mapping of class id to one-vs-rest accuracy.
    """
    tp, tn, fp, fn = calculate_metrics(output, target)
    per_class_accuracy: dict[int, float] = {}
    for class_id in range(output.size(1)):
        denom = tp[class_id] + tn[class_id] + fp[class_id] + fn[class_id]
        if denom.item() <= 0:
            per_class_accuracy[class_id] = 0.0
            continue
        per_class_accuracy[class_id] = float(((tp[class_id] + tn[class_id]) / denom).item())
    return per_class_accuracy


def precision(output: torch.Tensor, target: torch.Tensor) -> dict[int, float]:
    """Compute per-class precision values.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Mapping of class id to precision.
    """
    tp, _, fp, _ = calculate_metrics(output, target)
    result: dict[int, float] = {}
    for class_id in range(output.size(1)):
        denom = tp[class_id] + fp[class_id]
        if denom.item() <= 0:
            result[class_id] = 0.0
            continue
        result[class_id] = float((tp[class_id] / denom).item())
    return result


def recall(output: torch.Tensor, target: torch.Tensor) -> dict[int, float]:
    """Compute per-class recall values.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Mapping of class id to recall.
    """
    tp, _, _, fn = calculate_metrics(output, target)
    result: dict[int, float] = {}
    for class_id in range(output.size(1)):
        denom = tp[class_id] + fn[class_id]
        if denom.item() <= 0:
            result[class_id] = 0.0
            continue
        result[class_id] = float((tp[class_id] / denom).item())
    return result


def f1_score_per_class(output: torch.Tensor, target: torch.Tensor) -> dict[int, float]:
    """Compute per-class F1-score values.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Mapping of class id to F1-score.
    """
    per_class_precision = precision(output, target)
    per_class_recall = recall(output, target)

    result: dict[int, float] = {}
    for class_id in range(output.size(1)):
        p_val = per_class_precision[class_id]
        r_val = per_class_recall[class_id]
        if p_val + r_val <= 0:
            result[class_id] = 0.0
            continue
        result[class_id] = float((2.0 * p_val * r_val) / (p_val + r_val + _EPSILON))
    return result


def specificity(output: torch.Tensor, target: torch.Tensor) -> dict[int, float]:
    """Compute per-class specificity (true negative rate).

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Mapping of class id to specificity.
    """
    _, tn, fp, _ = calculate_metrics(output, target)
    result: dict[int, float] = {}
    for class_id in range(output.size(1)):
        denom = tn[class_id] + fp[class_id]
        if denom.item() <= 0:
            result[class_id] = 0.0
            continue
        result[class_id] = float((tn[class_id] / denom).item())
    return result


def specificity_all(output: torch.Tensor, target: torch.Tensor) -> float:
    """Compute overall specificity across all classes.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Overall specificity value.
    """
    _, tn, fp, _ = calculate_metrics(output, target)
    total_tn = tn.sum().item()
    total_fp = fp.sum().item()
    return float(total_tn / (total_tn + total_fp + _EPSILON))


def confusion_matrix_per_class(output: torch.Tensor, target: torch.Tensor) -> dict[int, list[list[float]]]:
    """Return one-vs-rest confusion matrix for each class.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Mapping of class id to ``[[tp, fp], [fn, tn]]`` matrix.
    """
    tp, tn, fp, fn = calculate_metrics(output, target)
    result: dict[int, list[list[float]]] = {}
    for class_id in range(output.size(1)):
        result[class_id] = [
            [float(tp[class_id].item()), float(fp[class_id].item())],
            [float(fn[class_id].item()), float(tn[class_id].item())],
        ]
    return result


def confusion_matrix(output: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Return full multiclass confusion matrix with shape ``[C, C]``.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.

    Returns:
        Confusion matrix tensor where rows are true labels and columns are predicted labels.
    """
    output, target = _normalize_classification_tensors(output, target)
    pred = torch.argmax(output, dim=1).to(dtype=torch.long, device="cpu")
    target_cpu = target.to(dtype=torch.long, device="cpu")
    num_classes = output.size(1)

    matrix = torch.zeros((num_classes, num_classes), dtype=torch.int64)
    if target_cpu.numel() == 0:
        return matrix

    valid = (
        (target_cpu >= 0)
        & (target_cpu < num_classes)
        & (pred >= 0)
        & (pred < num_classes)
    )
    if not torch.any(valid):
        return matrix

    indices = target_cpu[valid] * num_classes + pred[valid]
    counts = torch.bincount(indices, minlength=num_classes * num_classes)
    return counts.reshape(num_classes, num_classes)
