"""Common metric helpers shared across classification and detection tasks."""

from __future__ import annotations

from typing import TYPE_CHECKING

import torch


if TYPE_CHECKING:
    from collections.abc import Sequence


_EPSILON = 1e-10


def _as_1d_long_tensor(values: torch.Tensor) -> torch.Tensor:
    """Convert tensor-like input to flattened CPU ``int64`` tensor.

    Args:
        values: Tensor-like values to normalize.

    Returns:
        Flattened CPU tensor with ``torch.int64`` dtype.
    """
    if not isinstance(values, torch.Tensor):
        values = torch.as_tensor(values)
    if values.ndim == 0:
        values = values.unsqueeze(0)
    return values.detach().to(device="cpu", dtype=torch.long).reshape(-1)


def accuracy(output: torch.Tensor, target: torch.Tensor, topk: Sequence[int] = (1,)) -> list[torch.Tensor]:
    """Compute top-k accuracy values for logits and labels.

    Args:
        output: Classification logits tensor with shape ``[N, C]``.
        target: Target label tensor with ``N`` elements.
        topk: Sequence of k-values for top-k accuracy calculation.

    Returns:
        List of scalar tensors in ``[0, 1]`` matching requested ``topk`` order.
    """
    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")
    if not topk:
        raise ValueError("`topk` cannot be empty")

    with torch.no_grad():
        batch_size = target.size(0)
        if batch_size == 0:
            return [torch.tensor([0.0], device=output.device) for _ in topk]

        num_classes = output.size(1)
        adjusted_topk = tuple(min(int(k), num_classes) for k in topk)
        maxk = max(adjusted_topk)

        _, pred = output.topk(maxk, dim=1, largest=True, sorted=True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        results: list[torch.Tensor] = []
        for k in adjusted_topk:
            correct_k = correct[:k].reshape(-1).float().sum(dim=0, keepdim=True)
            results.append(correct_k.mul_(1.0 / batch_size))
        return results


def calculate_metrics_for_all_classes(predictions: torch.Tensor, target: torch.Tensor) -> tuple[float, float, float]:
    """Compute macro precision, recall, and F1 from class indices.

    Args:
        predictions: Predicted class index tensor.
        target: Ground-truth class index tensor.

    Returns:
        Tuple ``(macro_precision, macro_recall, macro_f1)``.
    """
    predictions_cpu = _as_1d_long_tensor(predictions)
    target_cpu = _as_1d_long_tensor(target)

    if predictions_cpu.numel() != target_cpu.numel():
        raise ValueError("`predictions` and `target` must have the same number of samples")
    if predictions_cpu.numel() == 0:
        return 0.0, 0.0, 0.0

    classes = torch.unique(torch.cat((predictions_cpu, target_cpu), dim=0))
    if classes.numel() == 0:
        return 0.0, 0.0, 0.0

    precision_values: list[float] = []
    recall_values: list[float] = []
    f1_values: list[float] = []
    for class_id in classes.tolist():
        predicted_positive = predictions_cpu == class_id
        actual_positive = target_cpu == class_id

        tp = (predicted_positive & actual_positive).sum().item()
        fp = (predicted_positive & ~actual_positive).sum().item()
        fn = (~predicted_positive & actual_positive).sum().item()

        precision = tp / (tp + fp + _EPSILON)
        recall = tp / (tp + fn + _EPSILON)
        f1_score = (2.0 * precision * recall) / (precision + recall + _EPSILON)

        precision_values.append(float(precision))
        recall_values.append(float(recall))
        f1_values.append(float(f1_score))

    macro_precision = sum(precision_values) / len(precision_values)
    macro_recall = sum(recall_values) / len(recall_values)
    macro_f1 = sum(f1_values) / len(f1_values)

    return float(macro_precision), float(macro_recall), float(macro_f1)
