"""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
    python check.py --strata  # how many are kept, by element set, size and number of reactant molecules

Needs only Python 3.
"""
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)
    kept = [r for r in rows if r["kept"] == "1"]
    left = [r for r in rows if r["status"] == "match" and r["kept"] == "0"]
    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"]
    soft = [r for r in kept if float(r["ts_imaginary_cm"]) > -200]
    return {
        "n_population": len(pop),
        "n_screened": n,
        "n_kept": len(kept),
        "n_not_kept": n - len(kept),
        "kept_pct": pct(len(kept), n),
        "n_left_out_irc": len(left),
        "left_out_irc_pct": pct(len(left), n),
        "one_end_pct": pct(status["start_only"] + status["end_only"], n),
        "mismatch_pct": pct(status["mismatch"], n),
        "stereo_undetermined_pct": pct(status["stereo_undetermined"], n),
        "degenerate_irc_pct": pct(status["degenerate_irc"], n),
        "irc_problem_pct": pct(status["irc_failed"] + status["irc_end_not_minimized"], n),
        "ts_wrong_order_pct": pct(status["ts_wrong_order"], n),
        "ts_not_converged_pct": pct(status["ts_not_converged"], n),
        "n_error": status["error"],
        "kept_pct_unimolecular": pct(sum(r["kept"] == "1" for r in uni), len(uni)),
        "kept_pct_multimolecular": pct(sum(r["kept"] == "1" for r in multi), len(multi)),
        "n_unbalanced_screened": len(unbalanced),
        "n_unbalanced_kept": sum(r["kept"] == "1" for r in unbalanced),
        "n_kept_soft_saddle": len(soft),
        "n_irc_retried": sum(r["irc_retried"] != "0" for r in rows),
        "n_irc_retry_failed": sum(r["irc_retried"] == "failed" for r in rows),
        "n_irc_retry_kept": sum(r["irc_retried"] == "1" and r["kept"] == "1" for r in rows),
        "n_irc_retry_kept_was_match": sum(r["irc_retried"] == "1" and r["kept"] == "1" and r["loose_irc_status"] == "match" for r in rows),
        "n_irc_retry_status_changed": sum(r["irc_retried"] == "1" and r["loose_irc_status"] != r["status"] for r in rows),
        "irc_retry_core_hours": round(sum(float(r["irc_retry_seconds"]) for r in rows if r["irc_retry_seconds"]) / 3600),
        "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),
        "core_hours": round(sum(float(r["seconds"]) for r in rows if r["seconds"]) / 3600),
        "median_seconds": round(st.median(float(r["seconds"]) for r in rows if r["seconds"])),
    }


def strata():
    rows = load("screen")
    groups = {
        "reactant molecules": 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["kept"] == "1"
        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}  kept {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]
    bad += [f"{k}: computed but not recorded" for k in computed if k not in numbers]
    if bad:
        sys.exit("mismatch:\n  " + "\n  ".join(bad))
    print(f"ok: {len(numbers)} numbers recomputed from data/")