#!/usr/bin/env python3
"""Like-for-like attribute comparison of two cross-validation runs.

    python scripts/compare_attributes.py --data-dir DATA --runs-a <folder with fold_*.pt> --runs-b <folder> [--device mps]

Pooled attribute accuracy in cv_report.md is computed over the nuclei each model
matched (IoU >= 0.5), so a model that matches more nuclei is scored on a harder
set. This script predicts every held-out image with both fold checkpoints and
reports the antibody / Nostril attribute accuracy on the nuclei matched by BOTH,
plus the per-label count errors. It makes no choice; it only reads out numbers.
"""
from __future__ import annotations

import argparse
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 ALL_LABELS, cell_attributes  # noqa: E402
from cellquant.metrics import match_by_iou  # noqa: E402
from cellquant.postprocess import PostConfig  # noqa: E402


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True)
    ap.add_argument("--runs-a", required=True)
    ap.add_argument("--runs-b", required=True)
    ap.add_argument("--names", default="v1,v2")
    ap.add_argument("--device", default="auto")
    ap.add_argument("--out", default="")
    a = ap.parse_args(argv)
    na, nb = a.names.split(",")
    samples, _ = discover(a.data_dir, verbose=False)
    tot = {na: {"ab": 0, "nos": 0, "both": 0}, nb: {"ab": 0, "nos": 0, "both": 0}, "n": 0}
    neg = {na: {"tp": 0, "fp": 0, "fn": 0}, nb: {"tp": 0, "fp": 0, "fn": 0}}
    counts = {lab: {"gt": 0, na: 0, nb: 0, f"abs_{na}": 0, f"abs_{nb}": 0} for lab in ALL_LABELS}
    rows = []
    for s in samples:
        img = load_image(s.image_path)
        R = rasterize(s.annotation, img.shape[0], img.shape[1])
        preds = {}
        for name, d in ((na, a.runs_a), (nb, a.runs_b)):
            ck = os.path.join(d, f"fold_{s.key}.pt")
            res, _ = Predictor(ck, device=a.device, tta=1).predict(img, PostConfig())
            preds[name] = res
            for lab in ALL_LABELS:
                counts[lab][name] += res.counts[lab]
                counts[lab][f"abs_{name}"] += abs(res.counts[lab] - s.annotation.count(lab))
        for lab in ALL_LABELS:
            counts[lab]["gt"] += s.annotation.count(lab)
        m = {name: {g: p for g, p, _ in match_by_iou(R.cell_inst, preds[name].cell_inst, 0.5)} for name in preds}
        common = sorted(set(m[na]) & set(m[nb]))
        row = {"image": s.key, "n_common": len(common), f"n_{na}": len(m[na]), f"n_{nb}": len(m[nb])}
        for name in (na, nb):
            ab = ns = both = 0
            for g in common:
                ga, gn = cell_attributes(R.cell_labels[g])
                pa, pn = cell_attributes(preds[name].cell_labels[m[name][g]])
                ab += ga == pa
                ns += gn == pn
                both += (ga == pa) and (gn == pn)
                if not ga and not pa:
                    neg[name]["tp"] += 1
                elif ga and not pa:
                    neg[name]["fp"] += 1
                elif not ga and pa:
                    neg[name]["fn"] += 1
            row[f"ab_{name}"] = ab / max(1, len(common))
            row[f"nos_{name}"] = ns / max(1, len(common))
            row[f"both_{name}"] = both / max(1, len(common))
            tot[name]["ab"] += ab
            tot[name]["nos"] += ns
            tot[name]["both"] += both
        tot["n"] += len(common)
        rows.append(row)
        print(f"[{s.key}] common {len(common)}: antibody acc {na} {row[f'ab_{na}']:.3f} {nb} {row[f'ab_{nb}']:.3f} | "
              f"nostril acc {row[f'nos_{na}']:.3f} {row[f'nos_{nb}']:.3f} | both {row[f'both_{na}']:.3f} {row[f'both_{nb}']:.3f}", flush=True)
    n = max(1, tot["n"])
    L = [f"## Attribute accuracy on the {tot['n']} held-out nuclei matched by both {na} and {nb} (IoU >= 0.5, base run)", "",
         f"| | {na} | {nb} |", "|---|---:|---:|"]
    for k, label in (("ab", "Antibody +/− correct"), ("nos", "Nostril attribute correct"), ("both", "Both correct")):
        L.append(f"| {label} | {100*tot[na][k]/n:.1f}% | {100*tot[nb][k]/n:.1f}% |")
    for name in (na, nb):
        t = neg[name]
        p = t["tp"] / max(1, t["tp"] + t["fp"])
        r = t["tp"] / max(1, t["tp"] + t["fn"])
        f = 2 * p * r / max(1e-9, p + r)
        L.append(f"| Antibody − on this set ({name}): TP / FP / FN, precision, recall, F1 | {t['tp']} / {t['fp']} / {t['fn']}, {100*p:.1f}%, {100*r:.1f}%, {100*f:.1f}% | |")
    L += ["", f"| Image | matched by both | {na} antibody acc | {nb} antibody acc | {na} Nostril acc | {nb} Nostril acc |", "|---|---:|---:|---:|---:|---:|"]
    for r in rows:
        L.append(f"| {r['image']} | {r['n_common']} | {100*r[f'ab_{na}']:.1f}% | {100*r[f'ab_{nb}']:.1f}% | {100*r[f'nos_{na}']:.1f}% | {100*r[f'nos_{nb}']:.1f}% |")
    L += ["", "## Counts summed over the nine held-out images (and summed absolute per-image count error)", "",
          f"| Label | GT | {na} | {nb} | abs error {na} | abs error {nb} |", "|---|---:|---:|---:|---:|---:|"]
    for lab in ALL_LABELS:
        c = counts[lab]
        L.append(f"| {lab} | {c['gt']} | {c[na]} | {c[nb]} | {c[f'abs_{na}']} | {c[f'abs_{nb}']} |")
    L.append("")
    text = "\n".join(L)
    print(text)
    if a.out:
        with open(a.out, "a") as f:
            f.write("\n" + text)
        print("appended to", a.out)


if __name__ == "__main__":
    main()