#!/usr/bin/env python3
"""Audit all 45 A4 trajectories and apply the frozen Oral-A gate."""
import argparse
import glob
import json
import math
import os
import statistics


METHODS = ("bp", "dfa", "sdil")
DEPTHS = (20, 32, 56)
SEEDS = tuple(range(10, 15))


def read(path):
    with open(path) as handle:
        record = json.load(handle)
    args = record["args"]
    key = (args["mode"], int(args["depth"]), int(args["seed"]))
    if key[0] not in METHODS or key[1] not in DEPTHS or key[2] not in SEEDS:
        raise ValueError(f"unexpected A4 key {key}: {path}")
    expected = {
        "width": 16, "epochs": 200, "val_examples": 0,
        "eval_split": "test", "eval_every": 0, "normalization": "batchnorm",
    }
    for name, value in expected.items():
        if args.get(name) != value:
            raise ValueError(f"{path}: {name} drift")
    if record["provenance"]["git_tracked_dirty"]:
        raise ValueError(f"tracked-dirty result: {path}")
    protocol = record["evaluation_protocol"]
    if protocol["test_evaluations"] != 1 or protocol["test_used_for_selection"]:
        raise ValueError(f"invalid test protocol: {path}")
    accuracy = float(record["final"]["accuracy"])
    loss = float(record["final"]["loss"])
    alignment = (None if record["diagnostics"] is None
                 else float(record["diagnostics"]["early_third_mean"]))
    return key, {
        "path": path, "accuracy": accuracy, "loss": loss,
        "finite": record["final"]["finite"] and math.isfinite(accuracy + loss),
        "early_third_alignment": alignment,
        "total_macs": record["work"]["total_macs_estimate"],
        "logical_queries": record["work"]["logical_batch_loss_queries"],
        "peak_memory": record["hardware"]["peak_memory_allocated_bytes"],
        "source_commit": record["provenance"]["git_commit"],
    }


def mean_sd(values):
    return {"mean": statistics.mean(values), "sample_sd": statistics.stdev(values)}


def frontier_members(points, cost):
    members = []
    for candidate in points:
        dominated = any(
            other["accuracy"] >= candidate["accuracy"]
            and other[cost] <= candidate[cost]
            and (other["accuracy"] > candidate["accuracy"]
                 or other[cost] < candidate[cost])
            for other in points if other is not candidate)
        if not dominated:
            members.append(candidate["name"])
    return members


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="results/oral_a_confirm")
    parser.add_argument("--out", default="results/oral_a_confirmation.json")
    args = parser.parse_args()
    rows = dict(read(path) for path in sorted(glob.glob(os.path.join(args.input, "*.json"))))
    expected = {(method, depth, seed) for method in METHODS
                for depth in DEPTHS for seed in SEEDS}
    if set(rows) != expected or len(rows) != 45:
        raise ValueError(f"incomplete A4 grid: missing={expected-set(rows)}, extra={set(rows)-expected}")
    summaries = {}
    for method in METHODS:
        for depth in DEPTHS:
            group = [rows[(method, depth, seed)] for seed in SEEDS]
            summaries[f"{method}_d{depth}"] = {
                "accuracy": mean_sd([row["accuracy"] for row in group]),
                "total_macs": mean_sd([row["total_macs"] for row in group]),
                "peak_memory": mean_sd([row["peak_memory"] for row in group]),
                "logical_queries": mean_sd([row["logical_queries"] for row in group]),
                "early_third_alignment": (None if method == "bp" else mean_sd(
                    [row["early_third_alignment"] for row in group])),
            }
    near_bp = all(
        summaries[f"sdil_d{depth}"]["accuracy"]["mean"]
        >= summaries[f"bp_d{depth}"]["accuracy"]["mean"] - 0.02
        for depth in DEPTHS)
    d56_advantage = (
        summaries["sdil_d56"]["accuracy"]["mean"]
        >= summaries["dfa_d56"]["accuracy"]["mean"] + 0.02)
    paired_depth_changes = [
        rows[("sdil", 56, seed)]["accuracy"] - rows[("sdil", 20, seed)]["accuracy"]
        for seed in SEEDS]
    depth_preserved = statistics.mean(paired_depth_changes) >= -0.02
    align20 = summaries["sdil_d20"]["early_third_alignment"]["mean"]
    align56 = summaries["sdil_d56"]["early_third_alignment"]["mean"]
    alignment_gate = align56 >= 0.05 and align56 >= 0.30 * align20
    local_points = []
    for method in ("dfa", "sdil"):
        for depth in DEPTHS:
            summary = summaries[f"{method}_d{depth}"]
            local_points.append({
                "name": f"{method}_d{depth}",
                "accuracy": summary["accuracy"]["mean"],
                "total_macs": summary["total_macs"]["mean"],
                "logical_queries": summary["logical_queries"]["mean"],
                "peak_memory": summary["peak_memory"]["mean"],
            })
    frontiers = {cost: frontier_members(local_points, cost) for cost in (
        "total_macs", "logical_queries", "peak_memory")}
    nondominated = all(
        any(name.startswith("sdil_") for name in members)
        for members in frontiers.values())
    checks = {
        "all_finite": all(row["finite"] for row in rows.values()),
        "within_2pt_bp_each_depth": near_bp,
        "d56_beats_dfa_2pt": d56_advantage,
        "paired_d20_to_d56_no_worse_than_minus_2pt": depth_preserved,
        "d56_alignment_level_and_retention": alignment_gate,
        "sdil_on_each_local_cost_frontier": nondominated,
    }
    output = {
        "protocol": "oral_a_A4_v1",
        "status": "passed" if all(checks.values()) else "failed",
        "checks": checks, "summaries": summaries,
        "local_frontiers": frontiers,
        "paired_sdil_d20_to_d56_accuracy_change": mean_sd(paired_depth_changes),
        "rows": {f"{method}_d{depth}_s{seed}": row
                 for (method, depth, seed), row in rows.items()},
    }
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w") as handle:
        json.dump(output, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps({"status": output["status"], "checks": checks,
                      "summaries": summaries}, indent=2))


if __name__ == "__main__":
    main()