#!/usr/bin/env python3
"""Audit the frozen nonlinear CLLN hardware confirmation."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


METHODS = (
    "clean",
    "raw",
    "constant",
    "sdil",
    "overclamp_clean",
    "overclamp",
    "overclamp_sdil",
)
DEVICE_SEEDS = (20260833, 20260834, 20260835)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input", type=Path,
        default=Path(
            "results/physical_bias/p11_nonlinear_hardware_confirmation.json"),
    )
    parser.add_argument(
        "--output", type=Path,
        default=Path(
            "results/physical_bias/p11_nonlinear_hardware_confirmation_gate.json"),
    )
    parser.add_argument("--bootstrap-replicates", type=int, default=20000)
    parser.add_argument("--bootstrap-seed", type=int, default=20260829)
    return parser.parse_args()


def task_means(records: list[dict], method: str) -> np.ndarray:
    return np.asarray([
        np.mean([
            record["methods"][method]["classification_error"]
            for record in records
            if record["task_index"] == task
        ])
        for task in range(40)
    ])


def summarize(
    values: np.ndarray, rng: np.random.Generator, replicates: int
) -> dict:
    samples = rng.integers(0, len(values), size=(replicates, len(values)))
    means = np.mean(values[samples], axis=1)
    return {
        "mean_classification_error": float(np.mean(values)),
        "task_bootstrap_95ci": [
            float(value) for value in np.percentile(means, (2.5, 97.5))
        ],
    }


def main() -> None:
    args = parse_args()
    report = json.loads(args.input.read_text())
    protocol = report["protocol"]
    records = report["records"]
    exact_cells = {
        (task, seed) for task in range(40) for seed in DEVICE_SEEDS
    }
    actual_cells = {
        (record["task_index"], record["device_seed"]) for record in records
    }
    protocol_checks = {
        "source_marked_confirmatory": report["confirmatory"] is True,
        "exact_120_paired_cells": actual_cells == exact_cells,
        "all_40_tasks": protocol["task_count"] == 40
        and protocol["rotations_per_input_diameter"] == 8,
        "new_component_seeds_exact": (
            tuple(protocol["device_seeds"]) == DEVICE_SEEDS),
        "seven_methods_exact": tuple(protocol["methods"]) == METHODS,
        "calibration_observations_16": (
            protocol["calibration_observations_per_trial"] == 16),
        "degree_two_predictor": protocol["sdil_polynomial_degree"] == 2,
        "standard_epochs_600": protocol["standard_epochs"] == 600,
        "overclamp_epochs_1000": (
            protocol["overclamp_epochs_maximum"] == 1000),
        "autodiff_unused": report["autodiff_used"] is False,
    }
    finite_complete = len(records) == 120 and all(
        set(record["methods"]) == set(METHODS)
        and all(
            np.isfinite(method["classification_error"])
            and np.isfinite(method["hinge_loss_v2"])
            for method in record["methods"].values()
        )
        for record in records
    )
    rng = np.random.default_rng(args.bootstrap_seed)
    by_method_values = {
        method: task_means(records, method) for method in METHODS
    }
    by_method = {
        method: summarize(values, rng, args.bootstrap_replicates)
        for method, values in by_method_values.items()
    }
    means = {
        method: by_method[method]["mean_classification_error"]
        for method in METHODS
    }
    standard_gap_closed = (
        (means["raw"] - means["sdil"])
        / (means["raw"] - means["clean"])
    )
    overclamp_gap_closed = (
        (means["overclamp"] - means["overclamp_sdil"])
        / (means["overclamp"] - means["overclamp_clean"])
    )
    checks = {
        **protocol_checks,
        "all_task_metrics_finite": finite_complete,
        "sdil_improves_raw_in_every_task": bool(np.all(
            by_method_values["sdil"] < by_method_values["raw"])),
        "sdil_closes_at_least_95_percent_standard_gap": (
            standard_gap_closed >= 0.95),
        "sdil_within_one_point_of_clean": (
            means["sdil"] - means["clean"] <= 0.01),
        "overclamp_sdil_closes_at_least_95_percent_gap": (
            overclamp_gap_closed >= 0.95),
        "constant_and_overclamp_retained": (
            "constant" in by_method and "overclamp" in by_method),
    }
    output = {
        "analysis": "nonlinear_clln_hardware_confirmation_gate",
        "source": str(args.input),
        "bootstrap": {
            "unit": "task; three device draws averaged within task",
            "task_clusters": 40,
            "replicates": args.bootstrap_replicates,
            "seed": args.bootstrap_seed,
            "interval": "percentile 95%",
        },
        "checks": checks,
        "by_method": by_method,
        "standard_raw_to_clean_gap_closed": float(standard_gap_closed),
        "overclamp_raw_to_clean_gap_closed": float(overclamp_gap_closed),
        "gate": "pass" if all(checks.values()) else "fail",
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2) + "\n")
    print(json.dumps(output, indent=2))
    print(f"wrote {args.output}")
    if output["gate"] != "pass":
        raise SystemExit(1)


if __name__ == "__main__":
    main()