#!/usr/bin/env python3
"""Validate and gate stagewise causally whitened feedback capture."""
import argparse
import json
import math
import os


SPLIT_HASH = "8328b206a97c420e49e54e3eca4abe3274c4756b084355784ea3fb8059e4515b"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input", default="results/oral_a_v6_calibration/result.json")
    parser.add_argument(
        "--out", default="results/oral_a_v6_calibration_gate.json")
    args = parser.parse_args()
    with open(args.input) as handle:
        record = json.load(handle)
    if record.get("protocol") != "oral_a_v6_stagewise_whitened_causal_capture_v1":
        raise ValueError("unexpected V6 protocol")
    expected = {
        "depth": 20, "width": 16, "seed": 0, "loader_seed": 0,
        "batch_size": 128, "train_limit": 10000,
        "val_examples": 5000, "split_seed": 2027,
        "normalization": "batchnorm", "residual_scale": 1.0,
        "feedback_scale": 1.0, "sigma": 0.01,
        "perturb_seed": 5000, "events_per_stage": 20,
        "readout_relative_ridge": 1e-6,
        "conv_diagonal_relative_ridge": 1e-3,
        "alignment_probe": 64, "calibration_augmentation": False,
    }
    if record.get("settings") != expected:
        raise ValueError("V6 settings drift")
    if record["provenance"]["git_tracked_dirty"]:
        raise ValueError("V6 result came from a tracked-dirty tree")
    if record["split"]["validation_index_sha256"] != SPLIT_HASH:
        raise ValueError("V6 split drift")
    if record["test_examples_touched"] or record["validation_endpoints_observed"]:
        raise ValueError("V6 touched a held-out endpoint")
    work = record["work"]
    audit = record["method_audit"]
    fixed = record["fixed_hfa"]
    learned = record["learned_scib"]
    finite_metrics = [
        fixed["early_third_alignment"], fixed["all_layer_alignment"],
        learned["early_third_alignment"], learned["all_layer_alignment"],
        learned["min_feedback_forward_norm_ratio"],
        learned["max_feedback_forward_norm_ratio"],
    ]
    checks = {
        "finite": bool(record["finite"])
                  and all(math.isfinite(value) for value in finite_metrics),
        "exactly_19_stages": work["stages"] == 19,
        "exactly_380_edge_events": work["edge_events"] == 380,
        "exactly_760_batch_loss_queries": (
            work["logical_batch_loss_queries"] == 760),
        "exactly_48640_per_example_observations": (
            work["per_example_causal_observations"] == 48640),
        "stage_order_readout_then_18_to_1": (
            audit["stage_order"] == ["readout"] + list(range(18, 0, -1))),
        "forward_state_bitwise_fixed": (
            audit["forward_state_max_absolute_difference"] == 0.0),
        "zero_forward_weight_reads_in_fit": (
            audit["forward_weight_reads_in_feedback_fit"] == 0),
        "zero_reverse_mode_learning_operations": (
            audit["reverse_mode_learning_operations"] == 0),
        "early_third_at_least_0.10": (
            learned["early_third_alignment"] >= 0.10),
        "all_layer_at_least_0.20": (
            learned["all_layer_alignment"] >= 0.20),
        "early_gain_over_fixed_hfa_at_least_0.08": (
            learned["early_third_alignment"]
            - fixed["early_third_alignment"] >= 0.08),
        "feedback_norm_ratios_in_0.1_to_3": (
            learned["min_feedback_forward_norm_ratio"] >= 0.1
            and learned["max_feedback_forward_norm_ratio"] <= 3.0),
    }
    output = {
        "protocol": "oral_a_v6_stagewise_whitened_causal_capture_gate_v1",
        "status": "passed" if all(checks.values()) else "failed",
        "checks": checks, "fixed_hfa": fixed, "learned_scib": learned,
        "work": work, "source_commit": record["provenance"]["git_commit"],
        "source_result": args.input,
        "conditional_short_task_gate_open": all(checks.values()),
        "confirmation_test_seeds_touched": False,
        "review_score_before": 5, "review_score_after": 5,
        "score_change_rule": "causal capture alone 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({
        "status": output["status"], "checks": checks,
        "fixed_hfa": fixed, "learned_scib": learned,
    }, indent=2))


if __name__ == "__main__":
    main()