"""Recompute every headline number of F0004 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)
"""

import csv
import gzip
import hashlib
import json
import statistics as st
import sys
from pathlib import Path

HERE = Path(__file__).parent


def path_blocks():
    lines = gzip.open(HERE / "data" / "paths.xyz.gz", "rt").read().splitlines()
    blocks, i = {}, 0
    while i < len(lines):
        n = int(lines[i]); head = lines[i + 1].split()
        blocks.setdefault((head[0], head[1].split("=")[1]), []).extend(lines[i:i + 2 + n])
        i += 2 + n
    return {k: hashlib.sha256("\n".join(v).encode()).hexdigest() for k, v in blocks.items()}


def compute():
    rows = list(csv.DictReader(open(HERE / "data" / "timings.csv", newline="")))
    calls = sorted({(r["idx"], r["nimages"]) for r in rows})
    stored = path_blocks()
    # median over the three rounds, per call and version
    t = {(c, v): st.median(float(r["seconds"]) for r in rows if (r["idx"], r["nimages"]) == c and r["version"] == v)
         for c in calls for v in ("before", "after")}
    hashes = {c: {r["path_sha256"] for r in rows if (r["idx"], r["nimages"]) == c} for c in calls}
    ratio = [t[c, "before"] / t[c, "after"] for c in calls]
    total = {v: {rnd: sum(float(r["seconds"]) for r in rows if r["version"] == v and r["round"] == rnd)
                 for rnd in ("1", "2", "3")} for v in ("before", "after")}
    tb, ta = st.median(total["before"].values()), st.median(total["after"].values())
    out = {
        "n_reactions": len({c[0] for c in calls}),
        "n_calls": len(calls),
        "n_rounds": len({r["round"] for r in rows}),
        "calls_bit_identical": sum(len(h) == 1 and h == {stored[c]} for c, h in hashes.items()),
        "total_before_s": round(tb, 1),
        "total_after_s": round(ta, 1),
        "total_speedup": round(tb / ta, 1),
        "median_call_before_s": round(st.median(t[c, "before"] for c in calls), 2),
        "median_call_after_s": round(st.median(t[c, "after"] for c in calls), 2),
        "median_call_speedup": round(st.median(ratio), 1),
        "min_call_speedup": round(min(ratio), 1),
        "max_call_speedup": round(max(ratio), 1),
        "max_round_spread_pct": round(100 * max((max(v.values()) - min(v.values())) / min(v.values())
                                                for v in total.values()), 1),
    }
    natoms = {r["idx"]: int(r["n_atoms"]) for r in csv.DictReader(open(HERE / "data" / "reactions.csv", newline=""))}
    big = max(natoms, key=natoms.get)
    out["n_atoms_min"], out["n_atoms_max"] = min(natoms.values()), natoms[big]
    out["n_reactions_over_26_atoms"] = sum(n > 26 for n in natoms.values())
    for n in ("12", "30"):
        out[f"largest_speedup_{n}_images"] = round(t[(big, n), "before"] / t[(big, n), "after"], 1)
    out["median_call_speedup_up_to_15_atoms"] = round(st.median(t[c, "before"] / t[c, "after"] for c in calls
                                                                if natoms[c[0]] <= 15), 1)
    out["median_call_speedup_over_15_atoms"] = round(st.median(t[c, "before"] / t[c, "after"] for c in calls
                                                               if natoms[c[0]] > 15), 1)
    out["n_reactions_up_to_15_atoms"] = sum(n <= 15 for n in natoms.values())
    for n in ("12", "30"):
        cs = [c for c in calls if c[1] == n]
        out[f"speedup_{n}_images"] = round(sum(t[c, "before"] for c in cs) / sum(t[c, "after"] for c in cs), 1)
    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]
    if bad:
        sys.exit("mismatch:\n  " + "\n  ".join(bad))
    print(f"ok: {len(numbers)} numbers recomputed from data/")