"""Reusable metric primitives extracted from the validated drift methodology."""

from __future__ import annotations

import math
from typing import Iterable, Sequence


def _as_float_list(values: Iterable[float]) -> list[float]:
    out = [float(v) for v in values]
    if not out:
        raise ValueError("expected at least one value")
    return out


def empirical_quantile(values: Sequence[float], q: float) -> float:
    """Return the q-quantile using linear interpolation over sorted samples."""
    if not 0.0 <= q <= 1.0:
        raise ValueError("q must be within [0, 1]")
    arr = sorted(_as_float_list(values))
    if len(arr) == 1:
        return arr[0]
    pos = q * (len(arr) - 1)
    low = int(math.floor(pos))
    high = int(math.ceil(pos))
    if low == high:
        return arr[low]
    frac = pos - low
    return (1.0 - frac) * arr[low] + frac * arr[high]


def mann_whitney_auroc(id_scores: Sequence[float], ood_scores: Sequence[float]) -> float:
    """Compute AUROC via the Mann-Whitney U statistic without external deps."""
    id_vals = _as_float_list(id_scores)
    ood_vals = _as_float_list(ood_scores)
    paired: list[tuple[float, int]] = [(x, 0) for x in id_vals] + [(x, 1) for x in ood_vals]
    paired.sort(key=lambda item: item[0])

    rank = 1
    i = 0
    ood_rank_sum = 0.0
    while i < len(paired):
        j = i + 1
        while j < len(paired) and paired[j][0] == paired[i][0]:
            j += 1
        mean_rank = (rank + (rank + (j - i) - 1)) / 2.0
        for k in range(i, j):
            if paired[k][1] == 1:
                ood_rank_sum += mean_rank
        rank += j - i
        i = j

    n_ood = len(ood_vals)
    n_id = len(id_vals)
    u = ood_rank_sum - n_ood * (n_ood + 1) / 2.0
    return float(u / (n_ood * n_id))


def ood_metrics_from_scores(
    id_scores: Sequence[float],
    ood_scores: Sequence[float],
    threshold_quantile: float = 0.95,
) -> dict[str, float]:
    """Compute AUROC/FPR/TPR metrics at a quantile threshold on ID scores."""
    id_vals = _as_float_list(id_scores)
    ood_vals = _as_float_list(ood_scores)
    thr = empirical_quantile(id_vals, threshold_quantile)
    fpr = sum(1 for x in id_vals if x > thr) / len(id_vals)
    tpr = sum(1 for x in ood_vals if x > thr) / len(ood_vals)
    return {
        "auroc": mann_whitney_auroc(id_vals, ood_vals),
        "fpr": float(fpr),
        "tpr": float(tpr),
        "threshold": float(thr),
    }


def rbf_mmd2_from_gaussian(mean_shift: float, var_ratio: float, gamma: float) -> float:
    """Closed-form RBF-MMD^2 for N(mean_shift, var_ratio) versus N(0, 1)."""
    if var_ratio <= 0:
        raise ValueError("var_ratio must be positive")
    if gamma <= 0:
        raise ValueError("gamma must be positive")

    p0_var = 1.0
    xx = 1.0 / math.sqrt(1.0 + 2.0 * gamma * var_ratio)
    yy = 1.0 / math.sqrt(1.0 + 2.0 * gamma * p0_var)
    denom = 1.0 + gamma * (var_ratio + p0_var)
    xy = (1.0 / math.sqrt(denom)) * math.exp(-gamma * (mean_shift**2) / denom)
    return float(xx + yy - 2.0 * xy)


def fused_discrepancy_score(
    energy_shift: float,
    mmd_shift: float,
    quantile_shift: float,
    weights: tuple[float, float, float] = (0.5, 0.3, 0.2),
) -> float:
    """Fuse drift heads (energy-like, discrepancy, quantile) into one score."""
    if len(weights) != 3:
        raise ValueError("weights must contain exactly three values")
    if any(w < 0 for w in weights):
        raise ValueError("weights must be non-negative")
    total = sum(weights)
    if total <= 0:
        raise ValueError("weights must not sum to zero")

    wn = tuple(w / total for w in weights)
    return float(
        wn[0] * abs(float(energy_shift))
        + wn[1] * abs(float(mmd_shift))
        + wn[2] * abs(float(quantile_shift))
    )