"""Recompute every headline number of F0001 from data/ and assert it matches the recorded values
(finding.yaml in the findings repo, expected_numbers.json in a downloaded copy).

    python check.py           # assert
    python check.py --print   # show computed values (for authoring)
    python check.py --strata  # survival by element set, size, fragment count
"""

import collections
import csv
import io
import json
import statistics as st
import sys
import zipfile
from pathlib import Path


HERE = Path(__file__).parent


def load(name):
    """data/<name>.csv.zip -> rows (a zip holding <name>.csv)."""
    with zipfile.ZipFile(HERE / "data" / f"{name}.csv.zip") as z, \
            io.TextIOWrapper(z.open(f"{name}.csv"), encoding="utf-8", newline="") as fh:
        return list(csv.DictReader(fh))


def pct(k, n):
    return round(100 * k / n, 1)


def compute():
    rows, pop = load("screen"), load("population")
    n = len(rows)
    status = collections.Counter(r["status"] for r in rows)
    match = [r for r in rows if r["status"] == "match"]
    err = [float(r["gxtb_barrier_sep_kcal"]) - float(r["ref_barrier_kcal"]) for r in match]
    uni = [r for r in rows if r["n_reactant_frags"] == "1"]
    multi = [r for r in rows if r["n_reactant_frags"] != "1"]
    unbalanced = [r for r in rows if r["balanced"] == "0"]
    sc = [p for p in pop if p["screened"] == "1"]
    un = [p for p in pop if p["screened"] == "0"]
    return {
        "n_population": len(pop),
        "n_screened": n,
        "n_match": status["match"],
        "n_nonmatch": n - status["match"],
        "match_pct": pct(status["match"], n),
        "mismatch_pct": pct(status["mismatch"], n),
        "degenerate_irc_pct": pct(status["degenerate_irc"], n),
        "irc_failed_pct": pct(status["irc_failed"], n),
        "ts_failed_pct": pct(status["ts_failed"], n),
        "match_pct_unimolecular": pct(sum(r["status"] == "match" for r in uni), len(uni)),
        "match_pct_multimolecular": pct(sum(r["status"] == "match" for r in multi), len(multi)),
        "n_matched_with_sep_barrier": len(err),
        "barrier_mae_kcal": round(st.mean(abs(e) for e in err), 2),
        "barrier_mean_signed_kcal": round(st.mean(err), 2),
        "barrier_median_signed_kcal": round(st.median(err), 2),
        "within_5_kcal_pct": pct(sum(abs(e) <= 5 for e in err), len(err)),
        "within_10_kcal_pct": pct(sum(abs(e) <= 10 for e in err), len(err)),
        "within_20_kcal_pct": pct(sum(abs(e) <= 20 for e in err), len(err)),
        "n_unbalanced_screened": len(unbalanced),
        "n_unbalanced_matched": sum(r["status"] == "match" for r in unbalanced),
        "median_ref_barrier_screened_kcal": round(st.median(float(p["ref_dE_dagger_kcal"]) for p in sc), 2),
        "median_ref_barrier_unscreened_kcal": round(st.median(float(p["ref_dE_dagger_kcal"]) for p in un), 2),
        "mean_atoms_screened": round(st.mean(int(p["n_atoms"]) for p in sc), 2),
        "mean_atoms_unscreened": round(st.mean(int(p["n_atoms"]) for p in un), 2),
    }


def strata():
    rows = load("screen")
    groups = {
        "reactant fragments": lambda r: r["n_reactant_frags"],
        "atoms": lambda r: f"{(int(r['n_atoms']) - 1) // 5 * 5 + 1:>2}-{(int(r['n_atoms']) - 1) // 5 * 5 + 5}",
        "element set (top 12)": lambda r: r["elements"],
    }
    for title, key in groups.items():
        c = collections.defaultdict(lambda: [0, 0])
        for r in rows:
            c[key(r)][0] += 1
            c[key(r)][1] += r["status"] == "match"
        items = sorted(c.items(), key=lambda kv: -kv[1][0])[:12] if "top" in title else sorted(c.items())
        print(f"\n{title}")
        for k, (n, m) in items:
            print(f"  {k:>10}  n={n:>6}  match {pct(m, n):5.1f}%")


def recorded_numbers() -> dict:
    if (HERE / "finding.yaml").exists():
        import yaml
        return yaml.safe_load((HERE / "finding.yaml").read_text())["evidence"]["numbers"]
    return json.loads((HERE / "expected_numbers.json").read_text())


if __name__ == "__main__":
    if "--strata" in sys.argv:
        strata()
        sys.exit()
    computed = compute()
    if "--print" in sys.argv:
        print(json.dumps(computed, indent=1))
        sys.exit()
    numbers = recorded_numbers()
    bad = [f"{k}: computed {computed.get(k)!r}, recorded {v!r}" for k, v in numbers.items() if computed.get(k) != v]
    if bad:
        sys.exit("mismatch:\n  " + "\n  ".join(bad))
    print(f"ok: {len(numbers)} numbers recomputed from data/")