"""Result formatter utilities shared across eval.py files.

This module converts task-specific evaluation artifacts into a normalized payload:

{
    "category": str,
    "splitType": str,
    "metricName": str,
    "metricValue": float,
    ...optional metadata...
}

Supported tasks:
- Classification (single-label)
- Detection (generic)
- Detection via Ultralytics YOLO result objects
- Keypoint/Pose
- Tracking (MOT-like)
- Multi-label classification
- OCR (text detection/recognition)
"""

from __future__ import annotations

from statistics import mean
from typing import TYPE_CHECKING, Any


if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping, MutableSequence, Sequence


MetricPayload = list[dict[str, Any]]


def _to_float(value: Any) -> float:
    """Convert tensor or scalar-like value to Python ``float``.

    Args:
        value: Metric value that may be tensor, numpy scalar, or numeric.

    Returns:
        Float-cast metric value.
    """
    if hasattr(value, "item"):
        return float(value.item())
    return float(value)


def _label(index_to_labels: Mapping[str, str], class_id: Any) -> str:
    """Resolve class identifier to a stable label string.

    Args:
        index_to_labels: Mapping from class id string to display label.
        class_id: Class identifier value to normalize.

    Returns:
        Resolved class label string.
    """
    key = str(class_id)
    return index_to_labels.get(key, key)


def _append_metric(
    payload: MutableSequence[dict[str, Any]],
    split: str,
    category: str,
    metric_name: str,
    metric_value: Any,
    **metadata: Any,
) -> None:
    """Append one normalized metric record into the payload collection.

    Args:
        payload: Mutable destination list for metric items.
        split: Dataset split label (for example ``train`` or ``val``).
        category: Metric category such as ``all`` or a class label.
        metric_name: Canonical metric name to store.
        metric_value: Raw metric value convertible to ``float``.
        **metadata: Optional extra fields merged into the metric record.

    Returns:
        None.
    """
    item: dict[str, Any] = {
        "category": category,
        "splitType": split,
        "metricName": metric_name,
        "metricValue": _to_float(metric_value),
    }
    if metadata:
        item.update(metadata)
    payload.append(item)


def _safe_mean(values: Sequence[Any]) -> float:
    """Compute mean robustly while handling empty input sequences.

    Args:
        values: Sequence of metric-like values.

    Returns:
        Mean value as float, or ``0.0`` for empty input.
    """
    if not values:
        return 0.0
    return float(mean(_to_float(v) for v in values))


def format_classification_results(
    split: str,
    predictions: Any,
    output: Any,
    target: Any,
    index_to_labels: Mapping[str, str],
    metric_functions: Mapping[str, Any],
) -> MetricPayload:
    """Build standardized single-label classification metric payload.

    Args:
        split: Split name such as ``train`` or ``val``.
        predictions: Predicted class ids from inference runner.
        output: Model output tensor or array used for metrics.
        target: Ground-truth target tensor or array.
        index_to_labels: Mapping of class indices to label names.
        metric_functions: Mapping containing required metric callables.

    Returns:
        Normalized classification metric payload records.
    """
    accuracy_fn = metric_functions["accuracy"]
    precision_fn = metric_functions["precision"]
    recall_fn = metric_functions["recall"]
    f1_score_fn = metric_functions["f1_score_per_class"]
    specificity_fn = metric_functions["specificity"]
    specificity_all_fn = metric_functions["specificity_all"]
    accuracy_per_class_fn = metric_functions["accuracy_per_class"]
    overall_metrics_fn = metric_functions["calculate_metrics_for_all_classes"]

    acc1 = accuracy_fn(output, target, topk=(1,))[0]
    try:
        acc5 = accuracy_fn(output, target, topk=(5,))[0]
    except Exception:
        acc5 = 1.0

    precision_all, recall_all, f1_score_all = overall_metrics_fn(predictions, target)
    results: MetricPayload = []
    _append_metric(results, split, "all", "acc@1", acc1)
    _append_metric(results, split, "all", "acc@5", acc5)
    _append_metric(results, split, "all", "specificity", specificity_all_fn(output, target))
    _append_metric(results, split, "all", "precision", precision_all)
    _append_metric(results, split, "all", "recall", recall_all)
    _append_metric(results, split, "all", "f1_score", f1_score_all)

    for class_id, value in precision_fn(output, target).items():
        _append_metric(results, split, _label(index_to_labels, class_id), "precision", value)
    for class_id, value in f1_score_fn(output, target).items():
        _append_metric(results, split, _label(index_to_labels, class_id), "f1_score", value)
    for class_id, value in recall_fn(output, target).items():
        _append_metric(results, split, _label(index_to_labels, class_id), "recall", value)
    for class_id, value in specificity_fn(output, target).items():
        _append_metric(results, split, _label(index_to_labels, class_id), "specificity", value)
    for class_id, value in accuracy_per_class_fn(output, target).items():
        _append_metric(results, split, _label(index_to_labels, class_id), "acc@1", value)

    return results


def format_multilabel_results(
    split: str,
    per_class_precision: Iterable[Any],
    per_class_recall: Iterable[Any],
    per_class_f1: Iterable[Any],
    macro_f1: Any | None,
    micro_f1: Any | None,
    mAP: Any | None,
    index_to_labels: Mapping[str, str],
) -> MetricPayload:
    """Build normalized payload for multi-label classification metrics.

    Args:
        split: Split name such as ``train`` or ``val``.
        per_class_precision: Iterable of per-class precision values.
        per_class_recall: Iterable of per-class recall values.
        per_class_f1: Iterable of per-class F1 values.
        macro_f1: Optional macro-averaged F1 value.
        micro_f1: Optional micro-averaged F1 value.
        mAP: Optional mean average precision value.
        index_to_labels: Mapping of class indices to label names.

    Returns:
        Normalized multi-label metric payload records.
    """
    precision_list = list(per_class_precision)
    recall_list = list(per_class_recall)
    f1_list = list(per_class_f1)
    results: MetricPayload = []

    _append_metric(results, split, "all", "precision", _safe_mean(precision_list))
    _append_metric(results, split, "all", "recall", _safe_mean(recall_list))
    _append_metric(results, split, "all", "f1_score", _safe_mean(f1_list))
    if macro_f1 is not None:
        _append_metric(results, split, "all", "macro_f1", macro_f1)
    if micro_f1 is not None:
        _append_metric(results, split, "all", "micro_f1", micro_f1)
    if mAP is not None:
        _append_metric(results, split, "all", "mAP", mAP)

    for i, value in enumerate(precision_list):
        _append_metric(results, split, _label(index_to_labels, i), "precision", value)
    for i, value in enumerate(recall_list):
        _append_metric(results, split, _label(index_to_labels, i), "recall", value)
    for i, value in enumerate(f1_list):
        _append_metric(results, split, _label(index_to_labels, i), "f1_score", value)

    return results


def format_detection_results(
    split: str,
    precision: Iterable[Any],
    recall: Iterable[Any],
    f1_score: Iterable[Any],
    index_to_labels: Mapping[str, str],
    mAP50: Any | None = None,
    mAP50_95: Any | None = None,
) -> MetricPayload:
    """Build standardized detection payload from per-class vectors.

    Args:
        split: Split name such as ``train`` or ``val``.
        precision: Iterable of per-class precision values.
        recall: Iterable of per-class recall values.
        f1_score: Iterable of per-class F1 values.
        index_to_labels: Mapping of class indices to label names.
        mAP50: Optional mean average precision at IoU 0.50.
        mAP50_95: Optional mean average precision over IoU range.

    Returns:
        Normalized detection metric payload records.
    """
    precision_list = list(precision)
    recall_list = list(recall)
    f1_list = list(f1_score)
    results: MetricPayload = []

    _append_metric(results, split, "all", "precision", _safe_mean(precision_list))
    _append_metric(results, split, "all", "recall", _safe_mean(recall_list))
    _append_metric(results, split, "all", "f1_score", _safe_mean(f1_list))
    if mAP50 is not None:
        _append_metric(results, split, "all", "mAP@50", mAP50)
    if mAP50_95 is not None:
        _append_metric(results, split, "all", "mAP@50-95", mAP50_95)

    for i, value in enumerate(precision_list):
        _append_metric(results, split, _label(index_to_labels, i), "precision", value)
    for i, value in enumerate(recall_list):
        _append_metric(results, split, _label(index_to_labels, i), "recall", value)
    for i, value in enumerate(f1_list):
        _append_metric(results, split, _label(index_to_labels, i), "f1_score", value)

    return results



def format_keypoint_results(
    split: str,
    keypoint_names: Sequence[str] | None = None,
    mAP: Any | None = None,
    mAP50: Any | None = None,
    oks_mAP: Any | None = None,
    pck: Any | None = None,
    per_keypoint_pck: Iterable[Any] | None = None,
) -> MetricPayload:
    """Build normalized payload for keypoint or pose estimation metrics.

    Args:
        split: Split name such as ``train`` or ``val``.
        keypoint_names: Optional names for each keypoint index.
        mAP: Optional mean average precision value.
        mAP50: Optional mean average precision at IoU 0.50.
        oks_mAP: Optional OKS-based mean average precision.
        pck: Optional overall PCK value.
        per_keypoint_pck: Optional iterable of per-keypoint PCK values.

    Returns:
        Normalized keypoint metric payload records.
    """
    results: MetricPayload = []

    if mAP is not None:
        _append_metric(results, split, "all", "mAP", mAP)
    if mAP50 is not None:
        _append_metric(results, split, "all", "mAP@50", mAP50)
    if oks_mAP is not None:
        _append_metric(results, split, "all", "OKS_mAP", oks_mAP)
    if pck is not None:
        _append_metric(results, split, "all", "PCK", pck)

    if per_keypoint_pck is not None:
        names = list(keypoint_names or [])
        for i, value in enumerate(per_keypoint_pck):
            kp_name = names[i] if i < len(names) else f"keypoint_{i}"
            _append_metric(results, split, kp_name, "PCK", value)

    return results


def format_tracking_results(
    split: str,
    mota: Any | None = None,
    motp: Any | None = None,
    idf1: Any | None = None,
    hota: Any | None = None,
    fp: Any | None = None,
    fn: Any | None = None,
    ids: Any | None = None,
) -> MetricPayload:
    """Build normalized payload for MOT-style tracking metrics.

    Args:
        split: Split name such as ``train`` or ``val``.
        mota: Optional MOTA value.
        motp: Optional MOTP value.
        idf1: Optional IDF1 value.
        hota: Optional HOTA value.
        fp: Optional false-positive count.
        fn: Optional false-negative count.
        ids: Optional identity-switch count.

    Returns:
        Normalized tracking metric payload records.
    """
    results: MetricPayload = []
    if mota is not None:
        _append_metric(results, split, "all", "MOTA", mota)
    if motp is not None:
        _append_metric(results, split, "all", "MOTP", motp)
    if idf1 is not None:
        _append_metric(results, split, "all", "IDF1", idf1)
    if hota is not None:
        _append_metric(results, split, "all", "HOTA", hota)
    if fp is not None:
        _append_metric(results, split, "all", "FP", fp)
    if fn is not None:
        _append_metric(results, split, "all", "FN", fn)
    if ids is not None:
        _append_metric(results, split, "all", "ID_SWITCHES", ids)
    return results


def format_ocr_results(
    split: str,
    cer: Any | None = None,
    wer: Any | None = None,
    precision: Any | None = None,
    recall: Any | None = None,
    f1_score: Any | None = None,
) -> MetricPayload:
    """Build normalized payload for OCR or text pipeline metrics.

    Args:
        split: Split name such as ``train`` or ``val``.
        cer: Optional character error rate.
        wer: Optional word error rate.
        precision: Optional precision value.
        recall: Optional recall value.
        f1_score: Optional F1 score value.

    Returns:
        Normalized OCR metric payload records.
    """
    results: MetricPayload = []
    if cer is not None:
        _append_metric(results, split, "all", "CER", cer)
    if wer is not None:
        _append_metric(results, split, "all", "WER", wer)
    if precision is not None:
        _append_metric(results, split, "all", "precision", precision)
    if recall is not None:
        _append_metric(results, split, "all", "recall", recall)
    if f1_score is not None:
        _append_metric(results, split, "all", "f1_score", f1_score)
    return results


def format_yolo_results(yolo_result: Any, split: str) -> MetricPayload:
    """Build detection payload from Ultralytics YOLO validation result object.

    Args:
        yolo_result: Ultralytics validation result object with metrics and per-class stats.
        split: Split name such as ``train`` or ``val``.

    Returns:
        Normalized detection payload records derived from YOLO result.
    """
    results: MutableSequence[dict[str, Any]] = []
    _append_metric(results, split, "all", "precision", yolo_result.results_dict["metrics/precision(B)"])
    _append_metric(results, split, "all", "recall", yolo_result.results_dict["metrics/recall(B)"])
    _append_metric(results, split, "all", "mAP@50", yolo_result.results_dict["metrics/mAP50(B)"])
    _append_metric(results, split, "all", "mAP@50-95", yolo_result.results_dict["metrics/mAP50-95(B)"])
    _append_metric(results, split, "all", "fitness", yolo_result.results_dict["fitness"])

    index_to_category = yolo_result.names
    for i in range(len(yolo_result.box.ap_class_index)):
        class_name = index_to_category[yolo_result.box.ap_class_index[i]]
        values = [
            yolo_result.box.p[i],
            yolo_result.box.r[i],
            yolo_result.box.f1[i],
            yolo_result.box.ap[i],
            yolo_result.box.maps[i],
        ]
        for metric_name, metric_value in zip(
            ["precision", "recall", "f1_score", "AP", "mAP"],
            values, strict=False,
        ):
            _append_metric(results, split, class_name, metric_name, metric_value)

    return list(results)
