#!/usr/bin/env python3
"""Audit the frozen HFA-S2 full ResNet-20 validation baseline."""
import argparse
import json
import math
import os


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--selection", default="results/oral_a_hfa_short_selection.json")
    parser.add_argument("--input", default="results/oral_a_hfa_full/hfa.json")
    parser.add_argument(
        "--out", default="results/oral_a_hfa_full_audit.json")
    args = parser.parse_args()
    with open(args.selection) as handle:
        selection = json.load(handle)
    if selection.get("status") != "selected_full_open":
        raise ValueError("HFA-S1 did not open HFA-S2")
    with open(args.input) as handle:
        record = json.load(handle)
    run_args = record["args"]
    expected = {
        "mode": "hfa", "depth": 20, "width": 16, "seed": 0,
        "loader_seed": 0, "epochs": 200, "train_limit": 0,
        "val_examples": 5000, "split_seed": 2027,
        "eval_split": "validation", "eval_every": 0,
        "augment_train": 1, "lr_schedule": "step",
        "lr_milestones": "100,150", "lr_gamma": 0.1,
        "warmup_epochs": 0, "momentum": 0.9, "weight_decay": 1e-4,
        "normalization": "batchnorm", "output_lr": 0.1,
        "a_scale": 1.0, "alignment_probe": 32,
    }
    for key, value in expected.items():
        if run_args.get(key) != value:
            raise ValueError(f"HFA-S2 {key} drift")
    if float(run_args["lr"]) != float(selection["selected"]["lr"]):
        raise ValueError("HFA-S2 did not copy the selected hidden rate")
    if record["provenance"]["git_tracked_dirty"]:
        raise ValueError("tracked-dirty HFA-S2 result")
    protocol = record["evaluation_protocol"]
    if protocol["test_evaluations"] or protocol["test_used_for_selection"]:
        raise ValueError("HFA-S2 touched test")
    accuracy = float(record["final"]["accuracy"])
    loss = float(record["final"]["loss"])
    early = float(record["diagnostics"]["early_third_mean"])
    finite = (bool(record["final"]["finite"])
              and math.isfinite(accuracy + loss + early))
    output = {
        "protocol": "convolutional_hfa_S2_v1",
        "status": ("strong_baseline" if finite and accuracy >= 0.80
                   else "weak_or_nonfinite_baseline"),
        "accuracy": accuracy, "loss": loss, "finite": finite,
        "early_third_alignment": early,
        "selected_hidden_lr": float(run_args["lr"]),
        "total_macs": int(record["work"]["total_macs_estimate"]),
        "logical_batch_loss_queries": int(
            record["work"]["logical_batch_loss_queries"]),
        "peak_memory_allocated_bytes": int(
            record["hardware"]["peak_memory_allocated_bytes"]),
        "wall_s": float(record["timing"]["total_timed_wall_s"]),
        "source_commit": record["provenance"]["git_commit"],
        "confirmation_test_seeds_touched": False,
    }
    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(output, indent=2))


if __name__ == "__main__":
    main()