#!/usr/bin/env python3
"""Sweep post-processing thresholds on held-out fold predictions (no retraining).

    python scripts/postprocess_sweep.py --data-dir DATA --runs-dir runs/cv [--keys A,B] [--device mps]

For every fold checkpoint the held-out image is predicted once; every PostConfig in
the grid is then scored on the same probabilities. Prints micro-averaged nucleus
detection F1 (IoU >= 0.5 and centroid rule) and per-class F1 per setting.
"""
from __future__ import annotations

import argparse
import itertools
import json
import os
import sys

import numpy as np

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


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True)
    ap.add_argument("--runs-dir", default="runs/cv")
    ap.add_argument("--keys", default="")
    ap.add_argument("--device", default="auto")
    ap.add_argument("--out", default="")
    a = ap.parse_args(argv)
    samples, _ = discover(a.data_dir, verbose=False)
    keys = a.keys.split(",") if a.keys else [s.key for s in samples]
    grid = [PostConfig(marker_thr=m, cell_fg_thr=f, min_marker_area=ma)
            for m, f, ma in itertools.product((0.5, 0.6, 0.7, 0.8), (0.4, 0.5), (12, 30))]
    per_cfg = {i: [] for i in range(len(grid))}
    for s in samples:
        if s.key not in keys:
            continue
        ck = os.path.join(a.runs_dir, f"fold_{s.key}.pt")
        if not os.path.exists(ck):
            print("missing", ck)
            continue
        img = load_image(s.image_path)
        R = rasterize(s.annotation, img.shape[0], img.shape[1])
        probs = Predictor(ck, device=a.device, tta=1).probs(img)
        for i, cfg in enumerate(grid):
            res = analyze(probs, cfg)
            ce = evaluate_cells(R.cell_inst, R.cell_labels, res.cell_inst, res.cell_labels)
            ne = evaluate_nostrils(R.nostril_inst, R.nostril_labels, res.nostril_inst, res.nostril_labels)
            per_cfg[i].append({"key": s.key, "det": ce["detection"], "cells": ce["per_class"], "nos": ne["per_class"],
                               "attr": ce["attribute_accuracy"]})
        print("done", s.key, flush=True)
    rows = []
    for i, cfg in enumerate(grid):
        rs = per_cfg[i]
        if not rs:
            continue
        tp = sum(r["det"]["tp"] for r in rs); fp = sum(r["det"]["fp"] for r in rs); fn = sum(r["det"]["fn"] for r in rs)
        ctp = sum(r["det"]["centroid"]["tp"] for r in rs); cfp = sum(r["det"]["centroid"]["fp"] for r in rs); cfn = sum(r["det"]["centroid"]["fn"] for r in rs)
        mc = micro_average([r["cells"] for r in rs], CELL_LABELS)
        mn = micro_average([r["nos"] for r in rs], NOSTRIL_LABELS)
        per_img = {r["key"]: round(r["det"]["f1"], 3) for r in rs}
        rows.append({"marker_thr": cfg.marker_thr, "cell_fg_thr": cfg.cell_fg_thr, "min_marker_area": cfg.min_marker_area,
                     "nuc_f1_iou50": prf(tp, fp, fn)["f1"], "nuc_f1_centroid": prf(ctp, cfp, cfn)["f1"],
                     "f1_ab_pos": mc["Cell, Antibody +"]["f1"], "f1_ab_neg": mc["Cell, Antibody -"]["f1"],
                     "f1_nos_pos": mc["Cell, Nostril, Antibody +"]["f1"], "f1_pin": mn["Nostril in Nucleus"]["f1"],
                     "per_image": per_img})
    rows.sort(key=lambda r: -r["nuc_f1_iou50"])
    print(f"{'marker':>6} {'fg':>4} {'minA':>4} | {'nucF1@.5':>8} {'nucF1cen':>8} | {'Ab+':>6} {'Ab-':>6} {'NosAb+':>6} {'PinNuc':>6} | per image")
    for r in rows:
        print(f"{r['marker_thr']:6.2f} {r['cell_fg_thr']:4.1f} {r['min_marker_area']:4d} | {100*r['nuc_f1_iou50']:8.1f} {100*r['nuc_f1_centroid']:8.1f} | "
              f"{100*r['f1_ab_pos']:6.1f} {100*r['f1_ab_neg']:6.1f} {100*r['f1_nos_pos']:6.1f} {100*r['f1_pin']:6.1f} | "
              + " ".join(f"{k}:{100*v:.0f}" for k, v in r['per_image'].items()))
    if a.out:
        with open(a.out, "w") as f:
            json.dump(rows, f, indent=1)


if __name__ == "__main__":
    main()