#!/usr/bin/env python3
"""Audit the same-path five-seed contrastive-bias confirmation."""
import argparse
import json
import math
from pathlib import Path
import statistics

from contrastive_bias_c1 import CONDITIONS, RESULT_ROOT, SEEDS


T95 = 2.131846786326649


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


def record_path(seed, condition):
    name = f"dp-bias-c1-s{seed}-{condition.replace('_', '-')}"
    return RESULT_ROOT / (name + ".json")


def bound(values, absolute_mean=False):
    mean = statistics.fmean(values)
    sem = statistics.stdev(values) / math.sqrt(len(values))
    center = abs(mean) if absolute_mean else mean
    return {"mean": mean, "sem": sem, "upper_95": center + T95 * sem}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=Path,
                        default=RESULT_ROOT.parent / "c1_gate.json")
    args = parser.parse_args()
    conditions = tuple(row[0] for row in CONDITIONS)
    missing = [
        f"s{seed}:{condition}" for seed in SEEDS for condition in conditions
        if not record_path(seed, condition).is_file()
    ]
    if missing:
        raise RuntimeError("missing C1 cells: " + ", ".join(missing))
    records = {}
    for seed in SEEDS:
        for condition in conditions:
            record = read_json(record_path(seed, condition))
            if (
                record.get("status") != "completed"
                or record.get("seed") != seed
                or record.get("condition") != condition
                or record.get("ratio") != 4.0
                or not math.isnan(float(
                    (record.get("history") or {}).get("test_accuracy", float("nan"))))
            ):
                raise RuntimeError(f"invalid C1 record s{seed}:{condition}")
            records[(seed, condition)] = record

    rows = []
    gains = []
    deficits = []
    innovation_oracle = []
    for seed in SEEDS:
        histories = {
            condition: records[(seed, condition)]["history"]
            for condition in conditions
        }
        clean = histories["same_path_clean"]["final_validation_accuracy"]
        raw = histories["raw"]["final_validation_accuracy"]
        innovation = histories["innovation"]["final_validation_accuracy"]
        oracle = histories["oracle"]["final_validation_accuracy"]
        raw_failed = (
            histories["raw"]["finite"] is not True or clean - raw >= 20.0)
        gains.append(innovation - raw)
        deficits.append(clean - innovation)
        innovation_oracle.append(innovation - oracle)
        uuids = {
            records[(seed, condition)]["hardware"]["uuid"]
            for condition in conditions
        }
        rows.append({
            "seed": seed, "same_path_clean": clean, "raw": raw,
            "innovation": innovation, "oracle": oracle,
            "raw_finite": histories["raw"]["finite"],
            "raw_failed": raw_failed, "innovation_minus_raw": innovation - raw,
            "clean_minus_innovation": clean - innovation,
            "innovation_minus_oracle": innovation - oracle,
            "single_physical_gpu": len(uuids) == 1,
            "physical_gpu_uuid": next(iter(uuids)) if len(uuids) == 1 else None,
        })

    gain_stats = bound(gains)
    gain_stats["lower_95"] = gain_stats["mean"] - T95 * gain_stats["sem"]
    deficit_stats = bound(deficits)
    oracle_stats = bound(innovation_oracle, absolute_mean=True)
    complete_finite = all(
        records[(seed, condition)]["history"]["finite"] is True
        and records[(seed, condition)]["history"]["epochs_completed"] == 130
        for seed in SEEDS
        for condition in ("same_path_clean", "innovation", "oracle")
    )
    clean_mean = statistics.fmean(row["same_path_clean"] for row in rows)
    post_bias_max = max(
        max(abs(value) for value in records[(seed, "innovation")]["history"]
            ["curves"]["post_bias_raw_bias_rms_ratio"])
        for seed in SEEDS
    )
    instruction_max = max(
        max(abs(value) for value in records[(seed, condition)]["history"]
            ["curves"]["instruction_observations_for_predictor"])
        for seed in SEEDS for condition in conditions
    )
    source_locks = {
        json.dumps(record["source"], sort_keys=True) for record in records.values()
    }
    registry_locks = {record["registry_sha256"] for record in records.values()}
    checks = {
        "complete_finite_clean_innovation_oracle": complete_finite,
        "mean_same_path_clean_at_least_80": clean_mean >= 80.0,
        "raw_failed_every_seed": all(row["raw_failed"] for row in rows),
        "innovation_above_raw_every_seed": all(value > 0 for value in gains),
        "gain_lower_95_above_20": gain_stats["lower_95"] > 20.0,
        "clean_deficit_upper_95_below_2": deficit_stats["upper_95"] < 2.0,
        "innovation_oracle_abs_mean_bound_below_2": oracle_stats["upper_95"] < 2.0,
        "post_bias_ratio_at_most_1e_3": post_bias_max <= 1e-3,
        "zero_instruction_observations": instruction_max == 0.0,
        "same_gpu_within_every_seed": all(row["single_physical_gpu"] for row in rows),
        "single_source_lock": len(source_locks) == 1,
        "single_registry_lock": len(registry_locks) == 1,
    }
    first = records[(SEEDS[0], conditions[0])]
    report = {
        "stage": "contrastive_bias_c1", "gate": (
            "pass" if all(checks.values()) else "fail"),
        "checks": checks, "rows": rows,
        "mean_same_path_clean_validation_accuracy": clean_mean,
        "paired_innovation_minus_raw": gain_stats,
        "paired_clean_minus_innovation": deficit_stats,
        "paired_innovation_minus_oracle": oracle_stats,
        "maximum_post_bias_ratio": post_bias_max,
        "maximum_instruction_observations": instruction_max,
        "num_expected_records": 20, "num_audited_records": len(records),
        "source": first["source"], "registry_sha256": first["registry_sha256"],
        "test_policy": "none",
    }
    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()