"""
common.py -- shared configuration, dataset, augmentation and detection metrics
for Part B (Kvasir-SEG polyp detection).

Seed = 16 (last two digits of the roll number), identical policy to Part A.

The detection metrics (AP at a given IoU, COCO-style mAP@[.50:.05:.95], matched
IoU, size-stratified recall) are implemented here from scratch rather than
imported from pycocotools, so every number in the report can be traced to code
the student can explain line by line.
"""

from __future__ import annotations

import json
import os
import random
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

SEED: int = 16
SEED_TAG: str = f"seed{SEED:02d}"

# Kvasir-SEG has a single foreground category. Faster R-CNN reserves index 0
# for background, so "polyp" is class index 1.
CATEGORIES: List[str] = ["__background__", "polyp"]
NUM_CLASSES: int = len(CATEGORIES)

SPLIT_FRACTIONS = (0.70, 0.15, 0.15)
DEFAULT_IMG_SIZE = 512

ROOT = Path(__file__).resolve().parents[1]
DATASET_DIR = ROOT / "Dataset" / "kvasir"
SPLIT_DIR = ROOT / "splits"
RESULTS_DIR = ROOT / "Results"
METRICS_DIR = RESULTS_DIR / "metrics"
PLOTS_DIR = RESULTS_DIR / "plots"
OVERLAY_DIR = RESULTS_DIR / "overlays"
WEIGHTS_DIR = RESULTS_DIR / "model_weights"

DET_SPLIT_CSV = SPLIT_DIR / f"kvasir_detection_{SEED_TAG}.csv"

for _d in (SPLIT_DIR, METRICS_DIR, PLOTS_DIR, OVERLAY_DIR, WEIGHTS_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# The bbox json ships with a typo in the official Kvasir-SEG archive
# ("kavsir_bboxes.json"); the assignment spells it "kvasir_bboxes.json".
# Both spellings are accepted.
BBOX_JSON_NAMES = ["kavsir_bboxes.json", "kvasir_bboxes.json", "bboxes.json"]
BBOX_DIR_NAMES = ["bbox", "bboxes", "Bbox", "annotations"]


# ---------------------------------------------------------------------------
def set_seed(seed: int = SEED, deterministic: bool = True) -> None:
    os.environ["PYTHONHASHSEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    try:
        import torch

        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        if deterministic:
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False
    except ImportError:
        pass


def seeded_rng(offset: int = 0) -> np.random.Generator:
    return np.random.default_rng(SEED + offset)


def worker_init_fn(worker_id: int) -> None:
    np.random.seed(SEED + worker_id)
    random.seed(SEED + worker_id)


def get_device():
    import torch

    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def save_json(obj, path: Path) -> None:
    from json import dump

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    def enc(o):
        if isinstance(o, (np.integer,)):
            return int(o)
        if isinstance(o, (np.floating,)):
            return float(o)
        if isinstance(o, np.ndarray):
            return o.tolist()
        if isinstance(o, Path):
            return str(o)
        raise TypeError(str(type(o)))

    with open(path, "w") as f:
        dump(obj, f, indent=2, default=enc)
    print(f"[saved] {path}")


def progress(iterable, desc: str = "", total: Optional[int] = None):
    try:
        from tqdm.auto import tqdm

        return tqdm(iterable, desc=desc, total=total, leave=False)
    except ImportError:
        return iterable


def _mpl():
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    return plt


# ---------------------------------------------------------------------------
# Dataset discovery + annotation loading
# ---------------------------------------------------------------------------

KAGGLE_SLUG = "debeshjha1/kvasirseg"


def _kagglehub_path(slug: str) -> Optional[Path]:
    try:
        import kagglehub
    except ImportError:
        print("[info] kagglehub is not installed (pip install kagglehub) -- "
              "falling back to a local Dataset/ folder")
        return None
    try:
        print(f"[info] resolving '{slug}' through kagglehub (downloads once, then cached) ...")
        return Path(kagglehub.dataset_download(slug))
    except Exception as e:
        print(f"[warn] kagglehub could not provide '{slug}': {e}")
        return None


def find_kvasir_root(user_root: Optional[str] = None, use_kagglehub: bool = True) -> Path:
    """Search order: --data-root, then Dataset/kvasir/, then the kagglehub cache."""
    def _ok(c: Path) -> bool:
        return bool(c) and (c / "images").is_dir() and (c / "masks").is_dir()

    cands: List[Path] = []
    if user_root:
        cands.append(Path(user_root))
    cands += [DATASET_DIR / "Kvasir-SEG", DATASET_DIR]
    if DATASET_DIR.exists():
        cands += [p for p in DATASET_DIR.rglob("Kvasir-SEG") if p.is_dir()]
    for c in cands:
        if _ok(c):
            return c.resolve()

    if use_kagglehub:
        kp = _kagglehub_path(KAGGLE_SLUG)
        if kp is not None and kp.exists():
            for c in [kp / "Kvasir-SEG", kp, *[p for p in kp.rglob("Kvasir-SEG") if p.is_dir()]]:
                if _ok(c):
                    _link_into_dataset_dir(c)
                    return c.resolve()

    raise FileNotFoundError(
        f"Kvasir-SEG not found under {DATASET_DIR} or the kagglehub cache.\n"
        "Either run  python download_data.py  from the submission root, or unzip the "
        f"Kaggle download so that {DATASET_DIR}/Kvasir-SEG/images and .../masks exist, "
        "or pass --data-root explicitly."
    )


def _link_into_dataset_dir(src: Path) -> None:
    dst = DATASET_DIR / "Kvasir-SEG"
    try:
        DATASET_DIR.mkdir(parents=True, exist_ok=True)
        if dst.exists() or dst.is_symlink():
            return
        dst.symlink_to(src, target_is_directory=True)
        print(f"[info] linked {dst} -> {src}")
    except Exception:
        pass


def find_bbox_json(root: Path) -> Path:
    """Locate the box annotations. Three real-world layouts are handled:

      1. a bbox/ directory of per-image CSV files, named <image_stem>.csv,
         each shaped:
             class_name,xmin,ymin,xmax,ymax
             polyp,38,5,430,338
         This is what the current Kaggle archive (debeshjha1/kvasirseg,
         version 3, confirmed by direct inspection) actually ships -- 1000
         files, one per image, no JSON anywhere in the archive.
      2. a bbox/ directory of per-image JSON files (seen in some
         re-packagings / older archive versions), each shaped:
             {"bbox": [{"label": "polyp", "xmin": ..., "ymin": ..., ...}],
              "width": ..., "height": ...}
      3. a single merged json, e.g. Kvasir-SEG/kavsir_bboxes.json (the layout
         described in some course materials, not seen in the live archive but
         supported in case a differently-packaged copy is used):
             {"cju0qkwl35piu...": {"bbox": [...], "height": ..., "width": ...}}

    Returns the bbox/ DIRECTORY for cases (1)/(2), or the file path for case
    (3); load_bboxes() below dispatches on whether it receives a directory or
    a file, and on the file extension of what is inside.
    """
    for name in BBOX_DIR_NAMES:
        for p in (root / name, root.parent / name):
            if p.is_dir() and (any(p.glob("*.csv")) or any(p.glob("*.json"))):
                return p
    for name in BBOX_JSON_NAMES:
        for p in (root / name, root.parent / name):
            if p.exists():
                return p
    hits_dir = [d for d in root.parent.rglob("*") if d.is_dir()
               and d.name.lower() in [n.lower() for n in BBOX_DIR_NAMES]
               and (any(d.glob("*.csv")) or any(d.glob("*.json")))]
    if hits_dir:
        return hits_dir[0]
    hits = list(root.parent.rglob("*bboxes*.json"))
    if hits:
        return hits[0]
    raise FileNotFoundError(
        f"Could not find box annotations under {root}. Tried a per-image "
        f"folder named one of {BBOX_DIR_NAMES} containing *.csv or *.json "
        f"files, and a merged file named one of {BBOX_JSON_NAMES}."
    )


def _load_bbox_csv_dir(csv_dir: Path) -> Dict[str, dict]:
    """Parse Kvasir-SEG's actual per-image CSV annotation format:
        class_name,xmin,ymin,xmax,ymax
        polyp,38,5,430,338
    One file per image, named <image_stem>.csv. Width/height are not recorded
    in these files, so they are left as None here; callers fill them in from
    the image itself (pre_process.py already reads each image's own size for
    other checks, so this costs nothing extra)."""
    import csv as _csv

    files = sorted(csv_dir.glob("*.csv"))
    if not files:
        raise FileNotFoundError(f"No .csv files found in bbox directory {csv_dir}")
    out: Dict[str, dict] = {}
    for f in files:
        boxes, labels = [], []
        with open(f, newline="") as fh:
            reader = _csv.DictReader(fh)
            fieldmap = {k.strip().lower(): k for k in (reader.fieldnames or [])}
            for row in reader:
                def _get(*names):
                    for n in names:
                        if n in fieldmap:
                            return row[fieldmap[n]]
                    return None
                xmin, ymin = _get("xmin", "x1"), _get("ymin", "y1")
                xmax, ymax = _get("xmax", "x2"), _get("ymax", "y2")
                if xmin is None or ymin is None or xmax is None or ymax is None:
                    continue
                boxes.append([float(xmin), float(ymin), float(xmax), float(ymax)])
                labels.append(str(_get("class_name", "label", "class") or "polyp"))
        out[f.stem] = {"boxes": boxes, "labels": labels, "width": None, "height": None}
    return out


def _load_bbox_json_dir(json_dir: Path) -> Dict[str, dict]:
    """Parse a bbox/ directory of one JSON file per image (an alternate
    layout seen in some Kvasir-SEG re-packagings)."""
    out: Dict[str, dict] = {}
    for f in sorted(json_dir.glob("*.json")):
        out[f.stem] = json.loads(f.read_text())
    return out


def load_bboxes(path: Path) -> Dict[str, dict]:
    """Returns {image_stem: {"boxes": [[x1,y1,x2,y2], ...], "labels": [...],
    "width": W or None, "height": H or None}}.

    Dispatches on what `path` (from find_bbox_json) actually is:
      - a directory of *.csv files  -> Kvasir-SEG's real per-image CSV format
      - a directory of *.json files -> per-image JSON (alternate layout)
      - a single merged json file   -> {image_stem: {...}, ...}
    """
    path = Path(path)
    if path.is_dir():
        if any(path.glob("*.csv")):
            return _load_bbox_csv_dir(path)
        raw = _load_bbox_json_dir(path)
    else:
        raw = json.loads(path.read_text())

    out: Dict[str, dict] = {}
    for key, val in raw.items():
        stem = Path(str(key)).stem
        if "boxes" in val:                   # already normalised (CSV path)
            out[stem] = val
            continue
        boxes, labels = [], []
        items = val.get("bbox", val.get("bboxes", []))
        if isinstance(items, dict):
            items = [items]
        for b in items:
            if isinstance(b, dict):
                x1, y1 = float(b.get("xmin", b.get("x1", 0))), float(b.get("ymin", b.get("y1", 0)))
                x2, y2 = float(b.get("xmax", b.get("x2", 0))), float(b.get("ymax", b.get("y2", 0)))
                lab = str(b.get("label", "polyp"))
            else:                                   # already [x1,y1,x2,y2]
                x1, y1, x2, y2 = [float(v) for v in b[:4]]
                lab = "polyp"
            boxes.append([x1, y1, x2, y2])
            labels.append(lab)
        out[stem] = {"boxes": boxes, "labels": labels,
                     "width": val.get("width"), "height": val.get("height")}
    return out


# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------


# ---------------------------------------------------------------------------
# Split CSVs store absolute paths; re-anchor them when the dataset lives
# elsewhere at read time (node-local /scratch, a different $HOME, a moved
# kagglehub cache).  Override with $KVASIR_DATA_ROOT.
KVASIR_ANCHOR = "Kvasir-SEG"


def _reanchor(path: str, new_root: Path, anchor: str) -> str:
    if not isinstance(path, str) or not path:
        return path
    parts = Path(path).parts
    if anchor in parts:
        i = len(parts) - 1 - parts[::-1].index(anchor)
        return str(Path(new_root).joinpath(*parts[i + 1:]))
    return str(Path(new_root) / Path(path).name)


def read_split_csv(csv_path: Path):
    """Load the detection split CSV with automatic path relocation."""
    import pandas as pd

    if not Path(csv_path).exists():
        raise FileNotFoundError(
            f"Split file {csv_path} not found. Run pre_process.py first.")
    df = pd.read_csv(csv_path)
    if "image_path" not in df.columns or len(df) == 0:
        return df
    sample = next((p for p in df["image_path"].head(50)
                   if isinstance(p, str) and p), None)
    if sample is None or Path(sample).exists():
        return df

    env_root = os.environ.get("KVASIR_DATA_ROOT")
    new_root = Path(env_root) if env_root else None
    if new_root is None:
        for c in [DATASET_DIR / KVASIR_ANCHOR, DATASET_DIR]:
            if (c / "images").is_dir():
                new_root = c
                break
        if new_root is None and DATASET_DIR.exists():
            hits = [p for p in DATASET_DIR.rglob(KVASIR_ANCHOR) if (p / "images").is_dir()]
            new_root = hits[0] if hits else None
    if new_root is None or not Path(new_root).is_dir():
        raise FileNotFoundError(
            f"Paths in the split file no longer resolve on this machine "
            f"(e.g. {sample}).\nSet KVASIR_DATA_ROOT to the folder containing "
            f"images/ and masks/, or place the dataset under {DATASET_DIR}.")
    for col in ("image_path", "mask_path"):
        if col in df.columns:
            df[col] = df[col].map(lambda p: _reanchor(p, Path(new_root), KVASIR_ANCHOR))
    print(f"[info] split paths re-anchored onto {new_root}")
    if not Path(df["image_path"].iloc[0]).exists():
        raise FileNotFoundError(
            f"Re-anchoring onto {new_root} still does not resolve "
            f"({df['image_path'].iloc[0]}). Check the folder layout.")
    return df


class KvasirDetectionDataset:
    """torchvision-detection-style dataset.

    __getitem__ returns (image_float_tensor[3,H,W] in [0,1], target dict) where
        target = {"boxes": FloatTensor[N,4] in xyxy,
                  "labels": Int64Tensor[N] (all 1 == polyp),
                  "image_id": Int64Tensor[1],
                  "area": FloatTensor[N],
                  "iscrowd": Int64Tensor[N],
                  "stem": str}
    """

    def __init__(self, csv_path: Path, split: str, augment: bool = False,
                 img_size: Optional[int] = None):
        df = read_split_csv(csv_path)
        if split != "all":
            df = df[df["split"] == split]
        self.df = df.reset_index(drop=True)
        self.augment = augment
        self.img_size = img_size

    def __len__(self) -> int:
        return len(self.df)

    def __getitem__(self, i: int):
        import torch
        from PIL import Image

        row = self.df.iloc[i]
        img = Image.open(row["image_path"]).convert("RGB")
        boxes = np.array(json.loads(row["boxes"]), dtype=np.float32).reshape(-1, 4)

        if self.img_size:
            w0, h0 = img.size
            img = img.resize((self.img_size, self.img_size), Image.BILINEAR)
            if len(boxes):
                boxes[:, [0, 2]] *= self.img_size / w0
                boxes[:, [1, 3]] *= self.img_size / h0

        arr = np.asarray(img, dtype=np.float32) / 255.0          # HWC
        if self.augment:
            arr, boxes = augment_detection(arr, boxes)

        x = torch.from_numpy(arr.transpose(2, 0, 1).copy())
        b = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
        target = {
            "boxes": b,
            "labels": torch.ones((len(b),), dtype=torch.int64),
            "image_id": torch.tensor([i]),
            "area": (b[:, 2] - b[:, 0]) * (b[:, 3] - b[:, 1]) if len(b) else torch.zeros((0,)),
            "iscrowd": torch.zeros((len(b),), dtype=torch.int64),
            "stem": str(row["image_id"]),
        }
        return x, target


def augment_detection(arr: np.ndarray, boxes: np.ndarray):
    """Seeded augmentation for colonoscopy frames.

    Horizontal AND vertical flips are used here (unlike the chest X-rays in Part
    A): an endoscope frame has no canonical anatomical orientation, so both are
    label-preserving. Photometric jitter mimics illumination and white-balance
    variation between scopes.
    """
    h, w = arr.shape[:2]
    if np.random.rand() < 0.5:                       # horizontal flip
        arr = arr[:, ::-1]
        if len(boxes):
            boxes = boxes.copy()
            boxes[:, [0, 2]] = w - boxes[:, [2, 0]]
    if np.random.rand() < 0.5:                       # vertical flip
        arr = arr[::-1]
        if len(boxes):
            boxes = boxes.copy()
            boxes[:, [1, 3]] = h - boxes[:, [3, 1]]
    if np.random.rand() < 0.7:                       # brightness / contrast
        g = 1.0 + np.random.uniform(-0.2, 0.2)
        b = np.random.uniform(-0.08, 0.08)
        arr = np.clip(arr * g + b, 0.0, 1.0)
    if np.random.rand() < 0.3:                       # mild saturation shift
        gray = arr.mean(axis=2, keepdims=True)
        s = 1.0 + np.random.uniform(-0.25, 0.25)
        arr = np.clip(gray + (arr - gray) * s, 0.0, 1.0)
    return np.ascontiguousarray(arr, dtype=np.float32), boxes.astype(np.float32)


def collate_fn(batch):
    return tuple(zip(*batch))


def make_loader(dataset, batch_size: int, shuffle: bool, num_workers: int = 2):
    import torch
    from torch.utils.data import DataLoader

    g = torch.Generator()
    g.manual_seed(SEED)
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle,
                      num_workers=num_workers, collate_fn=collate_fn,
                      worker_init_fn=worker_init_fn, generator=g,
                      pin_memory=torch.cuda.is_available())


# ---------------------------------------------------------------------------
# Detection metrics
# ---------------------------------------------------------------------------

def box_iou_np(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """a: (N,4), b: (M,4) in xyxy -> (N,M) IoU matrix."""
    if len(a) == 0 or len(b) == 0:
        return np.zeros((len(a), len(b)), dtype=np.float32)
    ax1, ay1, ax2, ay2 = [a[:, i][:, None] for i in range(4)]
    bx1, by1, bx2, by2 = [b[:, i][None, :] for i in range(4)]
    iw = np.clip(np.minimum(ax2, bx2) - np.maximum(ax1, bx1), 0, None)
    ih = np.clip(np.minimum(ay2, by2) - np.maximum(ay1, by1), 0, None)
    inter = iw * ih
    area_a = np.clip(ax2 - ax1, 0, None) * np.clip(ay2 - ay1, 0, None)
    area_b = np.clip(bx2 - bx1, 0, None) * np.clip(by2 - by1, 0, None)
    return (inter / np.clip(area_a + area_b - inter, 1e-9, None)).astype(np.float32)


def match_detections(preds: List[dict], gts: List[dict], iou_thr: float,
                     score_thr: float = 0.0):
    """Greedy score-ordered matching, one GT per detection (standard VOC/COCO rule).

    preds[i] = {"boxes": (N,4), "scores": (N,)}, gts[i] = {"boxes": (M,4)}
    Returns (tp, fp, scores, matched_ious, n_gt) as flat arrays over the dataset.
    """
    recs = []
    n_gt = 0
    for img_i, (p, g) in enumerate(zip(preds, gts)):
        pb = np.asarray(p["boxes"], dtype=np.float32).reshape(-1, 4)
        ps = np.asarray(p["scores"], dtype=np.float32).reshape(-1)
        keep = ps >= score_thr
        pb, ps = pb[keep], ps[keep]
        gb = np.asarray(g["boxes"], dtype=np.float32).reshape(-1, 4)
        n_gt += len(gb)
        if len(pb) == 0:
            continue
        order = np.argsort(-ps)
        pb, ps = pb[order], ps[order]
        ious = box_iou_np(pb, gb)
        taken = np.zeros(len(gb), dtype=bool)
        for d in range(len(pb)):
            best_j, best_iou = -1, 0.0
            for j in range(len(gb)):
                if taken[j]:
                    continue
                if ious[d, j] > best_iou:
                    best_iou, best_j = ious[d, j], j
            if best_j >= 0 and best_iou >= iou_thr:
                taken[best_j] = True
                recs.append((ps[d], 1, best_iou, img_i))
            else:
                recs.append((ps[d], 0, best_iou, img_i))
    if not recs:
        return (np.zeros(0), np.zeros(0), np.zeros(0), np.zeros(0), n_gt)
    recs.sort(key=lambda r: -r[0])
    scores = np.array([r[0] for r in recs], dtype=np.float64)
    tp = np.array([r[1] for r in recs], dtype=np.float64)
    fp = 1.0 - tp
    ious = np.array([r[2] for r in recs], dtype=np.float64)
    return tp, fp, scores, ious, n_gt


def average_precision(preds, gts, iou_thr: float = 0.5) -> dict:
    """All-point-interpolated AP (the COCO / VOC2010+ convention)."""
    tp, fp, scores, ious, n_gt = match_detections(preds, gts, iou_thr)
    if n_gt == 0:
        return {"ap": float("nan"), "precision": float("nan"), "recall": float("nan"),
                "n_gt": 0, "n_pred": int(len(tp))}
    if len(tp) == 0:
        return {"ap": 0.0, "precision": 0.0, "recall": 0.0, "n_gt": n_gt, "n_pred": 0}
    ctp, cfp = np.cumsum(tp), np.cumsum(fp)
    rec = ctp / n_gt
    prec = ctp / np.clip(ctp + cfp, 1e-9, None)
    # monotone-decreasing precision envelope
    mrec = np.concatenate([[0.0], rec, [rec[-1]]])
    mpre = np.concatenate([[1.0], prec, [0.0]])
    for i in range(len(mpre) - 2, -1, -1):
        mpre[i] = max(mpre[i], mpre[i + 1])
    idx = np.where(mrec[1:] != mrec[:-1])[0]
    ap = float(np.sum((mrec[idx + 1] - mrec[idx]) * mpre[idx + 1]))
    return {"ap": ap, "precision_curve": prec.tolist(), "recall_curve": rec.tolist(),
            "n_gt": int(n_gt), "n_pred": int(len(tp))}


def coco_map(preds, gts, thresholds: Optional[Sequence[float]] = None) -> dict:
    thresholds = list(thresholds) if thresholds is not None else \
        [round(x, 2) for x in np.arange(0.50, 0.96, 0.05)]
    per = {f"AP@{t:.2f}": average_precision(preds, gts, t)["ap"] for t in thresholds}
    vals = [v for v in per.values() if not np.isnan(v)]
    per["mAP@[.50:.05:.95]"] = float(np.mean(vals)) if vals else float("nan")
    return per


def prf_at_threshold(preds, gts, iou_thr: float = 0.5, score_thr: float = 0.5) -> dict:
    """Operating-point metrics at a fixed confidence threshold: precision, recall,
    F1, TP/FP/FN counts and the mean IoU of the matched detections."""
    tp, fp, scores, ious, n_gt = match_detections(preds, gts, iou_thr, score_thr=score_thr)
    n_tp, n_fp = int(tp.sum()), int(fp.sum())
    n_fn = int(n_gt - n_tp)
    prec = n_tp / max(1, n_tp + n_fp)
    rec = n_tp / max(1, n_gt)
    f1 = 0.0 if (prec + rec) == 0 else 2 * prec * rec / (prec + rec)
    matched_iou = float(ious[tp > 0].mean()) if (tp > 0).any() else float("nan")
    return {"iou_threshold": iou_thr, "score_threshold": score_thr,
            "precision": float(prec), "recall": float(rec), "f1": float(f1),
            "true_positives": n_tp, "false_positives": n_fp, "false_negatives": n_fn,
            "n_ground_truth": int(n_gt), "mean_iou_matched": matched_iou}


def size_stratified_recall(preds, gts, image_areas: Sequence[float],
                           iou_thr: float = 0.5, score_thr: float = 0.5,
                           bins=((0.0, 0.01, "small"), (0.01, 0.05, "medium"),
                                 (0.05, 1.01, "large"))) -> dict:
    """Recall per polyp-size bucket, where size = box area / image area.
    This is the 'performance on small polyps' number the loss ablation needs."""
    out = {name: {"n_gt": 0, "n_detected": 0} for _, _, name in bins}
    for p, g, img_area in zip(preds, gts, image_areas):
        pb = np.asarray(p["boxes"], dtype=np.float32).reshape(-1, 4)
        ps = np.asarray(p["scores"], dtype=np.float32).reshape(-1)
        keep = ps >= score_thr
        pb = pb[keep]
        gb = np.asarray(g["boxes"], dtype=np.float32).reshape(-1, 4)
        if len(gb) == 0:
            continue
        ious = box_iou_np(pb, gb) if len(pb) else np.zeros((0, len(gb)))
        best = ious.max(axis=0) if len(pb) else np.zeros(len(gb))
        for j in range(len(gb)):
            frac = float((gb[j, 2] - gb[j, 0]) * (gb[j, 3] - gb[j, 1]) / max(1.0, img_area))
            for lo, hi, name in bins:
                if lo <= frac < hi:
                    out[name]["n_gt"] += 1
                    out[name]["n_detected"] += int(best[j] >= iou_thr)
                    break
    for name in out:
        n = out[name]["n_gt"]
        out[name]["recall"] = out[name]["n_detected"] / n if n else float("nan")
    return out


# ----------------------------------------------------------------------------
# 9.  Version-compatibility shims (torch AMP API, matplotlib boxplot API)
# ----------------------------------------------------------------------------

def make_grad_scaler(enabled: bool):
    import torch
    try:
        return torch.amp.GradScaler("cuda", enabled=enabled)      # torch >= 2.4
    except (AttributeError, TypeError):
        return torch.cuda.amp.GradScaler(enabled=enabled)         # older torch


def autocast(enabled: bool):
    import torch
    try:
        return torch.amp.autocast("cuda", enabled=enabled)        # torch >= 2.4
    except (AttributeError, TypeError):
        return torch.cuda.amp.autocast(enabled=enabled)


def boxplot(ax, data, labels, **kw):
    """matplotlib renamed `labels` to `tick_labels` in 3.9."""
    try:
        return ax.boxplot(data, tick_labels=labels, **kw)
    except TypeError:
        return ax.boxplot(data, labels=labels, **kw)