"""Shared training utilities — metrics, averaging, and data-cleaning helpers.

Extracted from the ``ml_pytorch_vision_classification``, ``r2plus1d``,
and ``ml_yolov10_detection`` BYOM codebases where they were duplicated.
"""

from __future__ import annotations

import math
from typing import Any


# ------------------------------------------------------------------ #
#  AverageMeter                                                        #
# ------------------------------------------------------------------ #


class AverageMeter:
    """Running average tracker for a single scalar metric.

    Commonly used to track loss and accuracy across mini-batches within
    an epoch.

    Example::

        losses = AverageMeter("Loss")
        for inputs, targets in loader:
            loss = criterion(model(inputs), targets)
            losses.update(loss.item(), n=inputs.size(0))
        print(losses)  # Loss 0.3420 (0.4567)
        print(losses.avg)  # 0.4567

    Args:
        name: Human-readable label shown in ``__str__``.
        fmt: Format spec for ``val`` and ``avg`` (default ``":f"``).
    """

    def __init__(self, name: str = "", fmt: str = ":f") -> None:
        """Initialise the meter.

        Args:
            name: Display name for :meth:`__str__`.
            fmt: Float format spec (e.g. ``":.4f"``, ``":.2e"``).
        """
        self.name = name
        self.fmt = fmt
        self.val: float = 0.0
        self.avg: float = 0.0
        self.sum: float = 0.0
        self.count: int = 0

    def reset(self) -> None:
        """Zero all accumulators."""
        self.val = 0.0
        self.avg = 0.0
        self.sum = 0.0
        self.count = 0

    def update(self, val: float, n: int = 1) -> None:
        """Record *n* observations with value *val*.

        Args:
            val: Per-sample (or per-batch-mean) value.
            n: Number of observations this value represents
                (typically ``batch_size``).
        """
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self) -> str:
        """Format as ``'name val (avg)'``."""
        fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})"
        return fmtstr.format(**self.__dict__)

    def __repr__(self) -> str:
        """Return a developer-friendly representation."""
        return f"AverageMeter({self.name!r}, count={self.count}, avg={self.avg:.6f})"


# ------------------------------------------------------------------ #
#  Top-k accuracy                                                      #
# ------------------------------------------------------------------ #


def accuracy(
    output: Any,
    target: Any,
    topk: tuple[int, ...] = (1,),
) -> list[float]:
    """Compute top-k classification accuracy.

    Handles the edge case where *k* exceeds the number of classes by
    clamping to ``output.size(1)`` (from the r2plus1d BYOM repo).

    Args:
        output: Model logits with shape ``(N, C)``.
        target: Ground-truth class indices with shape ``(N,)``.
        topk: Tuple of *k* values to compute accuracy for.

    Returns:
        List of accuracy values (one per *k*), each in the **[0, 1]**
        range.
    """
    import torch

    with torch.no_grad():
        num_classes = output.size(1)
        maxk = min(max(topk), num_classes)
        batch_size = target.size(0)

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

        results: list[float] = []
        for k in topk:
            k_clamped = min(k, num_classes)
            if k_clamped <= 0:
                results.append(0.0)
            else:
                correct_k = correct[:k_clamped].reshape(-1).float().sum(0, keepdim=True)
                results.append(correct_k.mul_(1.0 / batch_size).item())

        return results


# ------------------------------------------------------------------ #
#  Metric sanitisation                                                 #
# ------------------------------------------------------------------ #


def sanitise_metrics(
    metrics: list[dict[str, Any]],
    *,
    precision: int = 4,
) -> list[dict[str, Any]]:
    """Clean metric values for JSON serialisation.

    Replaces ``inf`` / ``nan`` with ``0`` and rounds floats to
    *precision* decimal places.  This is the unified version of the
    ``round_metrics`` helper duplicated in the YOLOv10, D-FINE, and
    other BYOM codebases.

    Args:
        metrics: Raw epoch-metrics list — each entry is a dict with at
            least a ``metricValue`` key.
        precision: Number of decimal places to round to.

    Returns:
        A **new** list with sanitised numeric values (the original is
        not mutated).
    """
    cleaned: list[dict[str, Any]] = []
    for m in metrics:
        entry = dict(m)
        value = entry.get("metricValue")
        if isinstance(value, float):
            if math.isinf(value) or math.isnan(value):
                entry["metricValue"] = 0
            else:
                entry["metricValue"] = round(value, precision)
        cleaned.append(entry)
    return cleaned
