#!/usr/bin/env python3
"""Audit the frozen five-seed Rain EP layer-state confirmation."""

from __future__ import annotations

import json
import math
from pathlib import Path
import statistics

from rain_ep_bias_c1 import CONDITIONS, RESULT_ROOT, SEEDS


ROOT = Path(__file__).resolve().parents[1]
T95 = 2.131846786326649


def read(path: Path) -> dict:
    with path.open(encoding="utf-8") as handle:
        return json.load(handle)


def path_for(seed: int, file_mode: str) -> Path:
    return RESULT_ROOT / f"rain-ep-c1-s{seed}-{file_mode}.json"


def interval(values: list[float], absolute_mean: bool = False) -> dict:
    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,
        "lower_95": mean - T95 * sem,
        "upper_95": center + T95 * sem,
    }


def main() -> None:
    records = {}
    missing = []
    for seed in SEEDS:
        for mode, file_mode in CONDITIONS:
            path = path_for(seed, file_mode)
            if not path.is_file():
                missing.append(f"s{seed}:{mode}")
            else:
                records[(seed, mode)] = read(path)
    if missing:
        raise RuntimeError("missing C1 cells: " + ", ".join(missing))

    rows = []
    gain_raw = []
    gain_constant = []
    deficit_clean = []
    difference_oracle = []
    difference_noise = []
    wall_ratios = []
    for seed in SEEDS:
        final = {
            mode: records[(seed, mode)]["final"]
            for mode, _ in CONDITIONS
        }
        accuracy = {
            mode: final[mode]["test_accuracy"] for mode, _ in CONDITIONS
        }
        gain_raw.append(accuracy["innovation"] - accuracy["raw"])
        gain_constant.append(
            accuracy["innovation"] - accuracy["constant"])
        deficit_clean.append(accuracy["clean"] - accuracy["innovation"])
        difference_oracle.append(
            accuracy["innovation"] - accuracy["oracle"])
        difference_noise.append(
            accuracy["same_rms_noise"] - accuracy["clean"])
        wall_ratios.append(
            final["innovation"]["wall_seconds"] / final["clean"]["wall_seconds"])
        rows.append({
            "seed": seed,
            **{mode: value for mode, value in accuracy.items()},
            "innovation_minus_raw": gain_raw[-1],
            "innovation_minus_constant": gain_constant[-1],
            "clean_minus_innovation": deficit_clean[-1],
            "innovation_minus_oracle": difference_oracle[-1],
            "noise_minus_clean": difference_noise[-1],
            "innovation_wall_over_clean": wall_ratios[-1],
        })

    clean_mean = statistics.fmean(row["clean"] for row in rows)
    required_finite = (
        "clean", "same_rms_noise", "constant", "innovation", "oracle")
    finite_complete = all(
        len(records[(seed, mode)]["metrics"]) == 3
        and all(metric["finite"] for metric in records[(seed, mode)]["metrics"])
        for seed in SEEDS for mode in required_finite
    )
    innovation_residual_below_constant = all(
        records[(seed, "innovation")]["final"]["corrector"]
        ["residual_to_clean_state_difference_rms"]
        < records[(seed, "constant")]["final"]["corrector"]
        ["residual_to_clean_state_difference_rms"]
        for seed in SEEDS
    )
    neutral_matched = all(
        records[(seed, mode)]["final"]["corrector"]["neutral_observations"]
        == 128
        for seed in SEEDS for mode in ("constant", "innovation")
    )
    source_revisions = {
        record["sdil"]["revision"] for record in records.values()
    }
    author_revisions = {
        record["author"]["revision"] for record in records.values()
    }
    protocol_fixed = all(
        record["protocol"]["evaluation_split"] == "train_holdout"
        and record["protocol"]["data_seed"] == 6100
        and record["protocol"]["bias_ratio"] == 0.01
        and record["protocol"]["adapter"] == "layer"
        and record["protocol"]["autodiff_used_for_learning"] is False
        and record["protocol"]["extra_equilibrium_phases_for_predictor"] == 0
        for record in records.values()
    )
    stats = {
        "innovation_minus_raw": interval(gain_raw),
        "innovation_minus_constant": interval(gain_constant),
        "clean_minus_innovation": interval(deficit_clean),
        "innovation_minus_oracle": interval(
            difference_oracle, absolute_mean=True),
        "noise_minus_clean": interval(difference_noise, absolute_mean=True),
        "innovation_wall_over_clean": interval(wall_ratios),
    }
    checks = {
        "finite_complete_required_conditions": finite_complete,
        "mean_clean_validation_at_least_70": clean_mean >= 0.70,
        "raw_damage_at_least_10_points_every_seed": all(
            row["clean"] - row["raw"] >= 0.10 for row in rows),
        "innovation_beats_raw_every_seed": all(value > 0 for value in gain_raw),
        "innovation_raw_lower_95_at_least_10_points": (
            stats["innovation_minus_raw"]["lower_95"] >= 0.10),
        "innovation_beats_constant_every_seed": all(
            value > 0 for value in gain_constant),
        "innovation_constant_lower_95_at_least_5_points": (
            stats["innovation_minus_constant"]["lower_95"] >= 0.05),
        "clean_deficit_upper_95_below_3_points": (
            stats["clean_minus_innovation"]["upper_95"] < 0.03),
        "oracle_abs_mean_upper_95_below_3_points": (
            stats["innovation_minus_oracle"]["upper_95"] < 0.03),
        "noise_abs_mean_upper_95_below_3_points": (
            stats["noise_minus_clean"]["upper_95"] < 0.03),
        "innovation_residual_below_constant_every_seed": (
            innovation_residual_below_constant),
        "matched_128_neutral_observations": neutral_matched,
        "mean_wall_overhead_at_most_15_percent": (
            stats["innovation_wall_over_clean"]["mean"] <= 1.15),
        "single_sdil_revision": len(source_revisions) == 1,
        "single_author_revision": len(author_revisions) == 1,
        "protocol_fixed_and_bp_free": protocol_fixed,
    }
    report = {
        "stage": "rain_ep_bias_c1",
        "gate": "pass" if all(checks.values()) else "fail",
        "checks": checks,
        "rows": rows,
        "statistics": stats,
        "mean_clean_validation_accuracy": clean_mean,
        "num_records": len(records),
        "sdil_revision": next(iter(source_revisions)),
        "author_revision": next(iter(author_revisions)),
        "test_policy": "test_set_never_loaded; fixed epoch-3 train holdout",
    }
    output = RESULT_ROOT.parent / "c1_gate.json"
    output.write_text(
        json.dumps(report, indent=2, sort_keys=True, allow_nan=False) + "\n")
    print(json.dumps(report, indent=2, sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()