"""Recompute every number in the claim and report 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

Needs only Python 3.
"""
import csv
import json
import math
import sys
from pathlib import Path

HERE = Path(__file__).parent
HARTREE_KCAL = 627.5095
R_KCAL = 0.0019872043      # gas constant, kcal/(mol K)
T_K = 353.15               # 80 °C, the experiment's temperature


def rows(name):
    with open(HERE / "data" / f"{name}.csv", newline="", encoding="utf-8") as fh:
        return list(csv.DictReader(fh))


def compute():
    species = rows("species")
    adducts = [r for r in species if r["stereo"]]
    frac = lambda st: sum(float(r["final_fraction"]) for r in adducts if r["stereo"] == st)
    ts = {}
    for r in rows("cycloaddition_ts"):
        de_x = (float(r["gxtb_ts_hartree"]) - float(r["gxtb_reactants_separated_hartree"])) * HARTREE_KCAL
        de_d = (float(r["dft_ts_hartree"]) - float(r["dft_reactants_separated_hartree"])) * HARTREE_KCAL
        thermal = float(r["gxtb_dG_kcal"]) - de_x          # g-xTB free-energy correction of the barrier
        ts[int(r["adduct_species"])] = {"stereo": r["stereo"], "gxtb_dE": de_x, "gxtb_dG": float(r["gxtb_dG_kcal"]),
                                        "dft_dE": de_d, "dft_dG": de_d + thermal}
    # the mirror-image endo TS (higher of the two endo TSs) is reported as endo_mirror
    endo = sorted((k for k, t in ts.items() if t["stereo"] == "endo"), key=lambda k: ts[k]["gxtb_dG"])
    exo = [k for k, t in ts.items() if t["stereo"] == "exo"]
    names = {endo[0]: "endo", endo[1]: "endo_mirror", exo[0]: "exo"}
    out = {"n_species": len(species), "n_endo_adducts": len(endo), "n_exo_adducts": len(exo),
           "conversion_pct": round(100 * (1 - float(species[0]["final_fraction"])), 1),
           "gxtb_kinetics_endo_pct": round(100 * frac("endo") / (frac("endo") + frac("exo")))}
    for k, name in names.items():
        for q in ("gxtb_dE", "gxtb_dG", "dft_dE", "dft_dG"):
            out[f"{q}_{name}"] = round(ts[k][q], 1)
    for q in ("gxtb", "dft"):   # mirror images should be equal: the gap is a sampling shortfall
        out[f"{q}_endo_mirror_gap"] = round(ts[endo[1]][f"{q}_dG"] - ts[endo[0]][f"{q}_dG"], 1)
    # each diastereomer's lowest TS found; the experiment's temperature
    for q in ("gxtb", "dft"):
        gap = min(ts[k][f"{q}_dG"] for k in exo) - min(ts[k][f"{q}_dG"] for k in endo)
        out[f"{q}_ts_endo_pct"] = round(100 / (1 + math.exp(-gap / (R_KCAL * T_K))))
        out[f"{q}_ts_exo_pct"] = 100 - out[f"{q}_ts_endo_pct"]
    out["gxtb_kinetics_exo_pct"] = 100 - out["gxtb_kinetics_endo_pct"]
    exp = rows("experiment")[0]
    out["exp_endo_pct"] = int(exp["endo_pct"])
    out["exp_exo_pct"] = int(exp["exo_pct"])
    out["exp_temperature_c"] = int(exp["temperature_c"])
    out["exp_time_h"] = float(exp["time_h"])
    out["exp_yield_pct"] = int(exp["isolated_yield_pct"])
    return out


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__":
    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/")