#!/usr/bin/env python3
"""How far does test-time augmentation help? Held-out scores of the fold models for
no TTA, 4x, 8x (the full dihedral group) and 8x combined with several image scales
(24x = 3 scales, 40x = 5 scales), with the per-fold peak-distance thresholds.

    python scripts/tta_study.py --data-dir DATA --runs-dir runs/cv_v3 \
        --post-config runs/cv_v3/rule_study_post_config.json --out reports/cv_v3/tta_study.md
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from cellquant.annotations import discover, load_image, rasterize  # noqa: E402
from cellquant.infer import Predictor  # noqa: E402
from cellquant.metrics import evaluate_cells, evaluate_nostrils, prf  # noqa: E402
from cellquant.postprocess import PostConfig, analyze  # noqa: E402

SETTINGS = {
    "none (1x)": dict(tta=1, scales=(1.0,)),
    "4x flips": dict(tta=4, scales=(1.0,)),
    "8x rotations+flips": dict(tta=8, scales=(1.0,)),
    "8x + 3 scales 0.8/1/1.25 (24x)": dict(tta=8, scales=(0.8, 1.0, 1.25)),
    "8x + 5 scales 0.8..1.25 (40x)": dict(tta=8, scales=(0.8, 0.9, 1.0, 1.1, 1.25)),
}


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True)
    ap.add_argument("--runs-dir", default="runs/cv_v3")
    ap.add_argument("--post-config", default="")
    ap.add_argument("--out", default="reports/cv_v3/tta_study.md")
    ap.add_argument("--device", default="auto")
    a = ap.parse_args(argv)
    overrides = json.load(open(a.post_config)) if a.post_config else {}
    samples, _ = discover(a.data_dir, verbose=False)
    res: dict[str, dict] = {name: {"cells": [], "nos": [], "cen": [], "n_pred": 0, "n_gt": 0, "sec": 0.0, "per_image": {}} for name in SETTINGS}
    for s in samples:
        img = load_image(s.image_path)
        R = rasterize(s.annotation, img.shape[0], img.shape[1])
        cfg = PostConfig(**overrides.get(s.key, {}))
        ck = os.path.join(a.runs_dir, f"fold_{s.key}.pt")
        for name, st in SETTINGS.items():
            pred = Predictor(ck, device=a.device, tta=st["tta"], scales=st["scales"])
            t0 = time.time()
            probs = pred.probs(img)
            dt = time.time() - t0
            out = analyze(probs, cfg)
            ce = evaluate_cells(R.cell_inst, R.cell_labels, out.cell_inst, out.cell_labels)
            ne = evaluate_nostrils(R.nostril_inst, R.nostril_labels, out.nostril_inst, out.nostril_labels)
            r = res[name]
            r["cells"].append(ce["detection"]); r["cen"].append(ce["detection"]["centroid"]); r["nos"].append(ne["detection"])
            r["n_pred"] += ce["n_pred"]; r["n_gt"] += ce["n_gt"]; r["sec"] += dt
            r["per_image"][s.key] = ce["detection"]["f1"]
            print(f"[{s.key}] {name}: nuclei F1 {ce['detection']['f1']:.3f} puncta F1 {ne['detection']['f1']:.3f} ({dt:.1f}s)", flush=True)
    L = ["# Test-time augmentation on the held-out images (v3 rule, per-fold thresholds)", "",
         "Each row averages the network output over the listed transforms before post-processing; the dihedral group "
         "(4 rotations x flip) has exactly 8 members, so more than 8x means adding image scales (the image is resized, "
         "predicted, and the probabilities resized back). Micro-averaged over the nine held-out images; "
         "seconds = mean inference time per 1024x1024 image on this machine's GPU (the Space's CPU is roughly 20-40x slower).", "",
         "| Setting | passes | Nucleus P | R | F1 (IoU 0.5) | F1 (centroid) | Puncta P | R | F1 | predicted nuclei / annotated | s per image |",
         "|---|---:|---:|---:|---:|---:|---:|---:|---:|---|---:|"]
    summary = {}
    for name, st in SETTINGS.items():
        r = res[name]
        def tot(rows):
            return prf(sum(x["tp"] for x in rows), sum(x["fp"] for x in rows), sum(x["fn"] for x in rows))
        c, cc, n = tot(r["cells"]), tot(r["cen"]), tot(r["nos"])
        passes = st["tta"] * len(st["scales"])
        summary[name] = {"passes": passes, "cells": c, "centroid": cc, "puncta": n, "n_pred": r["n_pred"], "n_gt": r["n_gt"],
                         "sec_per_image": r["sec"] / len(samples), "per_image_f1": r["per_image"]}
        L.append(f"| {name} | {passes} | {100*c['precision']:.1f} | {100*c['recall']:.1f} | **{100*c['f1']:.1f}** | {100*cc['f1']:.1f} | "
                 f"{100*n['precision']:.1f} | {100*n['recall']:.1f} | **{100*n['f1']:.1f}** | {r['n_pred']} / {r['n_gt']} | {r['sec']/len(samples):.1f} |")
    L += ["", "## Nucleus F1 per held-out image", "", "| Image | " + " | ".join(SETTINGS) + " |", "|---|" + "---:|" * len(SETTINGS)]
    for s in samples:
        L.append(f"| {s.key} | " + " | ".join(f"{100*res[n]['per_image'][s.key]:.1f}" for n in SETTINGS) + " |")
    L.append("")
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    with open(a.out, "w") as f:
        f.write("\n".join(L))
    with open(os.path.splitext(a.out)[0] + ".json", "w") as f:
        json.dump(summary, f, indent=1)
    print("wrote", a.out)


if __name__ == "__main__":
    main()