"""Performance metrics and curve fitting.

Definitions follow ai/ANALYSIS_PLAN.md section 3.

Pure numpy/scipy: no plotting, no file I/O, no xarray. This is the part that has
to be numerically correct, so it is kept independently testable.
"""

from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Sequence, Tuple

import numpy as np

from evaluation import CLASS_NAMES

# numpy>=2 renamed trapz -> trapezoid; support both without a static reference to
# whichever name is missing in the installed version.
_trapezoid: Callable[..., Any] = getattr(np, "trapezoid", None) or getattr(np, "trapz")

#: Positive class for binary F1 (disease detection is the meaningful direction).
POSITIVE_CLASS = "AD"
POSITIVE_INDEX = CLASS_NAMES.index(POSITIVE_CLASS)

#: Default coverage grid: keep the most-confident 5%, 10%, ... 100%.
DEFAULT_COVERAGE_GRID: Tuple[float, ...] = tuple(
    round(0.05 * i, 2) for i in range(1, 21)
)

#: A coverage point with fewer retained samples than this is reported as NaN.
MIN_COVERAGE_SAMPLES = 10


# ---------------------------------------------------------------------------
# Classification metrics
# ---------------------------------------------------------------------------
def accuracy(pred: np.ndarray, true: np.ndarray) -> float:
    """Fraction of correct predictions; ``nan`` for an empty input."""
    pred = np.asarray(pred)
    true = np.asarray(true)
    if pred.size == 0:
        return float("nan")
    return float(np.mean(pred == true))


def f1_score(
    pred: np.ndarray, true: np.ndarray, positive: int = POSITIVE_INDEX
) -> float:
    """Binary F1 for the ``positive`` class: ``2TP / (2TP + FP + FN)``.

    Returns ``nan`` when F1 is genuinely undefined -- no true positives *and* no
    predicted positives -- rather than the misleading 0.0 some implementations
    report. A subset containing only true negatives is "no evidence", not
    "perfectly wrong".
    """
    pred = np.asarray(pred)
    true = np.asarray(true)
    if pred.size == 0:
        return float("nan")

    predicted_positive = pred == positive
    actually_positive = true == positive
    tp = int(np.sum(predicted_positive & actually_positive))
    fp = int(np.sum(predicted_positive & ~actually_positive))
    fn = int(np.sum(~predicted_positive & actually_positive))

    denominator = 2 * tp + fp + fn
    if denominator == 0:  # no positives predicted and none present
        return float("nan")
    return float(2 * tp / denominator)


def macro_f1(pred: np.ndarray, true: np.ndarray) -> float:
    """Unweighted mean of the per-class F1 scores (alternative to binary F1)."""
    scores = [
        f1_score(pred, true, positive=index)
        for index in range(len(CLASS_NAMES))
    ]
    finite = [s for s in scores if np.isfinite(s)]
    return float(np.mean(finite)) if finite else float("nan")


# ---------------------------------------------------------------------------
# Coverage / risk-coverage
# ---------------------------------------------------------------------------
@dataclass
class CoverageCurve:
    """Accuracy and F1 as a function of retained (most-confident) fraction."""

    coverage: List[float] = field(default_factory=list)
    accuracy: List[float] = field(default_factory=list)
    f1: List[float] = field(default_factory=list)
    n_retained: List[int] = field(default_factory=list)
    aurc: float = float("nan")


def coverage_curve(
    uncertainty: np.ndarray,
    pred: np.ndarray,
    true: np.ndarray,
    grid: Sequence[float] = DEFAULT_COVERAGE_GRID,
    min_samples: int = MIN_COVERAGE_SAMPLES,
) -> CoverageCurve:
    """Accuracy/F1 after discarding the most uncertain predictions.

    Samples are ranked by ``uncertainty`` ascending (most confident first) with a
    stable sort, so ties resolve deterministically. For each fraction in ``grid``
    the leading portion is retained and scored.

    Args:
        uncertainty: Per-sample uncertainty ``(S,)``; higher = less confident.
        pred: Predicted class indices ``(S,)``.
        true: Ground-truth class indices ``(S,)``.
        grid: Retained fractions in (0, 1].
        min_samples: Points retaining fewer samples than this report ``nan``.

    Returns:
        A :class:`CoverageCurve`, including its AURC summary.
    """
    uncertainty = np.asarray(uncertainty, dtype=float)
    pred = np.asarray(pred)
    true = np.asarray(true)
    n_samples = uncertainty.size

    # Ascending uncertainty = most confident first. A stable sort makes ties
    # resolve by original sample order, so the curve is reproducible.
    order = np.argsort(uncertainty, kind="stable")

    curve = CoverageCurve()
    for fraction in grid:
        n_keep = int(np.clip(round(float(fraction) * n_samples), 0, n_samples))
        retained = order[:n_keep]
        if n_keep < min_samples:
            acc = float("nan")
            f1 = float("nan")
        else:
            acc = accuracy(pred[retained], true[retained])
            f1 = f1_score(pred[retained], true[retained])
        curve.coverage.append(float(fraction))
        curve.accuracy.append(acc)
        curve.f1.append(f1)
        curve.n_retained.append(n_keep)

    curve.aurc = aurc(curve.coverage, curve.accuracy)
    return curve


def aurc(coverage: Sequence[float], accuracy_values: Sequence[float]) -> float:
    """Area under the risk-coverage curve (risk = 1 - accuracy); lower is better.

    Integrated with the trapezoid rule over the finite points and normalised by
    the covered range, so curves whose low-coverage end was suppressed by the
    ``min_samples`` guard remain comparable.
    """
    cov = np.asarray(coverage, dtype=float)
    acc = np.asarray(accuracy_values, dtype=float)
    finite = np.isfinite(acc)
    if finite.sum() < 2:
        return float("nan")

    cov = cov[finite]
    acc = acc[finite]
    span = float(cov[-1] - cov[0])
    if span <= 0:
        return float("nan")
    return float(_trapezoid(1.0 - acc, cov) / span)


# ---------------------------------------------------------------------------
# Binning + curve fitting (noise correlation)
# ---------------------------------------------------------------------------
def quantile_bins(
    uncertainty: np.ndarray, n_bins: int
) -> List[np.ndarray]:
    """Split sample indices into ``n_bins`` equal-count uncertainty bins.

    Equal-count (quantile) bins rather than equal-width, so every point in the
    scatter carries a comparable number of samples.
    """
    order = np.argsort(np.asarray(uncertainty, dtype=float), kind="stable")
    if order.size == 0 or n_bins < 1:
        return []
    return [chunk for chunk in np.array_split(order, n_bins) if chunk.size > 0]


@dataclass
class FitResult:
    """Outcome of fitting accuracy as a function of uncertainty."""

    model: str = "none"
    params: Dict[str, float] = field(default_factory=dict)
    r_squared: float = float("nan")
    spearman: float = float("nan")
    #: Points used, for plotting the fitted line.
    x_fit: List[float] = field(default_factory=list)
    y_fit: List[float] = field(default_factory=list)


def fit_accuracy_vs_uncertainty(
    uncertainty: np.ndarray,
    accuracy_values: np.ndarray,
    model: str = "exponential",
) -> FitResult:
    """Fit accuracy as a decreasing function of uncertainty.

    ``exponential``: ``acc = a + b * exp(-k * u)`` -- bounded and monotone, the
    expected shape. Falls back to a linear fit if the optimiser fails to
    converge. Reports R^2 and Spearman rho; rho is the robust headline because it
    does not assume the functional form.
    """
    from scipy.optimize import curve_fit  # local: keeps import cost off startup
    from scipy.stats import spearmanr

    x = np.asarray(uncertainty, dtype=float)
    y = np.asarray(accuracy_values, dtype=float)
    finite = np.isfinite(x) & np.isfinite(y)
    x, y = x[finite], y[finite]

    result = FitResult()
    if x.size < 3:
        return result  # not enough points to say anything

    if x.size >= 3 and np.ptp(x) > 0:
        rho = float(np.asarray(spearmanr(x, y)).ravel()[0])
        result.spearman = rho if np.isfinite(rho) else float("nan")

    def _exponential(u: np.ndarray, a: float, b: float, k: float) -> np.ndarray:
        return a + b * np.exp(-k * u)

    fitted: np.ndarray | None = None
    if model == "exponential" and x.size >= 4 and np.ptp(x) > 0:
        try:
            scale = float(np.mean(x)) or 1.0
            guess = (float(np.min(y)), float(np.ptp(y)) or 0.1, 1.0 / scale)
            params, _ = curve_fit(_exponential, x, y, p0=guess, maxfev=10000)
            fitted = _exponential(x, *params)
            result.model = "exponential"
            result.params = {
                "a": float(params[0]),
                "b": float(params[1]),
                "k": float(params[2]),
            }
        except Exception:  # noqa: BLE001 - fall back to a linear fit
            fitted = None

    if fitted is None and np.ptp(x) > 0:
        slope, intercept = np.polyfit(x, y, 1)
        fitted = slope * x + intercept
        result.model = "linear"
        result.params = {"slope": float(slope), "intercept": float(intercept)}

    if fitted is None:
        return result

    result.r_squared = r_squared(y, fitted)
    # A smooth line for plotting, evaluated on a dense grid over the fitted range.
    dense = np.linspace(float(x.min()), float(x.max()), 100)
    if result.model == "exponential":
        curve = _exponential(dense, *result.params.values())
    else:
        curve = result.params["slope"] * dense + result.params["intercept"]
    result.x_fit = [float(v) for v in dense]
    result.y_fit = [float(v) for v in curve]
    return result


def r_squared(observed: np.ndarray, predicted: np.ndarray) -> float:
    """Coefficient of determination between observed and fitted values."""
    obs = np.asarray(observed, dtype=float)
    pred = np.asarray(predicted, dtype=float)
    ss_tot = float(np.sum((obs - obs.mean()) ** 2))
    if ss_tot <= 0:
        return float("nan")
    ss_res = float(np.sum((obs - pred) ** 2))
    return float(1.0 - ss_res / ss_tot)