"""Audit the validation-only causal-calibration query screen."""
import glob
import json
import math
import os


ROOT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results")


def main():
    paths = sorted(glob.glob(os.path.join(ROOT, "query_dev_v1_*.json")))
    if not paths:
        raise RuntimeError("no query_dev_v1 results")
    rows = []
    commits = set()
    split_hashes = set()
    for path in paths:
        with open(path) as handle:
            row = json.load(handle)
        if row.get("final", {}).get("eval_split") != "validation":
            raise RuntimeError(f"non-validation result: {path}")
        if row.get("provenance", {}).get("git_dirty") is not False:
            raise RuntimeError(f"dirty or unknown provenance: {path}")
        commits.add(row["provenance"]["git_commit"])
        split_hashes.add(row["split"]["validation_index_sha256"])
        rows.append(row)
    if len(commits) != 1 or len(split_hashes) != 1:
        raise RuntimeError(f"mixed provenance/splits: commits={commits}, splits={split_hashes}")

    dfa = [row for row in rows if row["args"]["mode"] == "dfa"]
    ref = [row for row in rows if row["args"]["mode"] == "sdil"
           and row["args"]["pert_ndirs"] == 16 and row["args"]["pert_every"] == 4]
    if len(dfa) != 1 or len(ref) != 1:
        raise RuntimeError(f"need exactly one DFA and k16/e4 reference; got {len(dfa)}, {len(ref)}")
    dfa_acc = dfa[0]["final"]["val_acc"]
    ref_acc = ref[0]["final"]["val_acc"]
    ref_gain = ref_acc - dfa_acc
    ref_queries = ref[0]["cost"]["calibration_batch_loss_evaluations"]
    ref_work = ref[0]["cost"]["calibration_forward_equivalent_examples"]

    print(f"commit={next(iter(commits))} split={next(iter(split_hashes))}")
    print("| method | K | every | val (%) | batch loss queries | query reduction | "
          "forward-eq reduction | gain retention | wall (s) |")
    print("|:---|---:|---:|---:|---:|---:|---:|---:|---:|")
    for row in sorted(rows, key=lambda r: (r["args"]["mode"] == "dfa",
                                           r["args"]["pert_ndirs"],
                                           r["args"]["pert_every"])):
        args = row["args"]
        cost = row["cost"]
        queries = cost["calibration_batch_loss_evaluations"]
        work = cost["calibration_forward_equivalent_examples"]
        q_reduction = ref_queries / queries if queries else math.inf
        w_reduction = ref_work / work if work else math.inf
        retention = ((row["final"]["val_acc"] - dfa_acc) / ref_gain
                     if args["mode"] == "sdil" and ref_gain else float("nan"))
        print(f"| {args['mode']} | {args['pert_ndirs']} | {args['pert_every']} | "
              f"{100 * row['final']['val_acc']:.3f} | {queries} | {q_reduction:.1f}x | "
              f"{w_reduction:.1f}x | {retention:.3f} | {row['final']['wall_s']:.1f} |")


if __name__ == "__main__":
    main()