#!/usr/bin/env python3
"""Create task-clustered confidence intervals for the P7 auto-zero study."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np


KEY_CONDITIONS = (
    "ideal_sample_hold",
    "sample_noise_1",
    "sample_gain_0.99",
    "sample_gain_1.01",
    "pedestal_0.1",
    "pedestal_0.25",
    "refresh_every_4",
    "refresh_every_8",
    "combined_mild",
    "combined_strong",
    "overclamp_plus_combined_mild",
)


def task_means(records: list[dict], value_key: str) -> np.ndarray:
    task_indices = sorted({record["task_index"] for record in records})
    return np.asarray([
        np.mean([
            record[value_key] for record in records
            if record["task_index"] == task_index
        ])
        for task_index in task_indices
    ])


def interval(
    values: np.ndarray, bootstrap_indices: np.ndarray
) -> tuple[float, float]:
    replicates = np.mean(values[bootstrap_indices], axis=1)
    low, high = np.quantile(replicates, (0.025, 0.975))
    return float(low), float(high)


def reference_by_condition(reference: dict, method: str) -> list[dict]:
    return [{
        "task_index": record["task_index"],
        "device_seed": record["device_seed"],
        "classification_error": (
            record["methods"][method]["classification_error"]),
        "zero_error": (
            record["methods"][method]["classification_error"] == 0.0),
    } for record in reference["records"]]


def paired_differences(
    condition_records: list[dict], reference_records: list[dict]
) -> tuple[np.ndarray, np.ndarray]:
    lookup = {
        (record["task_index"], record["device_seed"]): record
        for record in reference_records
    }
    error_records = []
    zero_records = []
    for record in condition_records:
        reference = lookup[(record["task_index"], record["device_seed"])]
        error_records.append({
            "task_index": record["task_index"],
            "difference": (
                record["classification_error"]
                - reference["classification_error"]),
        })
        zero_records.append({
            "task_index": record["task_index"],
            "difference": float(record["zero_error"])
            - float(reference["zero_error"]),
        })
    return (
        task_means(error_records, "difference"),
        task_means(zero_records, "difference"),
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input", type=Path,
        default=Path(
            "results/physical_bias/p7_grid_autozero_robustness.json"))
    parser.add_argument(
        "--reference", type=Path,
        default=Path(
            "results/physical_bias/p5_full_grid_bias_crossover.json"))
    parser.add_argument(
        "--output", type=Path,
        default=Path(
            "results/physical_bias/p7_grid_autozero_key_results.json"))
    parser.add_argument("--bootstrap-replicates", type=int, default=20_000)
    parser.add_argument("--seed", type=int, default=20260829)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    report = json.loads(args.input.read_text())
    reference = json.loads(args.reference.read_text())
    task_count = report["protocol"]["task_count"]
    rng = np.random.default_rng(args.seed)
    bootstrap_indices = rng.integers(
        0, task_count, size=(args.bootstrap_replicates, task_count))
    reference_records = {
        method: reference_by_condition(reference, method)
        for method in ("raw", "overclamp")
    }
    results = {}
    for condition in KEY_CONDITIONS:
        selected = [
            record for record in report["records"]
            if record["condition"] == condition
        ]
        error_task_means = task_means(selected, "classification_error")
        zero_task_means = task_means(selected, "zero_error")
        comparisons = {}
        for method in ("raw", "overclamp"):
            error_difference, zero_difference = paired_differences(
                selected, reference_records[method])
            comparisons[method] = {
                "mean_classification_error_difference": float(np.mean(
                    error_difference)),
                "classification_error_difference_95ci": interval(
                    error_difference, bootstrap_indices),
                "zero_error_fraction_difference": float(np.mean(
                    zero_difference)),
                "zero_error_fraction_difference_95ci": interval(
                    zero_difference, bootstrap_indices),
            }
        results[condition] = {
            "trials": len(selected),
            "task_clusters": task_count,
            "mean_classification_error": float(np.mean(error_task_means)),
            "mean_classification_error_95ci": interval(
                error_task_means, bootstrap_indices),
            "zero_error_fraction": float(np.mean(zero_task_means)),
            "zero_error_fraction_95ci": interval(
                zero_task_means, bootstrap_indices),
            "comparisons": comparisons,
        }
    output = {
        "analysis": "physical_grid_hardware_autozero_p7_key_statistics",
        "source": str(args.input),
        "reference": str(args.reference),
        "bootstrap": {
            "unit": "task; four device draws are averaged within each task",
            "task_clusters": task_count,
            "replicates": args.bootstrap_replicates,
            "seed": args.seed,
            "interval": "percentile 95%",
        },
        "results": results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2) + "\n")
    print(json.dumps(results, indent=2))
    print(f"wrote {args.output}")


if __name__ == "__main__":
    main()