"""Repeatability + accuracy harness for the QUE-245 spike.

Runs each (approach, grid, model) config K times over every page of a
document and aggregates the three scored properties:

- independence: structural — no config calls Camelot (asserted in report).
- repeatability: across the K runs, the table-count stability (fraction of
  runs that yield the ground-truth count) and the per-table IoU spread
  (mean / std / min). The Cash Conversion Cycle drop shows up here as a
  count-stability < 1.0 on BHE_991 page 2.
- accuracy: mean IoU vs the full-anatomy ground truth, and the worst-case
  title clip in points.

Raw per-run boxes are saved to JSON so the report/overlays can be rebuilt
without re-calling the model.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import statistics
import tempfile
from pathlib import Path
from typing import Dict, List

import fitz

from experiments.que245.approaches import make_locator
from experiments.que245.gt import BHE_991_BANDS, ground_truth
from experiments.que245.score import iou, match, top_clip_pts

PDF = Path("documents/BHE_991.pdf")
OUT = Path("experiments/que245/results")


async def run_config(approach, model, temp, rows, cols, pages, k):
    loc = make_locator(approach, model=model, temperature=temp, rows=rows, cols=cols)
    doc = fitz.open(str(PDF))
    # pre-render page images once
    imgs = {}
    tmpdir = tempfile.mkdtemp(prefix="que245-rep-")
    for p in pages:
        ip = Path(tmpdir) / f"p{p}.png"
        doc[p - 1].get_pixmap(dpi=200).save(str(ip))
        imgs[p] = ip
    runs: Dict[int, List[List]] = {p: [] for p in pages}
    for _ in range(k):
        for p in pages:
            located = await loc.locate(imgs[p], PDF, p)
            runs[p].append([list(t.region) for t in located])
    return runs


def aggregate(runs, gt, page_h=792.0):
    rows_out = {}
    for p, run_list in runs.items():
        g = gt[p]
        counts = [len(r) for r in run_list]
        gt_n = len(g)
        count_stable = sum(1 for c in counts if c == gt_n) / len(counts)
        ious, clips = [], []
        for r in run_list:
            for gi, pj, v in match(g, r):
                if pj is not None:
                    ious.append(v)
                    clips.append(top_clip_pts(g[gi], r[pj], page_h))
        rows_out[p] = {
            "gt": gt_n,
            "counts": counts,
            "count_stable": count_stable,
            "mean_iou": statistics.mean(ious) if ious else 0.0,
            "iou_std": statistics.pstdev(ious) if len(ious) > 1 else 0.0,
            "min_iou": min(ious) if ious else 0.0,
            "max_clip_pts": max(clips) if clips else 0.0,
        }
    return rows_out


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--k", type=int, default=5)
    ap.add_argument("--pages", default="1-8")
    args = ap.parse_args()
    a, b = args.pages.split("-")
    pages = list(range(int(a), int(b) + 1))
    gt = ground_truth(PDF, BHE_991_BANDS)
    OUT.mkdir(parents=True, exist_ok=True)

    configs = [
        ("C raw-vision   ", "raw", "claude-haiku-4-5-20251001", 0.0, 24, 12),
        ("A data-grid r24", "data-grid", "claude-haiku-4-5-20251001", 0.0, 24, 12),
        ("B anatomy  r24", "anatomy", "claude-haiku-4-5-20251001", 0.0, 24, 12),
        ("B anatomy  r36", "anatomy", "claude-haiku-4-5-20251001", 0.0, 36, 12),
    ]
    all_results = {}
    for label, approach, model, temp, rows, cols in configs:
        print(f"\n### {label}  (k={args.k}) ###")
        runs = await run_config(approach, model, temp, rows, cols, pages, args.k)
        agg = aggregate(runs, gt)
        all_results[label] = {"agg": agg, "runs": runs}
        print(f"{'pg':>2} {'gt':>2} {'cnt-stable':>10} {'counts':>14} {'meanIoU':>8} {'std':>5} {'minIoU':>7} {'clip':>5}")
        for p in pages:
            a_ = agg[p]
            print(f"{p:>2} {a_['gt']:>2} {a_['count_stable']:>10.2f} {str(a_['counts']):>14} "
                  f"{a_['mean_iou']:>8.2f} {a_['iou_std']:>5.2f} {a_['min_iou']:>7.2f} {a_['max_clip_pts']:>5.0f}")
        macro = statistics.mean(a_["mean_iou"] for a_ in agg.values())
        stab = statistics.mean(a_["count_stable"] for a_ in agg.values())
        print(f"   macro mean-IoU={macro:.3f}  mean count-stability={stab:.3f}")

    (OUT / "bhe_repeat.json").write_text(json.dumps(all_results, indent=1))
    print(f"\nsaved {OUT / 'bhe_repeat.json'}")


if __name__ == "__main__":
    asyncio.run(main())