#!/usr/bin/env python3
"""Audit and gate the frozen RRM-3 full ResNet-20 validation baseline."""
import argparse
import json
import math
import os


SPLIT_HASH = "8328b206a97c420e49e54e3eca4abe3274c4756b084355784ea3fb8059e4515b"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--selection", default="results/residual_mirror_short_gate.json")
    parser.add_argument(
        "--bp_selection", default="results/oral_a_bp_selection.json")
    parser.add_argument("--input", default="results/residual_mirror_full/rrm.json")
    parser.add_argument("--out", default="results/residual_mirror_full_gate.json")
    args = parser.parse_args()
    with open(args.selection) as handle:
        selection = json.load(handle)
    if selection.get("protocol") != "residual_response_mirror_short_v1":
        raise ValueError("unexpected RRM-2 selection protocol")
    if selection.get("status") != "passed":
        raise ValueError("RRM-2 did not open RRM-3")
    with open(args.input) as handle:
        record = json.load(handle)
    run_args = record["args"]
    expected = {
        "mode": "rrm", "depth": 20, "width": 16, "seed": 0,
        "loader_seed": 0, "batch_size": 128, "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", "lr": 0.1, "output_lr": 0.1,
        "a_scale": 1.0, "alignment_probe": 32,
        "mirror_warmup_steps": 20, "mirror_every": 16,
        "mirror_batch_size": 1, "mirror_eta": 0.1,
        "mirror_noise_std": 1.0, "mirror_seed": 3000,
    }
    for key, value in expected.items():
        if run_args.get(key) != value:
            raise ValueError(f"RRM-3 {key} drift")
    if float(selection["selected_rrm"]["lr"]) != float(run_args["lr"]):
        raise ValueError("RRM-3 did not copy the selected hidden rate")
    if record["provenance"]["git_tracked_dirty"]:
        raise ValueError("tracked-dirty RRM-3 result")
    if record["split"]["validation_index_sha256"] != SPLIT_HASH:
        raise ValueError("RRM-3 split drift")
    protocol = record["evaluation_protocol"]
    if protocol["test_evaluations"] or protocol["test_used_for_selection"]:
        raise ValueError("RRM-3 touched test")
    if record.get("calibration_metric_space") != (
            "local_parent_child_response_residual"):
        raise ValueError("RRM-3 metric-space drift")

    with open(args.bp_selection) as handle:
        bp_selection = json.load(handle)
    if bp_selection.get("status") != "passed_primary":
        raise ValueError("matched full BP reference is not frozen")
    with open(bp_selection["selected"]["path"]) as handle:
        bp = json.load(handle)
    bp_macs = int(bp["work"]["total_macs_estimate"])
    bp_accuracy = float(bp["final"]["accuracy"])

    accuracy = float(record["final"]["accuracy"])
    loss = float(record["final"]["loss"])
    diagnostics = record["diagnostics"]
    early = float(diagnostics["early_third_mean"])
    values = diagnostics["teaching_negative_gradient_cosine"]
    all_layer = sum(values) / len(values)
    feedback_cosines = diagnostics["feedback_forward_cosine"]
    norm_ratios = diagnostics["feedback_forward_norm_ratio"]
    finite_values = [
        accuracy, loss, early, all_layer, *feedback_cosines, *norm_ratios]
    finite = (bool(record["final"]["finite"])
              and all(math.isfinite(value) for value in finite_values))
    total_macs = int(record["work"]["total_macs_estimate"])
    queries = int(record["work"]["logical_batch_loss_queries"])
    metrics = {
        "accuracy": accuracy, "loss": loss,
        "early_third_alignment": early, "all_layer_alignment": all_layer,
        "mean_feedback_forward_cosine": (
            sum(feedback_cosines) / len(feedback_cosines)),
        "min_feedback_forward_norm_ratio": min(norm_ratios),
        "max_feedback_forward_norm_ratio": max(norm_ratios),
        "finite": finite, "total_macs": total_macs,
        "bp_total_macs": bp_macs, "mac_ratio_to_bp": total_macs / bp_macs,
        "bp_accuracy": bp_accuracy, "logical_batch_loss_queries": queries,
        "mirror_events": int(record["counters"]["mirror_events"]),
        "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"],
    }
    checks = {
        "finite": finite,
        "accuracy_at_least_0.88": accuracy >= 0.88,
        "early_alignment_at_least_0.50": early >= 0.50,
        "zero_task_loss_queries": queries == 0,
        "macs_at_most_1.15x_bp": total_macs <= 1.15 * bp_macs,
    }
    status = "passed" if all(checks.values()) else "failed"
    output = {
        "protocol": "residual_response_mirror_full_v1",
        "status": status, "checks": checks, "metrics": metrics,
        "innovation_experiment_opened": status == "passed",
        "confirmation_test_seeds_touched": False,
        "review_score_before": 5, "review_score_after": 5,
        "score_change_rule": "inherited full baseline cannot raise score",
    }
    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()