#!/usr/bin/env python3
"""Audit the frozen contrastive state-bias B1 gate."""
import argparse
import json
import math
from pathlib import Path

from contrastive_bias_b1 import RESULT_ROOT, bias_cells


def read_json(path):
    with open(path, encoding="utf-8") as handle:
        return json.load(handle)


def cell_path(cell_id):
    return RESULT_ROOT / ("dp-bias-b1-" + cell_id + ".json")


def valid_record(record, cell):
    history = record.get("history") or {}
    return (
        record.get("status") == "completed"
        and record.get("cell_id") == cell["cell_id"]
        and record.get("kind") == cell["kind"]
        and record.get("rule") == cell["rule"]
        and record.get("ratio") == cell["ratio"]
        and history.get("finite") is True
        and history.get("epochs_completed") == 20
        and math.isnan(float(history.get("test_accuracy", float("nan"))))
    )


def metric_max(record, key):
    values = record["history"]["curves"][key]
    return max(abs(float(value)) for value in values)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--out", type=Path,
        default=RESULT_ROOT.parent / "b1_gate.json")
    args = parser.parse_args()
    cells = bias_cells()
    missing = [cell["cell_id"] for cell in cells if not cell_path(cell["cell_id"]).is_file()]
    if missing:
        raise RuntimeError("missing B1 cells: " + ", ".join(missing))
    records = {}
    for cell in cells:
        record = read_json(cell_path(cell["cell_id"]))
        if not valid_record(record, cell):
            raise RuntimeError(f"invalid or incomplete B1 cell: {cell['cell_id']}")
        records[cell["cell_id"]] = record
    clean = records["clean"]
    clean_acc = float(clean["history"]["final_validation_accuracy"])
    common = records["common-activity-r4-raw"]
    common_acc = float(common["history"]["final_validation_accuracy"])
    common_error = metric_max(
        common, "maximum_used_clean_difference_relative_error")
    predictor_instruction_max = max(
        metric_max(record, "instruction_observations_for_predictor")
        for record in records.values()
    )
    innovation_post_max = max(
        metric_max(record, "post_bias_raw_bias_rms_ratio")
        for cell_id, record in records.items() if "-innovation" in cell_id
    )
    candidates = []
    table = []
    for ratio in (0.25, 1.0, 4.0):
        tag = rate_tag = f"{ratio:g}".replace(".", "p")
        row = {"ratio": ratio}
        for rule in ("raw", "innovation", "oracle"):
            record = records[f"activity-r{tag}-{rule}"]
            row[rule] = float(record["history"]["final_validation_accuracy"])
        row["raw_degradation"] = clean_acc - row["raw"]
        row["innovation_clean_gap"] = abs(row["innovation"] - clean_acc)
        row["innovation_oracle_gap"] = abs(row["innovation"] - row["oracle"])
        row["innovation_post_bias_ratio_max"] = metric_max(
            records[f"activity-r{tag}-innovation"],
            "post_bias_raw_bias_rms_ratio")
        row["passes"] = (
            row["raw_degradation"] >= 5.0
            and row["innovation_clean_gap"] <= 2.0
            and row["innovation_oracle_gap"] <= 1.0
            and row["innovation_post_bias_ratio_max"] <= 1e-3
        )
        if row["passes"]:
            candidates.append(ratio)
        table.append(row)
    checks = {
        "clean_at_least_70": clean_acc >= 70.0,
        "common_within_0p2": abs(common_acc - clean_acc) <= 0.2,
        "common_difference_error_at_most_1e_6": common_error <= 1e-6,
        "some_activity_ratio_passes": bool(candidates),
        "all_innovation_post_bias_at_most_1e_3": innovation_post_max <= 1e-3,
        "predictor_saw_zero_instruction_observations": predictor_instruction_max == 0.0,
    }
    source_values = {json.dumps(row["source"], sort_keys=True) for row in records.values()}
    registry_values = {row["registry_sha256"] for row in records.values()}
    checks["single_source_lock"] = len(source_values) == 1
    checks["single_registry_lock"] = len(registry_values) == 1
    report = {
        "stage": "contrastive_bias_b1", "gate": (
            "pass" if all(checks.values()) else "fail"),
        "checks": checks, "clean_final_validation_accuracy": clean_acc,
        "common_final_validation_accuracy": common_acc,
        "common_maximum_difference_relative_error": common_error,
        "innovation_maximum_post_bias_ratio": innovation_post_max,
        "predictor_maximum_instruction_observations": predictor_instruction_max,
        "activity_table": table,
        "selected_confirmation_ratio": max(candidates) if candidates else None,
        "num_expected_records": 17, "num_audited_records": len(records),
        "source": clean["source"], "registry_sha256": clean["registry_sha256"],
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w", encoding="utf-8") as handle:
        json.dump(report, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()