"""Audit the frozen BP useful-depth validation curve for C2."""
import glob
import json
import os
import statistics


ROOT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "results")
PREFIX = "c2_bp_curve_val_v1_"
DEPTHS = (1, 2, 3, 4, 6)
TASK_SEEDS = (0, 1, 2)


def audit_row(path, row):
    args = row["args"]
    required = {
        "mode": "bp", "dataset": "tentmap", "width": 8, "act": "relu",
        "residual": 1, "epochs": 80, "batch_size": 256,
        "eta": 0.03, "momentum": 0.9,
        "task_train_examples": 10000, "task_test_examples": 5000,
        "task_levels": 2, "task_n_in": 1,
        "val_examples": 2000, "split_seed": 2027,
        "eval_split": "validation", "eval_every": 0,
        "diagnostics": "none", "diagnostics_schedule": "final",
        "probe_bs": 512, "seed": 0,
    }
    mismatches = {key: (args.get(key), expected) for key, expected in required.items()
                  if args.get(key) != expected}
    depth = args["depth"]
    expected_lesion = 0.0 if depth == 1 else 1.0 / 3.0
    if abs(args.get("residual_lesion_fraction", 0.0) - expected_lesion) > 1e-12:
        mismatches["residual_lesion_fraction"] = (
            args.get("residual_lesion_fraction"), expected_lesion)
    if mismatches:
        raise RuntimeError(f"protocol mismatch {path}: {mismatches}")
    if row["final"].get("eval_split") != "validation":
        raise RuntimeError(f"test-contaminated development row: {path}")
    if any("eval_acc" in step for step in row.get("steps", [])):
        raise RuntimeError(f"intermediate validation metric: {path}")
    split = row.get("split", {})
    if (not split.get("split_from_training_only")
            or split.get("validation_examples") != 2000
            or split.get("evaluation_split") != "validation"):
        raise RuntimeError(f"invalid validation split {path}: {split}")
    if row.get("provenance", {}).get("git_dirty") is not False:
        raise RuntimeError(f"dirty or unknown source provenance: {path}")
    if depth > 1:
        lesion = row["final"].get("residual_lesion")
        if not lesion or abs(row["lesion_protocol"]["fraction_of_interior_blocks"] - 1 / 3) > 1e-12:
            raise RuntimeError(f"missing final-third lesion: {path}")


def main():
    paths = sorted(glob.glob(os.path.join(ROOT, PREFIX + "*.json")))
    rows = {}
    commits = set()
    split_hashes = {}
    for path in paths:
        with open(path) as handle:
            row = json.load(handle)
        audit_row(path, row)
        args = row["args"]
        key = (args["task_seed"], args["depth"])
        if key in rows:
            raise RuntimeError(f"duplicate row: {key}")
        rows[key] = row
        commits.add(row["provenance"]["git_commit"])
        split_hashes.setdefault(args["task_seed"], set()).add(
            row["split"]["validation_index_sha256"])
    expected = {(task_seed, depth) for task_seed in TASK_SEEDS for depth in DEPTHS}
    if set(rows) != expected or len(commits) != 1:
        raise RuntimeError(f"incomplete/mixed screen: rows={len(rows)}, "
                           f"missing={expected - set(rows)}, extra={set(rows) - expected}, "
                           f"commits={commits}")
    if any(len(hashes) != 1 for hashes in split_hashes.values()):
        raise RuntimeError(f"depths did not share splits within tasks: {split_hashes}")

    print(f"commit={next(iter(commits))} rows={len(rows)} task_seeds={list(TASK_SEEDS)}")
    print("| depth | mean validation (%) | SD | per-task validation (%) | lesion drop (%) |")
    print("|---:|---:|---:|:---|:---|")
    for depth in DEPTHS:
        accs = [100 * rows[(task_seed, depth)]["final"]["val_acc"]
                for task_seed in TASK_SEEDS]
        if depth == 1:
            lesion_text = "--"
        else:
            drops = [100 * rows[(task_seed, depth)]["final"]["residual_lesion"]["lesion_acc_drop"]
                     for task_seed in TASK_SEEDS]
            lesion_text = ", ".join(f"{value:+.2f}" for value in drops)
        print(f"| {depth} | {statistics.mean(accs):.3f} | {statistics.stdev(accs):.3f} | "
              f"{', '.join(f'{value:.2f}' for value in accs)} | {lesion_text} |")

    gains = [100 * (rows[(task_seed, 4)]["final"]["val_acc"]
                     - rows[(task_seed, 1)]["final"]["val_acc"])
             for task_seed in TASK_SEEDS]
    lesion_drops = [100 * rows[(task_seed, 4)]["final"]["residual_lesion"]["lesion_acc_drop"]
                    for task_seed in TASK_SEEDS]
    passed = (statistics.mean(gains) >= 5.0 and all(gain > 0 for gain in gains)
              and statistics.mean(lesion_drops) >= 2.0
              and all(drop > 0 for drop in lesion_drops))
    print(f"frozen d1->d4 gains={gains}, mean={statistics.mean(gains):+.3f}")
    print(f"d4 final-third lesion drops={lesion_drops}, mean={statistics.mean(lesion_drops):+.3f}")
    print(f"C2 BP useful-depth screen: {'PASS' if passed else 'FAIL'}")
    if not passed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()