"""Appendix B.3: finite-sample stability of cFID and csFID (the ordering of the models and their relative gaps at N = 10,000).

    python reproduce/appendix/fid_sample_stability.py --system euler --samples direct_si=DIR --samples edm=DIR \\
        --samples cascaded_si=DIR --out DIR [--floor-ns 1000,2500,5000] [--model-ns 2000,5000,10000] [--device cuda]

  floor rows   FID(A_N, B_N), A and B disjoint halves of the truth ensemble (N <= 5,000: the bank holds 10,000 members)
  model rows   FID(gen_N, truth_N) for each Table 2 model
Members are the first N in the stored order, so every N is a nested subset; features are extracted once at the largest
N. Writes OUT/sample_stability.{csv,json}.
"""
from __future__ import annotations

import argparse
import csv
import math
import sys
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

import numpy as np

REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO))

from reproduce.appendix.common import SystemInputs, add_system_args, both_features, extractors, write_json  # noqa: E402
from bifurcation.fid import backend as BE  # noqa: E402
from bifurcation.fid.evaluate import bank_members  # noqa: E402


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    add_system_args(ap)
    ap.add_argument("--floor-ns", default="1000,2500,5000")
    ap.add_argument("--model-ns", default="2000,5000,10000")
    a = ap.parse_args(argv)

    inp = SystemInputs(a)
    outdir = a.out.resolve(); outdir.mkdir(parents=True, exist_ok=True)
    ext, dev, layer_rec = extractors(inp.fid_cfg, a.device)
    batch = int(inp.fid_cfg.get("feature_batch", 200))
    adapter, gids, t_idx = inp.adapter, inp.gids, inp.t_idx
    models = inp.models()
    floor_ns = [int(x) for x in a.floor_ns.split(",")]
    model_ns = [int(x) for x in a.model_ns.split(",")]
    nmax = max(model_ns + [2 * max(floor_ns)])
    print(f"{a.system}: {len(gids)} test conditions, features at N = {nmax}")
    print("models: " + "; ".join(f"{m['label']} ({m['key']}, n* = {m['n_steps']})" for m in models))
    print(f"floor N {floor_ns} (at most 5000: an independent floor at N needs 2N of the 10,000 members); model N {model_ns}\n")

    acc = defaultdict(lambda: defaultdict(list))        # acc[(row, N)][metric] -> per-condition values
    for g in gids:
        tv = bank_members(inp.bank, g, t_idx, nmax)
        imgs, _ = adapter.to_images(tv)
        tc, ts = both_features(imgs, ext, dev, batch)
        for N in floor_ns:
            acc[("floor (truth vs truth)", N)]["cfid"].append(BE.fid_from_features(tc[:N], tc[N:2 * N]))
            acc[("floor (truth vs truth)", N)]["sfid"].append(BE.fid_from_features(ts[:N], ts[N:2 * N]))
        line = (f"  gid {g:4d} floor({floor_ns[-1]}) c{acc[('floor (truth vs truth)', floor_ns[-1])]['cfid'][-1]:6.2f}"
                f"/s{acc[('floor (truth vs truth)', floor_ns[-1])]['sfid'][-1]:7.2f}")
        for m in models:
            gi, _ = adapter.to_images(inp.generated(m["methods"][0], g, nmax))
            gc, gs = both_features(gi, ext, dev, batch)
            for N in model_ns:
                acc[(m["label"], N)]["cfid"].append(BE.fid_from_features(gc[:N], tc[:N]))
                acc[(m["label"], N)]["sfid"].append(BE.fid_from_features(gs[:N], ts[:N]))
            line += f" | {m['label'][:20]} c{acc[(m['label'], model_ns[-1])]['cfid'][-1]:6.2f}/s{acc[(m['label'], model_ns[-1])]['sfid'][-1]:7.2f}"
        print(line, flush=True)

    def ms(v):
        x = np.asarray(v, float); return x.mean(), x.std(ddof=1) / math.sqrt(len(x))

    rows = []
    print(f"\n{'row':28s} {'N':>6s} {'cFID':>17s} {'csFID':>18s}")
    for (label, N), d in sorted(acc.items(), key=lambda kv: (kv[0][0], kv[0][1])):
        mc, sc = ms(d["cfid"]); msf, ssf = ms(d["sfid"])
        rows.append({"row": label, "N": N, "n_ic": len(d["cfid"]), "cfid": mc, "cfid_sem": sc, "sfid": msf, "sfid_sem": ssf,
                     "per_gid_cfid": {int(g): float(v) for g, v in zip(gids, d["cfid"])},
                     "per_gid_sfid": {int(g): float(v) for g, v in zip(gids, d["sfid"])}})
        print(f"{label:28s} {N:6d} {mc:8.3f} +- {sc:6.3f} {msf:9.3f} +- {ssf:6.3f}")

    print("\nordering and relative gaps across N (model rows only)")
    summary = {}
    for metric, key in (("cFID", "cfid"), ("csFID", "sfid")):
        summary[metric] = {}
        for N in model_ns:
            vals = {r["row"]: r[key] for r in rows if r["N"] == N and not r["row"].startswith("floor")}
            order = sorted(vals, key=vals.get)
            best = vals[order[0]]
            summary[metric][N] = {"order": order, "values": vals, "ratio_to_best": {k: vals[k] / best for k in order}}
            print(f"  {metric} N={N:6d}  order: {' < '.join(order)}   ratios to best: "
                  + ", ".join(f"{k}={v:.2f}" for k, v in summary[metric][N]["ratio_to_best"].items()))

    write_json(outdir / "sample_stability.json", {"at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "system": a.system,
                                                  "gids": gids, "t_idx": t_idx, "floor_ns": floor_ns, "model_ns": model_ns,
                                                  "models": models, "adapter": adapter.params, "layer": layer_rec, "pins": BE.pins(),
                                                  "rows": rows, "stability": summary})
    with open(outdir / "sample_stability.csv", "w", newline="") as fh:
        w = csv.writer(fh); w.writerow(["row", "N", "n_ic", "cfid", "cfid_sem", "sfid", "sfid_sem"])
        for r in rows:
            w.writerow([r["row"], r["N"], r["n_ic"], r["cfid"], r["cfid_sem"], r["sfid"], r["sfid_sem"]])
    print(f"\nwrote {outdir / 'sample_stability.csv'}\n      {outdir / 'sample_stability.json'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())