#!/usr/bin/env python3
"""Task-clustered confidence intervals for correlated auto-zero sampling."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))

from summarize_physical_autozero_p7 import (  # noqa: E402
    interval,
    paired_differences,
    reference_by_condition,
    task_means,
)


KEY_CONDITIONS = (
    "ideal_cds",
    "common_pedestal_10",
    "pedestal_mismatch_0.05",
    "pedestal_mismatch_0.1",
    "pedestal_mismatch_0.25",
    "gain_mismatch_0.01",
    "gain_mismatch_0.05",
    "sample_noise_1",
    "refresh_every_4",
    "combined_mild",
    "combined_mild_refresh4",
    "combined_strong",
    "overclamp_plus_combined_mild",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input", type=Path,
        default=Path(
            "results/physical_bias/p9_grid_correlated_autozero.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/p9_grid_correlated_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_correlated_autozero_p9_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()