"""Score the trained models under a ground-truth-selected slice protocol.

The paper argues that an evaluation which searches only the slices picked out by
the reference standard reverses the ranking of both factors under study. Stated
as a before-and-after that argument is confounded: the earlier analysis differed
from this one in the protocol *and* in five implementation defects *and* in
every trained weight.

This script removes the confound. It takes the models reported in Section 3.1 -
the same checkpoints, the same official folds, the same aggregation, the same
evaluator, the same exclusion list - and changes exactly one thing: instead of
searching every axial slice, it searches only the slices that contain an
annotated nodule centre. Scans with no annotation are searched not at all and
still count in the false-positive-per-scan denominator, which is what such a
protocol does.

Whatever differs between this output and Table 1 is therefore attributable to
the slice set and to nothing else.

Output
------
``candidates_restricted/<experiment>.parquet``
    Candidate lists under the restricted protocol.
``results/table_protocol_effect.csv``
    CPM under both protocols, per configuration, with the ranking under each.

Usage
-----
    python scripts/15_restricted_protocol.py [--force]
"""

from __future__ import annotations

import argparse
import sys
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import pandas as pd
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from luna_rev import config as cfg
from luna_rev import evaluate as ev
from luna_rev.io_luna import (all_uids, group_by_uid, load_annotations,
                              load_excluded, normalize_hu, read_volume)
from luna_rev.predict import (candidates_to_world, cluster_to_3d, detect_volume,
                              load_candidates)
from luna_rev.splits import get_folds
from luna_rev.train import weights_path

OUT_DIR = cfg.ROOT / "candidates_restricted"
MAIN = ["Exp1_2D_Loose", "Exp2_MIP_Loose", "Exp3_2p5D_Loose", "Exp4_2p5D_Strict"]


def restricted_slices(annotations: pd.DataFrame, uid: str, meta) -> list[int]:
    """Axial slices holding an annotated nodule centre, in voxel indices.

    This is the sharpest reading of "the reference standard chooses where to
    look": one slice per annotation, and nothing at all in a scan that has no
    annotation.
    """
    rows = annotations[annotations.seriesuid == uid]
    zs = []
    for _, r in rows.iterrows():
        _, _, cz = meta.world_to_voxel(np.array([r.coordX, r.coordY, r.coordZ]))
        zs.append(int(round(cz)))
    return zs


def predict_restricted(force: bool) -> None:
    from ultralytics import YOLO

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    todo = [n for n in MAIN if force or not (OUT_DIR / f"{n}.parquet").exists()]
    if not todo:
        print("restricted candidates already present")
        return

    annotations = load_annotations()
    rows: dict[str, list[pd.DataFrame]] = {n: [] for n in todo}
    timing: dict[str, float] = defaultdict(float)

    for fold in get_folds("official"):
        models = {n: YOLO(str(weights_path(cfg.EXPERIMENTS_BY_NAME[n], fold.index)))
                  for n in todo}
        for uid in tqdm(fold.test, desc=f"fold{fold.index} restricted"):
            t0 = time.time()
            vol_hu, meta = read_volume(uid)
            vol_u8 = normalize_hu(vol_hu)
            del vol_hu
            timing["read"] += time.time() - t0

            zs = restricted_slices(annotations, uid, meta)
            if not zs:
                continue                      # nodule-free scan contributes nothing
            for name in todo:
                exp = cfg.EXPERIMENTS_BY_NAME[name]
                t1 = time.time()
                det = detect_volume(models[name], vol_u8, exp.representation, slices=zs)
                timing["infer"] += time.time() - t1
                clusters = cluster_to_3d(det)[:150]
                if len(clusters) == 0:
                    continue
                world = candidates_to_world(clusters, meta)
                rows[name].append(pd.DataFrame({
                    "seriesuid": uid,
                    "coordX": world[:, 0], "coordY": world[:, 1], "coordZ": world[:, 2],
                    "probability": clusters[:, 5],
                    "slice_z": clusters[:, 0].astype(int),
                    "fold": fold.index,
                }))
        del models

    for name in todo:
        df = pd.concat(rows[name], ignore_index=True) if rows[name] else pd.DataFrame(
            columns=["seriesuid", "coordX", "coordY", "coordZ", "probability",
                     "slice_z", "fold"])
        df.to_parquet(OUT_DIR / f"{name}.parquet", index=False)
        print(f"  {name:20s} {len(df):6d} candidates over {df.seriesuid.nunique()} scans")
    total = sum(timing.values()) or 1.0
    print("[restricted] " + ", ".join(f"{k}={v / 60:.1f}min ({100 * v / total:.0f}%)"
                                      for k, v in timing.items()))


def searched_slices() -> tuple[int, float, float]:
    """How many slices the restricted protocol actually searches.

    Read from the image headers rather than from the candidate lists: a slice
    that produced no detection was still searched, so counting candidates would
    understate it.
    """
    import SimpleITK as sitk

    ann = load_annotations()
    by_uid = {u: g for u, g in ann.groupby("seriesuid")}
    paths = {q.stem: q for q in cfg.RAW_PATH.rglob("*.mhd")}
    per = []
    for uid in all_uids():
        g = by_uid.get(uid)
        if g is None:
            per.append(0)
            continue
        r = sitk.ImageFileReader()
        r.SetFileName(str(paths[uid]))
        r.ReadImageInformation()
        o = np.array(r.GetOrigin())
        sp = np.array(r.GetSpacing())
        D = np.array(r.GetDirection()).reshape(3, 3)
        zs = {int(round((np.linalg.inv(D) @
                         (np.array([row.coordX, row.coordY, row.coordZ]) - o) / sp)[2]))
              for _, row in g.iterrows()}
        per.append(len(zs))
    v = np.asarray(per)
    return int(v.sum()), float(v.mean()), float(v[v > 0].mean())


def load_restricted(name: str) -> dict[str, np.ndarray]:
    df = pd.read_parquet(OUT_DIR / f"{name}.parquet")
    return {str(uid): g[["coordX", "coordY", "coordZ", "probability"]].to_numpy(float)
            for uid, g in df.groupby("seriesuid", sort=False)}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    predict_restricted(args.force)

    uids = list(all_uids())
    included, excluded = group_by_uid(load_annotations()), group_by_uid(load_excluded())
    n_searched, mean_all, mean_pos = searched_slices()
    print(f"\nrestricted protocol searches {n_searched} slices "
          f"({mean_all:.2f} per scan over all {len(uids)}, {mean_pos:.2f} per scan "
          f"over the scans that have an annotation)")

    rows = []
    for name in MAIN:
        whole = ev.evaluate(ev.match(load_candidates(name), included, excluded, uids))
        restr_c = load_restricted(name)
        restr = ev.evaluate(ev.match(restr_c, included, excluded, uids))
        n_slices = sum(len(v) for v in restr_c.values())
        rows.append({
            "experiment": name,
            "label": cfg.EXPERIMENTS_BY_NAME[name].label,
            "cpm_whole_volume": round(whole["cpm"], 4),
            "cpm_restricted": round(restr["cpm"], 4),
            "delta": round(restr["cpm"] - whole["cpm"], 4),
            "candidates_whole_volume": int(whole["total_candidates"]),
            "candidates_restricted": int(restr["total_candidates"]),
            "scans_with_candidates_restricted": len(restr_c),
            "candidates_per_scan_restricted": round(n_slices / len(uids), 2),
            "slices_searched_total": n_searched,
            "sens_at_saturation_whole": round(whole["sensitivity_at_saturation"], 4),
            "sens_at_saturation_restricted": round(restr["sensitivity_at_saturation"], 4),
        })
    df = pd.DataFrame(rows)
    df["rank_whole_volume"] = df.cpm_whole_volume.rank(ascending=False).astype(int)
    df["rank_restricted"] = df.cpm_restricted.rank(ascending=False).astype(int)
    out = cfg.RESULTS_DIR / "table_protocol_effect.csv"
    df.to_csv(out, index=False)

    pd.set_option("display.width", 200)
    print("\n" + df.drop(columns=["label"]).to_string(index=False))
    print(f"\nwrote {out}")

    # Paired bootstrap within each protocol, on resamples shared by all four
    # configurations, so each contrast carries an interval instead of being read
    # off two point estimates.
    vec_whole = {n: ev.match(load_candidates(n), included, excluded, uids) for n in MAIN}
    vec_restr = {n: ev.match(load_restricted(n), included, excluded, uids) for n in MAIN}
    boot_whole = ev.bootstrap_cpm(vec_whole, uids, n_iter=cfg.N_BOOTSTRAP,
                                  seed=cfg.ANALYSIS_SEED)
    boot_restr = ev.bootstrap_cpm(vec_restr, uids, n_iter=cfg.N_BOOTSTRAP,
                                  seed=cfg.ANALYSIS_SEED)

    w = df.set_index("experiment")
    lines = []
    print("\nContrasts that the paper's conclusions rest on:")
    for a, b, what in [("Exp3_2p5D_Loose", "Exp1_2D_Loose", "representation (2.5D - 2D)"),
                       ("Exp4_2p5D_Strict", "Exp3_2p5D_Loose", "supervision (strict - loose)")]:
        dw = w.loc[a, "cpm_whole_volume"] - w.loc[b, "cpm_whole_volume"]
        dr = w.loc[a, "cpm_restricted"] - w.loc[b, "cpm_restricted"]
        pw = ev.paired_difference(boot_whole[a], boot_whole[b], dw)
        pr = ev.paired_difference(boot_restr[a], boot_restr[b], dr)
        flip = "SIGN FLIPS" if dw * dr < 0 else "same sign"
        print(f"  {what}")
        print(f"    whole-volume {dw:+.4f} [{pw['ci_low']:+.4f}, {pw['ci_high']:+.4f}] "
              f"p={pw['p_bootstrap_two_sided']:.3f}")
        print(f"    restricted   {dr:+.4f} [{pr['ci_low']:+.4f}, {pr['ci_high']:+.4f}] "
              f"p={pr['p_bootstrap_two_sided']:.3f}   -> {flip}")
        lines.append({"contrast": what,
                      "delta_whole_volume": round(dw, 4),
                      "ci_low_whole": round(pw["ci_low"], 4),
                      "ci_high_whole": round(pw["ci_high"], 4),
                      "p_whole": round(pw["p_bootstrap_two_sided"], 4),
                      "delta_restricted": round(dr, 4),
                      "ci_low_restricted": round(pr["ci_low"], 4),
                      "ci_high_restricted": round(pr["ci_high"], 4),
                      "p_restricted": round(pr["p_bootstrap_two_sided"], 4),
                      "sign_flips": "reverses" if dw * dr < 0 else "preserved"})
    out2 = cfg.RESULTS_DIR / "table_protocol_contrasts.csv"
    pd.DataFrame(lines).to_csv(out2, index=False)
    print(f"\nwrote {out2}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())