"""How sure is a difference between two models measured on the same gold frames?

The gold set is small (149 unseen frames on 2026-10-02: 40 crystals, 10 ice objects), so a
difference of one or two points in F1 between two models can be the luck of which frames are
gold. A paired bootstrap answers that: draw the gold frames again with replacement, the SAME
draw for both models, and look at the spread of the difference. ``win`` is the share of draws
in which the candidate beats the baseline.

Pure Python on purpose (the gate's tests run without numpy); only ``per_frame`` needs the
evaluation code, and imports it when called.
"""
from __future__ import annotations

import random
from pathlib import Path
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

Counts = Tuple[float, float, float, float]          # tp, fp, fn, sum of matched mask IoU
FrameCounts = Dict[int, Counts]                      # class id -> counts, for one frame


def per_frame(pred_dir, gold_dir, iou_thr: float = 0.5, stems: Optional[Iterable[str]] = None) -> Dict[str, FrameCounts]:
    """Per gold frame and class: tp, fp, fn and the summed mask IoU of the matches.

    ``stems`` restricts the frames; a gold frame without a prediction file is left out
    (the caller pairs on the frames both models have)."""
    from src.eval.metrics import evaluate_image
    from src.eval.yolo_labels import read_label_file

    pred_dir, gold_dir = Path(pred_dir), Path(gold_dir)
    wanted = set(stems) if stems is not None else None
    out: Dict[str, FrameCounts] = {}
    for g in sorted(gold_dir.glob("*.txt")):
        if wanted is not None and g.stem not in wanted:
            continue
        p = pred_dir / g.name
        if not p.exists():
            continue
        frame: FrameCounts = {}
        for c, cc in evaluate_image(read_label_file(g), read_label_file(p), iou_thr).items():
            frame[int(c)] = (float(cc.tp), float(cc.fp), float(cc.fn), float(sum(cc.mask_ious)))
        out[g.stem] = frame
    return out


def _total(frames: Sequence[FrameCounts], class_id: Optional[int]) -> Counts:
    tp = fp = fn = iou = 0.0
    for f in frames:
        for c, (a, b, d, e) in f.items():
            if class_id is None or c == class_id:
                tp += a; fp += b; fn += d; iou += e
    return tp, fp, fn, iou


def f1_of(c: Counts) -> float:
    tp, fp, fn, _ = c
    return 2 * tp / (2 * tp + fp + fn) if tp else 0.0


def mask_score_of(c: Counts) -> float:
    """Mean mask IoU over every gold object, a missed object counting as zero."""
    tp, _, fn, iou = c
    return iou / (tp + fn) if (tp + fn) else 0.0


def _summary(point: float, deltas: List[float]) -> Dict[str, float]:
    d = sorted(deltas)
    n = len(d)
    lo, hi = d[int(0.025 * (n - 1))], d[int(0.975 * (n - 1))]
    return {"delta": round(point, 4), "lo": round(lo, 4), "hi": round(hi, 4),
            "win": round(sum(1 for x in d if x > 0) / n, 3)}


def paired_bootstrap(cand: Mapping[str, FrameCounts], base: Mapping[str, FrameCounts], *, draws: int = 2000,
                     seed: int = 0, classes: Optional[Mapping[str, int]] = None) -> Dict:
    """Candidate minus baseline on the frames both have: point difference, 95 % interval and win share
    for F1 and the mask score, overall and for each class in ``classes`` (name -> class id)."""
    keys = sorted(set(cand) & set(base))
    n = len(keys)
    if n == 0:
        return {"frames": 0, "draws": 0}
    a = [cand[k] for k in keys]; b = [base[k] for k in keys]
    targets: List[Tuple[str, Optional[int]]] = [("overall", None)] + [(name, cid) for name, cid in (classes or {}).items()]
    point = {}
    for name, cid in targets:
        ca, cb = _total(a, cid), _total(b, cid)
        point[name] = (f1_of(ca) - f1_of(cb), mask_score_of(ca) - mask_score_of(cb))
    rng = random.Random(seed)
    d_f1 = {name: [] for name, _ in targets}; d_ms = {name: [] for name, _ in targets}
    for _ in range(draws):
        idx = [rng.randrange(n) for _ in range(n)]
        sa = [a[i] for i in idx]; sb = [b[i] for i in idx]
        for name, cid in targets:
            ca, cb = _total(sa, cid), _total(sb, cid)
            d_f1[name].append(f1_of(ca) - f1_of(cb)); d_ms[name].append(mask_score_of(ca) - mask_score_of(cb))
    out = {"frames": n, "draws": draws,
           "f1": _summary(point["overall"][0], d_f1["overall"]),
           "mask_score": _summary(point["overall"][1], d_ms["overall"])}
    if classes:
        out["per_class"] = {name: {"f1": _summary(point[name][0], d_f1[name]), "mask_score": _summary(point[name][1], d_ms[name])}
                            for name in classes}
    return out