"""Reproduce one entry of Table 2, or one depth point of Table 6 / Fig. 3: csFID-10k of a model on the ten test conditions.

    python -m reproduce.table2 list
    python -m reproduce.table2 check    --system SYS --method M [--S S] [--checkpoint FILE] [--samples DIR] [--hash-large]
    python -m reproduce.table2 run      --system SYS --method M [--S S] --checkpoint FILE --out DIR [--device cuda]
    python -m reproduce.table2 generate --system SYS --method M [--S S] --checkpoint FILE --out DIR [--gids g1,g2] [--n-samples N]
    python -m reproduce.table2 score    --system SYS --method M [--S S] --out DIR [--samples DIR]
    python -m reproduce.table2 compare  --system SYS --method M [--S S] --out DIR [--tolerance REL]

SYS is euler | sqg | waves | nlkg | passive_scalar; M is direct_si | edm | cascaded_si (without --S the selected depth S*,
with --S the Cascaded SI model of that depth). The package and the truth bank are taken from <SYS>_DATA and
<SYS>_TRUTH_BANK or given with --data and --truth-bank. Everything else comes from records/<SYS>/: the run configuration
of the model, its sampler resolution n*, the image adapter and the recorded values.

  generate  10,000 samples per test condition (bifurcation.sample with the settings of configs/evaluation.yaml:
            n* sampler steps, noise stream 777 * 100003 + gid, chunks of 2,000) -> OUT/samples/<variant>/group_<gid>.npz
  score     csFID per condition (bifurcation.fid.evaluate.evaluate_csfid_run) of OUT/samples, or of the archived samples
            given with --samples -> OUT/csfid/{csfid_per_ic.csv, csfid_summary.json, MANIFEST_CSFID.json}
  compare   per-condition values, mean and SEM against the recorded values -> OUT/COMPARE.json
  run       generate + score + compare
  check     the package, truth bank, checkpoint and archived samples against the recorded sha256
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import json
import sys
from datetime import datetime, timezone
from pathlib import Path

from bifurcation.config import RECORDS, REPO, SYSTEMS, data_paths, dataset_card, evaluation_settings, records, run_config, unit


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 24), b""):
            h.update(chunk)
    return h.hexdigest()


def asset_hashes(system: str) -> dict:
    """{section: {relative path: sha256}} from records/<system>/assets.sha256."""
    out, sec = {}, None
    for line in (RECORDS / system / "assets.sha256").read_text().splitlines():
        if line.startswith("# "):
            sec = line[2:].split(":")[0].strip()
            continue
        h, q = line.split("  ", 1)
        out.setdefault(sec, {})[q] = h
    return out


def cmd_list(a):
    print(f"{'system':15s} {'model':16s} {'paper':22s} {'n*':>4s} {'NFE':>6s}  csFID-10k")
    for s in SYSTEMS:
        rec = records(s)
        for u in rec["units"]:
            names = ", ".join({"direct_si": "Direct SI", "edm": "EDM", "cascaded_si": f"Cascaded SI (S*={rec['s_star']})"}[n]
                              for n in u["paper_names"]) or f"depth S={u['S']}"
            e = u["expected_csfid"]
            print(f"{s:15s} {u['key']:16s} {names:22s} {u['sampler']['n_steps']:4d} {u['sampler']['nfe']:6d}  {e['mean']:.2f} +- {e['sem']:.2f}")


def cmd_check(a) -> int:
    rec, u = unit(a.system, a.method, a.S)
    data, bank = data_paths(a.system, a.data, a.truth_bank)
    card = dataset_card(a.system)
    problems, checked = [], 0
    big = 200 * 1024 * 1024
    for rel, want in card["expected_sha256"].items():
        p = data / rel
        if not p.is_file():
            problems.append(f"--data: missing {rel}"); continue
        if p.stat().st_size > big and not a.hash_large:
            continue
        checked += 1
        if sha256(p) != want:
            problems.append(f"--data: {rel} sha256 differs from configs/{a.system}/dataset.yaml")
    hashes = asset_hashes(a.system)
    for gid in rec["test_gids"]:
        rel = f"group_{gid:04d}.npz"
        p = bank / rel
        if not p.is_file():
            problems.append(f"--truth-bank: missing {rel}"); continue
        if a.hash_large:
            checked += 1
            if sha256(p) != hashes["truth_bank"][rel]:
                problems.append(f"--truth-bank: {rel} sha256 differs")
    if a.checkpoint:
        got = sha256(Path(a.checkpoint)); checked += 1
        if u["checkpoint"] is None:
            problems.append(f"--checkpoint: no checkpoint of {u['key']} is distributed")
        elif got != u["checkpoint"]["sha256"]:
            problems.append(f"--checkpoint: sha256 {got[:16]} is not the recorded {u['checkpoint']['file']}")
    if a.samples:
        v = u["sampler"]["variant"]
        sec = f"samples/{u['paper_names'][0]}" if u["paper_names"] else None   # archived samples exist for the Table 2 models
        for gid in rec["test_gids"]:
            rel = f"{v}/group_{gid:04d}.npz"
            p = Path(a.samples) / rel
            if not p.is_file():
                problems.append(f"--samples: missing {rel}"); continue
            want = (hashes.get(sec) or {}).get(rel) if sec else None
            if want and a.hash_large:
                checked += 1
                if sha256(p) != want:
                    problems.append(f"--samples: {rel} sha256 differs from the archived samples")
    print(json.dumps({"system": a.system, "model": u["key"], "files_hashed": checked, "problems": problems}, indent=1))
    return 2 if problems else 0


def _fresh(out: Path, sub: str) -> Path:
    d = out / sub
    if d.exists() and any(d.iterdir()):
        raise SystemExit(f"refusing: {d} exists and is not empty (use a fresh --out)")
    d.mkdir(parents=True, exist_ok=True)
    return d


def cmd_generate(a):
    from bifurcation import sample
    import yaml
    rec, u = unit(a.system, a.method, a.S)
    data, bank = data_paths(a.system, a.data, a.truth_bank)
    ev = evaluation_settings()
    out = Path(a.out).resolve()
    out.mkdir(parents=True, exist_ok=True)
    cfg_path = out / "config.yaml"
    cfg_path.write_text(yaml.safe_dump(run_config(a.system, u, data), sort_keys=False))
    n_samples = int(a.n_samples) if getattr(a, "n_samples", None) else int(ev["fid"]["n_generated"])
    argv = ["--config", str(cfg_path), "--data", str(data), "--checkpoint", str(Path(a.checkpoint).resolve()),
            "--out", str(out / "samples"), "--split", "test", "--n-samples", str(n_samples),
            "--n-steps", str(u["sampler"]["n_steps"]), "--seed-base", str(ev["sampling"]["crn_base_seed"]),
            "--chunk", str(ev["sampling"]["sample_chunk"]), "--device", a.device]
    if rec["x0_from_truth_bank"]:
        argv += ["--truth-bank", str(bank)]
    if a.gids:
        argv += ["--gids", a.gids]
    print("$ python -m bifurcation.sample " + " ".join(argv), flush=True)
    sample.main(argv)
    log = out / "GENERATED.json"
    entries = json.loads(log.read_text()) if log.is_file() else []
    entries.append({"system": a.system, "model": u["key"], "checkpoint": str(Path(a.checkpoint).resolve()), "gids": a.gids or "all",
                    "argv": argv, "at": datetime.now(timezone.utc).isoformat(timespec="seconds")})
    log.write_text(json.dumps(entries, indent=1) + "\n")


def cmd_score(a):
    from bifurcation.fid.adapters import FIDAdapter
    from bifurcation.fid.evaluate import evaluate_csfid_run, generated_states
    rec, u = unit(a.system, a.method, a.S)
    data, bank = data_paths(a.system, a.data, a.truth_bank)
    card, fid_cfg = dataset_card(a.system), evaluation_settings()["fid"]
    out = Path(a.out).resolve()
    samples = Path(a.samples).resolve() if a.samples else out / "samples"
    variant = u["sampler"]["variant"]
    missing = [g for g in rec["test_gids"] if not (samples / variant / f"group_{g:04d}.npz").is_file()]
    if missing:
        raise SystemExit(f"{samples / variant}: no samples for test conditions {missing}")
    n_gen = int(fid_cfg["n_generated"])
    gen_raw = generated_states(samples, variant, data, n_gen)
    gen_raw = {g: x for g, x in gen_raw.items() if g in set(rec["test_gids"])}
    for g, x in gen_raw.items():
        if len(x) != n_gen:
            raise SystemExit(f"condition {g}: {len(x)} generated states, the paper's read uses {n_gen}")
    adapter = FIDAdapter.load(RECORDS / a.system / "fid_adapter.json", **(card["fid_adapter"].get("params") or {}))
    csdir = _fresh(out, "csfid")
    summary = evaluate_csfid_run(gen_raw, bank, int(rec["fine_horizon_idx"]), int(fid_cfg["n_truth"]), adapter, csdir,
                                 out / "fid_cache", fid_cfg, a.device,
                                 provenance={"system": a.system, "model": u["key"], "samples": str(samples), "source": "reproduce.table2 score"})
    print(f"csFID-10k {a.system} {u['key']}: {summary['csFID']['mean']:.6f} +- {summary['csFID']['sem']:.6f} (n = {summary['csFID']['n']})")


def cmd_compare(a) -> int:
    rec, u = unit(a.system, a.method, a.S)
    out = Path(a.out).resolve()
    got = {int(r["gid"]): float(r["sfid"]) for r in csv.DictReader(open(out / "csfid" / "csfid_per_ic.csv"))}
    summ = json.loads((out / "csfid" / "csfid_summary.json").read_text())["csFID"]
    exp = u["expected_csfid"]
    rows, worst = [], 0.0
    for g in rec["test_gids"]:
        e, v = exp["per_condition"][str(g)], got.get(g, float("nan"))
        rel = abs(v - e) / abs(e)
        worst = max(worst, rel)
        rows.append({"gid": g, "recorded": e, "this_run": v, "rel_diff": rel})
    agg = {k: {"recorded": exp[k], "this_run": summ[k], "rel_diff": abs(summ[k] - exp[k]) / abs(exp[k])} for k in ("mean", "sem")}
    for r in rows:
        print(f"  gid {r['gid']:4d}  recorded {r['recorded']:.6f}  this run {r['this_run']:.6f}  rel diff {r['rel_diff']:.2e}")
    print(f"  mean {summ['mean']:.6f} (recorded {exp['mean']:.6f}, rel diff {agg['mean']['rel_diff']:.2e}); "
          f"SEM {summ['sem']:.6f} (recorded {exp['sem']:.6f}, rel diff {agg['sem']['rel_diff']:.2e})")
    verdict = {"system": a.system, "model": u["key"], "max_rel_diff_per_condition": worst, "aggregate": agg, "per_condition": rows,
               "bitwise_equal": all(r["rel_diff"] == 0.0 for r in rows) and all(v["rel_diff"] == 0.0 for v in agg.values())}
    if a.tolerance is not None:
        verdict["tolerance"] = a.tolerance
        verdict["within_tolerance"] = worst <= a.tolerance and all(v["rel_diff"] <= a.tolerance for v in agg.values())
    (out / "COMPARE.json").write_text(json.dumps(verdict, indent=1) + "\n")
    print(f"  max relative difference per condition: {worst:.3e}" + (f"; within {a.tolerance:g}: {verdict['within_tolerance']}" if a.tolerance is not None else ""))
    return 1 if (a.tolerance is not None and not verdict["within_tolerance"]) else 0


def cmd_run(a) -> int:
    cmd_generate(a)
    a.samples = None
    cmd_score(a)
    return cmd_compare(a)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(prog="python -m reproduce.table2", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("list")

    def unit_args(p, need_data=True):
        p.add_argument("--system", required=True, choices=SYSTEMS)
        p.add_argument("--method", required=True, choices=("direct_si", "edm", "cascaded_si"))
        p.add_argument("--S", type=int, default=None, help="depth of a Cascaded SI model (Table 6); default: the selected S*")
        if need_data:
            p.add_argument("--data", default=None, help="dataset package directory (default: $<SYSTEM>_DATA)")
            p.add_argument("--truth-bank", default=None, help="truth-bank directory (default: $<SYSTEM>_TRUTH_BANK)")
    c = sub.add_parser("check"); unit_args(c)
    c.add_argument("--checkpoint"); c.add_argument("--samples")
    c.add_argument("--hash-large", action="store_true", help="also hash the large files (trajectories, truth banks, samples)")
    for name in ("generate", "run"):
        g = sub.add_parser(name); unit_args(g)
        g.add_argument("--checkpoint", required=True); g.add_argument("--out", required=True)
        g.add_argument("--gids", default=None, help="comma-separated subset of the test conditions (split generation across jobs)")
        g.add_argument("--device", default="cuda")
        if name == "generate":
            g.add_argument("--n-samples", type=int, default=None, help="fewer samples per condition (a smoke test; `score` refuses fewer than 10,000)")
        if name == "run":
            g.add_argument("--tolerance", type=float, default=None)
    s = sub.add_parser("score"); unit_args(s)
    s.add_argument("--out", required=True); s.add_argument("--samples", default=None, help="archived samples (default OUT/samples)")
    s.add_argument("--device", default="cuda")
    m = sub.add_parser("compare"); unit_args(m, need_data=False)
    m.add_argument("--out", required=True); m.add_argument("--tolerance", type=float, default=None)
    a = ap.parse_args(argv)
    fn = {"list": cmd_list, "check": cmd_check, "generate": cmd_generate, "score": cmd_score, "compare": cmd_compare, "run": cmd_run}[a.cmd]
    return int(fn(a) or 0)


if __name__ == "__main__":
    sys.exit(main())