# SPDX-License-Identifier: Apache-2.0
"""C06 metrics: device-vs-reference agreement measures (host side, numpy only, float64 arithmetic).

Inputs may be numpy arrays, torch tensors (bf16 included) or nested lists. Gates are applied by
:mod:`.golden` (``GateRegistry``); these functions only measure.

==========================  =====================================================================================
function                    use
==========================  =====================================================================================
``pcc``                     dense Pearson correlation (the default gate metric)
``masked_pcc``              PCC over the elements where a mask is true (valid pillars, in-range points)
``valid_row_pcc``           PCC over the first ``n_valid`` rows of a fixed-capacity buffer
``error_stats``             max / mean abs error, relative L2, PCC in one dict
``argmax_agreement``        fraction of positions whose argmax over ``axis`` agrees (segmentation logits)
``label_agreement``         fraction of equal labels (decided class maps)
``mask_iou`` ``mean_iou``   binary-mask IoU; per-class and mean IoU of label maps
``box_iou_xyxy``            pairwise IoU of 2-D boxes
``topk_overlap``            |A ∩ B| / k of two index sets (data-dependent selections)
``topk_set_overlap``        the same, computed from two score arrays
``match_detections``        greedy same-label matching by BEV-centre distance: recall, precision, errors
``ade_fde``                 average / final displacement error of trajectories
==========================  =====================================================================================
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

__all__ = [
    "as_array",
    "pcc",
    "masked_pcc",
    "valid_row_pcc",
    "error_stats",
    "argmax_agreement",
    "label_agreement",
    "mask_iou",
    "mean_iou",
    "box_iou_xyxy",
    "topk_overlap",
    "topk_set_overlap",
    "DetectionMatch",
    "match_detections",
    "ade_fde",
]


def as_array(x: Any, dtype: Optional[Any] = None) -> np.ndarray:
    """numpy / torch (any dtype, bf16 widened to float32) / list -> numpy, without importing torch."""
    if hasattr(x, "detach") and hasattr(x, "cpu"):
        t = x.detach().cpu()
        if str(t.dtype) == "torch.bfloat16":
            t = t.float()
        x = t.numpy()
    a = np.asarray(x)
    return a if dtype is None else a.astype(dtype, copy=False)


def _pcc_flat(a: np.ndarray, b: np.ndarray) -> float:
    if a.shape != b.shape:
        raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}")
    if a.size == 0:
        return float("nan")
    if not (np.isfinite(a).all() and np.isfinite(b).all()):
        return float("nan")
    # A constant side is detected exactly (max == min), never from the centred sum of squares: the mean of a
    # constant array is not exact in floating point (np.full(1000, 0.1).mean() != 0.1), so the centred values hold
    # rounding noise that would correlate perfectly with the other side's noise (two different constants gave +-1).
    const_a, const_b = a.max() == a.min(), b.max() == b.min()
    if const_a or const_b:  # equal constants agree perfectly, anything else does not
        return 1.0 if (const_a and const_b and np.allclose(a, b)) else 0.0
    da = a - a.mean()
    db = b - b.mean()
    sa = float(np.dot(da, da))
    sb = float(np.dot(db, db))
    return float(np.dot(da, db) / math.sqrt(sa * sb))


def pcc(test: Any, ref: Any) -> float:
    """Pearson correlation of two equally shaped tensors (float64). NaN if any value is non-finite or both are
    empty; when a side is constant: 1.0 for two equal constants (``np.allclose``), else 0.0 (two different
    constants included)."""
    return _pcc_flat(as_array(test, np.float64).reshape(-1), as_array(ref, np.float64).reshape(-1))


def masked_pcc(test: Any, ref: Any, mask: Any) -> float:
    """PCC over ``mask`` (broadcastable to the tensors' shape)."""
    a, b = as_array(test, np.float64), as_array(ref, np.float64)
    if a.shape != b.shape:
        raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}")
    m = np.broadcast_to(as_array(mask).astype(bool), a.shape)
    return _pcc_flat(a[m], b[m])


def valid_row_pcc(test: Any, ref: Any, n_valid: Optional[int] = None, *, rows: Any = None, axis: int = 0) -> float:
    """PCC over the valid rows of fixed-capacity buffers: the first ``n_valid`` along ``axis``, or the rows where
    the boolean vector ``rows`` is true."""
    a, b = as_array(test, np.float64), as_array(ref, np.float64)
    if (n_valid is None) == (rows is None):
        raise ValueError("give exactly one of n_valid and rows")
    if rows is None:
        idx = np.arange(int(n_valid))
    else:
        idx = np.flatnonzero(as_array(rows).astype(bool))
    return _pcc_flat(np.take(a, idx, axis=axis).reshape(-1), np.take(b, idx, axis=axis).reshape(-1))


def error_stats(test: Any, ref: Any) -> Dict[str, float]:
    """``{"pcc", "max_abs", "mean_abs", "rel_l2", "n"}`` (``rel_l2 = ||t - r|| / ||r||``)."""
    a, b = as_array(test, np.float64), as_array(ref, np.float64)
    if a.shape != b.shape:
        raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}")
    d = np.abs(a - b)
    ref_norm = float(np.linalg.norm(b))
    return {"pcc": _pcc_flat(a.reshape(-1), b.reshape(-1)),
            "max_abs": float(d.max()) if d.size else 0.0,
            "mean_abs": float(d.mean()) if d.size else 0.0,
            "rel_l2": float(np.linalg.norm(a - b) / ref_norm) if ref_norm > 0 else float(np.linalg.norm(a - b)),
            "n": int(a.size)}


def argmax_agreement(test_logits: Any, ref_logits: Any, *, axis: int = -1, mask: Any = None) -> float:
    """Fraction of positions where ``argmax`` over ``axis`` agrees (ties resolve to the lowest index, as numpy
    and Autoware's decision rules do). ``mask`` selects positions (shape without ``axis``)."""
    a = np.argmax(as_array(test_logits), axis=axis)
    b = np.argmax(as_array(ref_logits), axis=axis)
    return label_agreement(a, b, mask=mask)


def label_agreement(test_labels: Any, ref_labels: Any, *, mask: Any = None) -> float:
    """Fraction of equal labels (optionally over ``mask``). NaN when nothing is selected."""
    a, b = as_array(test_labels), as_array(ref_labels)
    if a.shape != b.shape:
        raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}")
    eq = a == b
    if mask is not None:
        eq = eq[np.broadcast_to(as_array(mask).astype(bool), eq.shape)]
    return float(eq.mean()) if eq.size else float("nan")


def mask_iou(test_mask: Any, ref_mask: Any) -> float:
    """IoU of two binary masks (1.0 when both are empty)."""
    a, b = as_array(test_mask).astype(bool), as_array(ref_mask).astype(bool)
    union = np.logical_or(a, b).sum()
    return 1.0 if union == 0 else float(np.logical_and(a, b).sum() / union)


def mean_iou(test_labels: Any, ref_labels: Any, num_classes: int, *,
             ignore_index: Optional[int] = None) -> Tuple[float, List[float]]:
    """``(mIoU over classes present in either map, per-class IoU with NaN for absent classes)``."""
    a, b = as_array(test_labels).reshape(-1), as_array(ref_labels).reshape(-1)
    if a.shape != b.shape:
        raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}")
    if ignore_index is not None:
        keep = b != ignore_index
        a, b = a[keep], b[keep]
    per = []
    for c in range(num_classes):
        inter = np.logical_and(a == c, b == c).sum()
        union = np.logical_or(a == c, b == c).sum()
        per.append(float(inter / union) if union else float("nan"))
    valid = [v for v in per if not math.isnan(v)]
    return (float(np.mean(valid)) if valid else float("nan")), per


def box_iou_xyxy(boxes_a: Any, boxes_b: Any) -> np.ndarray:
    """Pairwise IoU ``[N, M]`` of axis-aligned boxes ``[x1, y1, x2, y2]`` (float64)."""
    a = as_array(boxes_a, np.float64).reshape(-1, 4)
    b = as_array(boxes_b, np.float64).reshape(-1, 4)
    ix1 = np.maximum(a[:, None, 0], b[None, :, 0])
    iy1 = np.maximum(a[:, None, 1], b[None, :, 1])
    ix2 = np.minimum(a[:, None, 2], b[None, :, 2])
    iy2 = np.minimum(a[:, None, 3], b[None, :, 3])
    inter = np.clip(ix2 - ix1, 0, None) * np.clip(iy2 - iy1, 0, None)
    area_a = np.clip(a[:, 2] - a[:, 0], 0, None) * np.clip(a[:, 3] - a[:, 1], 0, None)
    area_b = np.clip(b[:, 2] - b[:, 0], 0, None) * np.clip(b[:, 3] - b[:, 1], 0, None)
    union = area_a[:, None] + area_b[None, :] - inter
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.where(union > 0, inter / union, 0.0)


def topk_overlap(test_indices: Any, ref_indices: Any, k: Optional[int] = None) -> float:
    """``|set(test[:k]) ∩ set(ref[:k])| / k`` (``k`` defaults to ``len(ref)``)."""
    a, b = as_array(test_indices).reshape(-1), as_array(ref_indices).reshape(-1)
    k = len(b) if k is None else int(k)
    if k <= 0:
        return float("nan")
    return len(set(a[:k].tolist()) & set(b[:k].tolist())) / k


def topk_set_overlap(test_scores: Any, ref_scores: Any, k: int) -> float:
    """Overlap of the top-``k`` index sets of two score arrays (flattened; stable descending order)."""
    a, b = as_array(test_scores, np.float64).reshape(-1), as_array(ref_scores, np.float64).reshape(-1)
    ta = np.argsort(-a, kind="stable")[:k]
    tb = np.argsort(-b, kind="stable")[:k]
    return topk_overlap(ta, tb, k)


@dataclass
class DetectionMatch:
    """Result of :func:`match_detections` (``pairs`` are ``(test_index, ref_index)``)."""

    pairs: List[Tuple[int, int]]
    n_test: int
    n_ref: int
    center_errors: List[float] = field(default_factory=list)
    score_errors: List[float] = field(default_factory=list)

    @property
    def matched(self) -> int:
        return len(self.pairs)

    @property
    def recall(self) -> float:
        return self.matched / self.n_ref if self.n_ref else 1.0

    @property
    def precision(self) -> float:
        return self.matched / self.n_test if self.n_test else 1.0

    def to_dict(self) -> Dict[str, Any]:
        def stat(v: List[float], fn) -> Optional[float]:
            return float(fn(v)) if v else None

        return {"n_test": self.n_test, "n_ref": self.n_ref, "matched": self.matched, "recall": self.recall,
                "precision": self.precision, "center_err_mean": stat(self.center_errors, np.mean),
                "center_err_max": stat(self.center_errors, np.max), "score_err_max": stat(self.score_errors, np.max),
                "score_err_mean": stat(self.score_errors, np.mean)}


def _centers(x: Any, n: int, dims: int) -> np.ndarray:
    a = as_array(x, np.float64)
    if a.size == 0:
        return np.zeros((n, dims))
    if a.size % max(n, 1) or n == 0:
        raise ValueError(f"{a.shape[0] if a.ndim else 0} centres for {n} labels: lengths differ")
    return a.reshape(n, -1)[:, :dims]


def match_detections(test_centers: Any, test_labels: Any, test_scores: Any, ref_centers: Any, ref_labels: Any,
                     ref_scores: Any, *, max_dist: float = 0.5, dims: int = 2,
                     same_label: bool = True) -> DetectionMatch:
    """Greedy matching: reference boxes in descending score order each take the nearest unused test box (same
    label unless ``same_label=False``) within ``max_dist`` of its centre, measured over the first ``dims``
    coordinates (2 = BEV x/y). ``*_centers`` are ``[N, >=dims]`` (box arrays ``[N, 7]`` work as-is). A box with a
    non-finite centre never matches (it costs recall / precision; it does not block the other boxes)."""
    tl, rl = as_array(test_labels).reshape(-1), as_array(ref_labels).reshape(-1)
    ts, rs = as_array(test_scores, np.float64).reshape(-1), as_array(ref_scores, np.float64).reshape(-1)
    if len(ts) != len(tl) or len(rs) != len(rl):
        raise ValueError(f"labels / scores lengths differ: test {len(tl)}/{len(ts)}, ref {len(rl)}/{len(rs)}")
    tc, rc = _centers(test_centers, len(tl), dims), _centers(ref_centers, len(rl), dims)
    used = np.zeros(len(tl), bool)
    result = DetectionMatch([], len(tl), len(rl))
    for i in np.argsort(-rs, kind="stable"):
        if not len(tl):
            break
        d = np.linalg.norm(tc - rc[i], axis=1)
        bad = used | ~np.isfinite(d)          # np.argmin would pick a NaN distance before any real one
        if same_label:
            bad |= tl != rl[i]
        d[bad] = np.inf
        j = int(np.argmin(d))
        if d[j] <= max_dist:
            used[j] = True
            result.pairs.append((j, int(i)))
            result.center_errors.append(float(d[j]))
            result.score_errors.append(abs(float(ts[j] - rs[i])))
    return result


def ade_fde(pred: Any, ref: Any, *, dims: int = 2) -> Tuple[float, float]:
    """Average and final displacement errors of trajectories ``[..., T, >=dims]`` (Euclidean over the first
    ``dims`` coordinates; averaged over all leading dims)."""
    p, r = as_array(pred, np.float64), as_array(ref, np.float64)
    if p.shape[:-1] != r.shape[:-1]:
        raise ValueError(f"trajectory shapes differ: {p.shape} vs {r.shape}")
    disp = np.linalg.norm(p[..., :dims] - r[..., :dims], axis=-1)
    return float(disp.mean()), float(disp[..., -1].mean())