"""Compares a reproduction working copy with the shipped artifacts and writes a summary.

usage: python3 code/tools/compare_reproduction.py WORK ROOT MODE OUT.json

Checks (none of them modifies the shipped artifacts):
  1. exact values: every (n, Z) in WORK/results/exact_values.csv equals the shipped one;
  2. regenerated boundary certificates (records in WORK whose paths are outside
     code/proofs/certificates): the regenerated CNF has the same SHA-256 as the shipped
     CNF for the same (n, p) (the formula is deterministic), the solver reported UNSAT
     and drat-trim verified the regenerated proof;
  3. regenerated fort-cover certificates: the solver reported UNSAT at the same bound as
     the shipped record and the regenerated proof was verified; the fort list produced
     by the lazy loop may legitimately differ from the shipped one, so CNF digests are
     reported but not required to agree;
  4. every certificate the independent checker actually checked in WORK was accepted.
Exit code 0 iff checks 1, 2 (UNSAT + verified), 3 (UNSAT + verified) and 4 pass.
"""
import csv
import datetime
import json
import os
import sys


def jl(path):
    return [json.loads(l) for l in open(path)] if os.path.exists(path) else []


def main(work, root, mode, out_path):
    rep = {"utc": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"), "mode": mode, "work": work}
    ok = True
    # 1 exact values
    def vals(p):
        return {int(r["n"]): (r["Z"], r["status"]) for r in csv.DictReader(open(p))} if os.path.exists(p) else {}
    shipped, regen = vals(os.path.join(root, "results", "exact_values.csv")), vals(os.path.join(work, "results", "exact_values.csv"))
    diff = {n: {"shipped": shipped.get(n), "reproduced": regen.get(n)} for n in sorted(set(shipped) | set(regen)) if shipped.get(n, (None,))[0] != regen.get(n, (None,))[0]}
    rep["exact_values"] = {"rows_shipped": len(shipped), "rows_reproduced": len(regen), "value_differences": diff}
    ok &= not diff and bool(regen)
    # 2 regenerated boundary certificates
    shipped_b = {(r["n"], r["p"]): r for r in jl(os.path.join(root, "results", "boundary_certificates.jsonl")) if r.get("sat") is False}
    rows = []
    for r in jl(os.path.join(work, "results", "boundary_certificates.jsonl")):
        if r.get("sat") is not False or r.get("cnf_path", "").startswith("code/proofs/certificates"):
            continue  # shipped record copied into WORK, or SAT/timeout
        s = shipped_b.get((r["n"], r["p"]))
        row = {"n": r["n"], "p": r["p"], "unsat": True, "verified": bool(r.get("checker_verified")) and r.get("checker_exit") == 0,
               "cnf_sha_matches_shipped": (s is not None and s["cnf_sha256"] == r["cnf_sha256"]),
               "proof_sha_matches_shipped": (s is not None and s.get("proof_sha256") == r.get("proof_sha256"))}
        rows.append(row)
        ok &= row["verified"] and row["cnf_sha_matches_shipped"]
    rep["regenerated_boundary"] = {"count": len(rows), "all_verified": all(x["verified"] for x in rows), "all_cnf_match": all(x["cnf_sha_matches_shipped"] for x in rows),
                                   "proofs_identical_to_shipped": sum(1 for x in rows if x["proof_sha_matches_shipped"]), "rows": rows}
    # 3 regenerated fort-cover certificates
    shipped_f = {(r["n"], r["bound"]): r for r in jl(os.path.join(root, "results", "lower_bound_certificates.jsonl"))}
    rows = []
    for r in jl(os.path.join(work, "results", "lower_bound_certificates.jsonl")):
        if r.get("cnf_path", "").startswith("code/proofs/certificates"):
            continue
        s = shipped_f.get((r["n"], r["bound"]))
        row = {"n": r["n"], "bound": r["bound"], "shipped_record_exists_at_same_bound": s is not None,
               "verified": bool(r.get("checker_verified")) and r.get("checker_exit") == 0,
               "cnf_sha_matches_shipped": (s is not None and s["cnf_sha256"] == r["cnf_sha256"]), "num_forts": r.get("num_forts")}
        rows.append(row)
        ok &= row["verified"] and row["shipped_record_exists_at_same_bound"]
    rep["regenerated_fortcover"] = {"count": len(rows), "all_verified": all(x["verified"] for x in rows), "rows": rows}
    # 4 independent checker verdicts in WORK
    checker = {}
    for name in ("lower_bound_verification.json", "lower_bound_verification_regenerated.json"):
        p = os.path.join(work, "results", name)
        if os.path.exists(p):
            v = json.load(open(p))
            checker[name] = {"counts": v.get("counts"), "all_checked_ok": v.get("all_checked_ok", v.get("all_ok"))}
            ok &= bool(checker[name]["all_checked_ok"])
    p = os.path.join(work, "results", "upper_bound_verification.json")
    if os.path.exists(p):
        v = json.load(open(p))
        checker["upper_bound_verification.json"] = {"witnesses_ok": sum(1 for x in v["witnesses"] if x["replay_ok"]), "witnesses": len(v["witnesses"]),
                                                    "P7_all_ok": all(x["replay_ok"] for x in v["P7"])}
        ok &= checker["upper_bound_verification.json"]["witnesses_ok"] == checker["upper_bound_verification.json"]["witnesses"]
    rep["checker"] = checker
    rep["ok"] = bool(ok)
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    json.dump(rep, open(out_path, "w"), indent=1)
    print(json.dumps({k: v for k, v in rep.items() if k not in ("regenerated_boundary", "regenerated_fortcover")}, indent=1))
    print("regenerated boundary:", {k: v for k, v in rep["regenerated_boundary"].items() if k != "rows"})
    print("regenerated fort-cover:", {k: v for k, v in rep["regenerated_fortcover"].items() if k != "rows"})
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main(*sys.argv[1:5])