#!/usr/bin/env python3
"""Audit the frozen Oral-A failure and quantify its calibration precursor."""
import argparse
import json
import math
import os


def load(path):
    with open(path) as handle:
        return json.load(handle)


def finite(value):
    return isinstance(value, (int, float)) and math.isfinite(value)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--full", default="results/oral_a_dev/sdil_full_r20_s0.json")
    parser.add_argument(
        "--short", default="results/oral_a_short/sdil_channel_gated_lr0.03.json")
    parser.add_argument(
        "--gate", default="results/oral_a_full_gate.json")
    parser.add_argument(
        "--out", default="results/oral_a_failure_diagnosis.json")
    args = parser.parse_args()

    full = load(args.full)
    short = load(args.short)
    gate = load(args.gate)
    expected = {
        "mode": "sdil", "depth": 20, "epochs": 200,
        "vectorizer_mode": "channel_gated", "a_scale": 1.0,
        "eta_A": 0.001, "pert_directions": 1, "pert_every": 4,
    }
    for key, value in expected.items():
        if full["args"].get(key) != value:
            raise ValueError(f"full A3 drift: {key}={full['args'].get(key)!r}")
    if full["provenance"]["git_tracked_dirty"]:
        raise ValueError("full A3 result has tracked source edits")
    if gate["status"] != "failed" or gate["confirmation_test_seeds_touched"]:
        raise ValueError("expected a failed A3 gate with untouched confirmation")

    calibration_rows = []
    first_nonfinite_epoch = None
    eval_rows = []
    for row in full["epochs"]:
        loss = row["train_loss"]
        if first_nonfinite_epoch is None and not finite(loss):
            first_nonfinite_epoch = row["epoch"]
        if finite(row.get("eval_accuracy")):
            eval_rows.append({
                "epoch": row["epoch"],
                "accuracy": row["eval_accuracy"],
                "loss": row["eval_loss"] if finite(row["eval_loss"]) else None,
                "finite": finite(row["eval_loss"]),
            })
        calibration = row.get("calibration")
        if calibration and all(finite(calibration.get(key)) for key in (
                "calibration_mse", "target_power", "prediction_target_cosine")):
            power = calibration["target_power"]
            calibration_rows.append({
                "epoch": row["epoch"],
                "target_power": power,
                "mse_to_target_power": calibration["calibration_mse"] / power,
                "prediction_target_cosine": calibration["prediction_target_cosine"],
            })
    if first_nonfinite_epoch is None or not calibration_rows:
        raise ValueError("A3 failure trajectory is incomplete")
    finite_eval_rows = [row for row in eval_rows if row["finite"]]

    row_by_epoch = {row["epoch"]: row for row in calibration_rows}
    checkpoints = [row_by_epoch[epoch] for epoch in (1, 20, 40, 60, 80, 87)]
    target_growth_1_to_80 = (
        row_by_epoch[80]["target_power"] / row_by_epoch[1]["target_power"])
    target_growth_1_to_87 = (
        row_by_epoch[87]["target_power"] / row_by_epoch[1]["target_power"])
    max_abs_cosine = max(abs(row["prediction_target_cosine"])
                         for row in calibration_rows)
    max_mse_power_gap = max(abs(row["mse_to_target_power"] - 1.0)
                            for row in calibration_rows)
    short_calibration = short["apical_warmup"]["last"]
    output = {
        "protocol": "oral_a_A3_failure_diagnosis_v1",
        "source_paths": {"full": args.full, "short": args.short, "gate": args.gate},
        "source_commits": {
            "full": full["provenance"]["git_commit"],
            "short": short["provenance"]["git_commit"],
        },
        "frozen_outcome": {
            "first_nonfinite_epoch": first_nonfinite_epoch,
            "final_accuracy": full["final"]["accuracy"],
            "final_finite": full["final"]["finite"],
            "last_finite_validation": finite_eval_rows[-1],
            "last_reported_validation": eval_rows[-1],
            "confirmation_test_seeds_touched": False,
        },
        "short_screen_contrast": {
            "accuracy": short["final"]["accuracy"],
            "early_third_alignment": short["diagnostics"]["early_third_mean"],
            "warmup_prediction_target_cosine": (
                short_calibration["prediction_target_cosine"]),
            "warmup_mse_to_target_power": (
                short_calibration["calibration_mse"]
                / short_calibration["target_power"]),
        },
        "calibration_precursor": {
            "checkpoints": checkpoints,
            "target_power_growth_epoch1_to_80": target_growth_1_to_80,
            "target_power_growth_epoch1_to_87": target_growth_1_to_87,
            "max_abs_prediction_target_cosine_before_nonfinite": max_abs_cosine,
            "max_abs_mse_to_target_power_minus_one": max_mse_power_gap,
        },
        "diagnosis": {
            "classification": "causal_target_not_captured_before_runaway",
            "evidence": [
                "prediction-target cosine remains effectively zero",
                "calibration MSE remains indistinguishable from target power",
                "target power grows by orders of magnitude before nonfiniteness",
                "short-run teaching alignment therefore cannot be attributed to learned A",
            ],
            "next_test": (
                "reduce estimator variance by perturbing the representable "
                "channel-gated feedback subspace, not every spatial unit"),
        },
    }
    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, sort_keys=True))


if __name__ == "__main__":
    main()