#!/usr/bin/env python3
"""Collect a final-eval directory into one mean +/- std table across seeds.

Reads every `*/summary.json` under the run root, pulls the held-out metrics for
each method, and aggregates across seeds. LatticeSpike and the baselines store
their metrics differently -- LatticeSpike as a dict under `eval`, the baseline
runner as a list of per-method dicts -- so both shapes are handled here rather
than at the call site.

Two reporting rules, both because AP is base-rate dependent:

* the random-AP floor (positive prevalence) is printed alongside, since an AP is
  meaningless without it and is never comparable across datasets;
* `n_metric_tokens` is printed per method, because a table is only a comparison
  if every row scored the same tokens. A mismatch is flagged loudly.
"""

from __future__ import annotations

import glob
import json
import os
import statistics
import sys


def load(root):
    """method -> seed -> metrics, from both summary shapes."""
    out = {}
    prevalence = {}
    for path in sorted(glob.glob(os.path.join(root, "*", "summary.json"))):
        name = os.path.basename(os.path.dirname(path))
        seed = name.rsplit("_s", 1)[-1]
        try:
            summary = json.load(open(path))
        except Exception:
            continue
        evaluation = summary.get("eval")
        rows = []
        if isinstance(evaluation, dict):          # LatticeSpike
            rows = [dict(evaluation, method="LatticeSpike (ours)")]
        elif isinstance(evaluation, list):        # baseline runner
            rows = evaluation
        for row in rows:
            method = row.get("method")
            if not method:
                continue
            out.setdefault(method, {})[seed] = row
            if row.get("positive_prevalence") is not None:
                prevalence[method] = row["positive_prevalence"]
    return out, prevalence


def agg(values):
    values = [v for v in values if isinstance(v, (int, float))]
    if not values:
        return None, None
    if len(values) == 1:
        return values[0], None
    return statistics.mean(values), statistics.stdev(values)


def fmt(mean, sd):
    if mean is None:
        return "--"
    return f"{mean:.4f}" + (f" ±{sd:.4f}" if sd is not None else "")


def main(argv):
    if len(argv) < 2:
        raise SystemExit("usage: collect_final_eval.py <run-root>")
    root = argv[1]
    data, prevalence = load(root)
    if not data:
        print("(no summaries found)")
        return 0

    supports = {
        m: {r.get("n_metric_tokens") for r in seeds.values()}
        for m, seeds in data.items()
    }
    all_supports = {s for values in supports.values() for s in values if s}

    ordered = sorted(
        data.items(),
        key=lambda kv: -(agg([r.get("auprc") for r in kv[1].values()])[0] or 0),
    )

    print(f"# Final evaluation — {os.path.basename(root)}\n")
    print(f"{'method':<30}{'AP':>18}{'AUROC':>18}{'P':>10}{'R':>10}{'F1':>10}{'seeds':>7}")
    print("-" * 103)
    for method, seeds in ordered:
        rows = list(seeds.values())
        cells = []
        for key in ("auprc", "auroc", "precision", "recall", "f1"):
            mean, sd = agg([r.get(key) for r in rows])
            cells.append((mean, sd))
        star = "*" if method.startswith("LatticeSpike") else " "
        print(
            f"{star}{method:<29}"
            f"{fmt(*cells[0]):>18}{fmt(*cells[1]):>18}"
            f"{cells[2][0] if cells[2][0] is None else f'{cells[2][0]:.4f}':>10}"
            f"{cells[3][0] if cells[3][0] is None else f'{cells[3][0]:.4f}':>10}"
            f"{cells[4][0] if cells[4][0] is None else f'{cells[4][0]:.4f}':>10}"
            f"{len(rows):>7}"
        )

    floor = statistics.median(prevalence.values()) if prevalence else None
    print()
    if floor is not None:
        print(f"random-AP floor (positive prevalence): {floor:.4f}")
    if len(all_supports) == 1:
        print(f"support: all methods scored {all_supports.pop():,} metric tokens — comparable")
    else:
        print(f"!! SUPPORT MISMATCH across methods: {sorted(all_supports)}")
        print("!! rows scored on different token sets are NOT directly comparable")
    seed_counts = {len(s) for s in data.values()}
    if len(seed_counts) > 1:
        print(f"!! uneven seed counts across methods: {sorted(seed_counts)} — a mean over")
        print("!! more seeds is not comparable to a single draw; equalise before reporting")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))