"""Recompute every number of F0003 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
"""
import csv
import json
import statistics as st
import sys
from pathlib import Path

HERE = Path(__file__).parent


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


def compute():
    rows = list(csv.DictReader(open(HERE / "data" / "search_outcomes.csv", newline="")))
    sset = list(csv.DictReader(open(HERE / "data" / "benchmark_set.csv", newline="")))
    done = [r for r in rows if r["status"] == "done"]
    lower = [r for r in done if r["outcome"] == "lower"]
    kind = lambda k: [float(r["lower_by_kcal"]) for r in lower if r["best_kind"] == k]
    gap = [float(r["lower_by_kcal"]) for r in lower]
    med = lambda xs: round(st.median(xs), 2) if xs else None
    atoms = [int(r["n_atoms"]) for r in sset]
    left = list(csv.DictReader(open(HERE / "data" / "left_out_irc.csv", newline="")))
    nr = [r for r in csv.DictReader(open(HERE / "data" / "search_diagnostics.csv", newline="")) if r["status"] == "no_route"]
    q = lambda r: r["irc_pair_queued"] == "1"
    dn = lambda r: r["irc_pair_path_done"] == "1"
    sp = lambda r: r["irc_pair_path_split"] == "1"
    return {
        "set": len(sset), "searched": len(rows), "with_route": len(done),
        "lower": len(lower), "lower_pct_of_searched": pct(len(lower), len(rows)),
        "lower_single_step": len(kind("direct")), "lower_multistep": len(kind("multistep")),
        "matched": sum(r["outcome"] == "matched" for r in done), "higher": sum(r["outcome"] == "higher" for r in done),
        "no_route": sum(r["status"] == "no_route" for r in rows),
        "no_route_pct_of_searched": pct(sum(r["status"] == "no_route" for r in rows), len(rows)),
        "time_limit": sum(r["status"] == "timeout" for r in rows),
        "crashed": sum(r["status"] not in ("done", "no_route", "timeout") for r in rows),
        "lower_median_kcal": med(gap), "lower_median_single_step_kcal": med(kind("direct")),
        "lower_median_multistep_kcal": med(kind("multistep")), "lower_max_kcal": round(max(gap), 2) if gap else None,
        "top_n": 10, "n_atoms_min": min(atoms), "n_atoms_max": max(atoms),
        "left_out_irc": len(left),
        "no_route_irc_pair_not_queued": sum(not q(r) for r in nr),
        "no_route_irc_ends_not_in_pools": sum(r["irc_ends_in_pools"] != "1" for r in nr),
        "no_route_irc_pair_not_reached": sum(q(r) and not dn(r) for r in nr),
        "no_route_irc_pair_split": sum(dn(r) and sp(r) for r in nr),
        "no_route_irc_pair_single_step": sum(dn(r) and not sp(r) for r in nr),
        "no_route_budget_used_up": sum(r["budget_reached"] == "1" for r in nr),
        "no_route_found_other_ts": sum(r["n_offtarget"] not in ("", "0") for r in nr),
        "left_out_drop_over_20_kcal": sum(max(float(r["start_end_drop_kcal"]), float(r["end_end_drop_kcal"])) > 20
                                          for r in left),
        "left_out_soft_saddle": sum(float(r["ts_imaginary_cm"]) > -200 for r in left),
        "left_out_clear_saddle_small_drop": sum(float(r["ts_imaginary_cm"]) <= -200 and
                                                max(float(r["start_end_drop_kcal"]), float(r["end_end_drop_kcal"])) <= 5
                                                for r in left),
    }


def expected():
    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__":
    got = compute()
    if "--print" in sys.argv:
        print(json.dumps(got, indent=1))
        sys.exit(0)
    want = expected()
    bad = [f"{k}: recorded {want.get(k)!r}, data gives {v!r}" for k, v in got.items() if want.get(k) != v]
    bad += [f"{k}: recorded but not computed" for k in want if k not in got]
    if bad:
        print("\n".join(bad))
        sys.exit(1)
    print(f"all {len(got)} numbers reproduce from data/")