#!/usr/bin/env python3
"""Build the hardware-realistic CLLN evidence figure from frozen results."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np


METHODS = (
    ("clean", "Clean", "#222222"),
    ("raw", "Raw", "#D55E00"),
    ("constant", "Static\ncalibration", "#777777"),
    ("overclamp", "Overclamp", "#E69F00"),
    ("sdil", "SDIL", "#0072B2"),
    ("overclamp_sdil", "Overclamp\n+ SDIL", "#56B4E9"),
)

COMBINED_CONDITIONS = (
    ("ideal_cds", "Ideal\nCDS"),
    ("common_pedestal_10", "Common\npedestal"),
    ("sample_noise_1", "Sample\nnoise"),
    ("refresh_every_4", "Refresh\nevery 4"),
    ("combined_mild_refresh4", "Mild\ncombined"),
    ("combined_strong", "Strong\ncombined"),
    ("overclamp_plus_combined_mild", "Overclamp\n+ mild"),
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--physical",
        type=Path,
        default=Path(
            "results/physical_bias/p5_full_grid_bias_crossover.json"),
    )
    parser.add_argument(
        "--sampler",
        type=Path,
        default=Path(
            "results/physical_bias/p9_grid_correlated_autozero.json"),
    )
    parser.add_argument(
        "--spice",
        type=Path,
        default=Path(
            "results/physical_bias/p8_spice_autozero_primitive.json"),
    )
    parser.add_argument(
        "--output-analysis",
        type=Path,
        default=Path(
            "results/physical_bias/p10_hardware_evidence_analysis.json"),
    )
    parser.add_argument(
        "--output-csv",
        type=Path,
        default=Path(
            "results/physical_bias/p10_hardware_evidence_source.csv"),
    )
    parser.add_argument(
        "--output-figure",
        type=Path,
        default=Path("results/figs/figure_physical_hardware_evidence"),
    )
    parser.add_argument("--bootstrap-replicates", type=int, default=20000)
    parser.add_argument("--bootstrap-seed", type=int, default=20260829)
    return parser.parse_args()


def clustered_summary(
    records: list[dict],
    value,
    *,
    rng: np.random.Generator,
    replicates: int,
) -> dict:
    tasks = sorted({record["task_index"] for record in records})
    task_means = np.asarray([
        np.mean([value(record) for record in records
                 if record["task_index"] == task])
        for task in tasks
    ])
    samples = rng.integers(
        0, len(task_means), size=(replicates, len(task_means)))
    bootstrap = np.mean(task_means[samples], axis=1)
    return {
        "mean": float(np.mean(task_means)),
        "task_bootstrap_95ci": [
            float(bound) for bound in np.percentile(bootstrap, (2.5, 97.5))
        ],
        "task_clusters": len(tasks),
        "trials": len(records),
    }


def physical_method_summaries(
    report: dict, rng: np.random.Generator, replicates: int
) -> dict:
    output = {}
    for method, _, _ in METHODS:
        output[method] = clustered_summary(
            report["records"],
            lambda record, name=method: record["methods"][name][
                "classification_error"],
            rng=rng,
            replicates=replicates,
        )
    return output


def sampler_condition_summary(
    report: dict,
    condition: str,
    rng: np.random.Generator,
    replicates: int,
) -> dict:
    selected = [
        record for record in report["records"]
        if record["condition"] == condition
    ]
    summary = clustered_summary(
        selected,
        lambda record: record["classification_error"],
        rng=rng,
        replicates=replicates,
    )
    summary["zero_error_fraction"] = float(np.mean([
        record["classification_error"] == 0.0 for record in selected
    ]))
    summary["solver_failure_fraction"] = float(np.mean([
        record["status"] != "completed" for record in selected
    ]))
    finite_rmse = [
        record["applied_rate_rmse_v_per_s"] for record in selected
        if record["applied_rate_rmse_v_per_s"] is not None
    ]
    summary["median_applied_rate_rmse_v_per_s"] = float(np.median([
        value for value in finite_rmse
    ]))
    return summary


def build_analysis(
    physical: dict,
    sampler: dict,
    spice: dict,
    *,
    replicates: int,
    seed: int,
) -> dict:
    rng = np.random.default_rng(seed)
    sampler_conditions = {
        condition["name"] for condition in sampler["protocol"]["conditions"]
    }
    required_conditions = {
        name for name, _ in COMBINED_CONDITIONS
    } | {
        f"pedestal_mismatch_{value:g}"
        for value in (0.01, 0.025, 0.05, 0.1, 0.25, 0.5)
    } | {
        f"gain_mismatch_{value:g}"
        for value in (0.001, 0.005, 0.01, 0.025, 0.05)
    }
    missing = required_conditions - sampler_conditions
    if missing:
        raise ValueError(f"sampler report is missing {sorted(missing)}")
    condition_summaries = {
        condition: sampler_condition_summary(
            sampler, condition, rng, replicates)
        for condition in sorted(required_conditions)
    }
    return {
        "analysis": "physical_clln_hardware_evidence_figure",
        "confirmatory": False,
        "bootstrap": {
            "unit": "task; four device draws averaged within task",
            "task_clusters": 40,
            "replicates": replicates,
            "seed": seed,
            "interval": "percentile 95%",
        },
        "physical_methods": physical_method_summaries(
            physical, rng, replicates),
        "sampler_conditions": condition_summaries,
        "spice_primitive": {
            "publication_evidence": spice["publication_evidence"],
            "scope": spice["scope"],
            "configuration_count": spice["configuration_count"],
            "fraction_below_1_percent_error_at_start": spice[
                "fraction_below_1_percent_error_at_start"],
            "fraction_below_1_percent_error_at_end": spice[
                "fraction_below_1_percent_error_at_end"],
            "reference_configuration": spice["reference_configuration"],
        },
    }


def write_csv(path: Path, analysis: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow((
            "panel", "series", "condition", "x", "x_unit", "mean_error",
            "ci_low", "ci_high", "zero_error_fraction",
        ))
        for method, label, _ in METHODS:
            summary = analysis["physical_methods"][method]
            writer.writerow((
                "a", "physical_method", label.replace("\n", " "), "", "",
                summary["mean"], *summary["task_bootstrap_95ci"], "",
            ))
        for family, values in (
            ("pedestal mismatch", (0.01, 0.025, 0.05, 0.1, 0.25, 0.5)),
            ("gain mismatch", (0.001, 0.005, 0.01, 0.025, 0.05)),
        ):
            prefix = family.replace(" ", "_")
            for value in values:
                summary = analysis["sampler_conditions"][f"{prefix}_{value:g}"]
                writer.writerow((
                    "b", family, f"{prefix}_{value:g}",
                    summary["median_applied_rate_rmse_v_per_s"], "V/s",
                    summary["mean"], *summary["task_bootstrap_95ci"],
                    summary["zero_error_fraction"],
                ))
        for condition, label in COMBINED_CONDITIONS:
            summary = analysis["sampler_conditions"][condition]
            writer.writerow((
                "c", "sampling condition", label.replace("\n", " "), "", "",
                summary["mean"], *summary["task_bootstrap_95ci"],
                summary["zero_error_fraction"],
            ))


def errorbar(axis, x, summaries, **kwargs) -> None:
    means = np.asarray([summary["mean"] for summary in summaries]) * 100.0
    intervals = np.asarray([
        summary["task_bootstrap_95ci"] for summary in summaries
    ]) * 100.0
    axis.errorbar(
        x,
        means,
        yerr=np.vstack((means - intervals[:, 0], intervals[:, 1] - means)),
        capsize=2.2,
        **kwargs,
    )


def plot(path: Path, analysis: dict) -> None:
    mpl.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 8.5,
        "axes.labelsize": 9,
        "axes.titlesize": 9.5,
        "legend.fontsize": 7.5,
        "xtick.labelsize": 7.7,
        "ytick.labelsize": 8,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "svg.fonttype": "none",
        "pdf.fonttype": 42,
        "figure.facecolor": "white",
        "axes.facecolor": "white",
    })
    figure, axes = plt.subplots(1, 3, figsize=(10.7, 3.15))

    method_summaries = [
        analysis["physical_methods"][method] for method, _, _ in METHODS
    ]
    method_means = np.asarray([
        summary["mean"] for summary in method_summaries]) * 100.0
    method_intervals = np.asarray([
        summary["task_bootstrap_95ci"] for summary in method_summaries
    ]) * 100.0
    positions = np.arange(len(METHODS))[::-1]
    axes[0].barh(
        positions,
        method_means,
        color=[color for _, _, color in METHODS],
        height=0.70,
        edgecolor="white",
        linewidth=0.4,
    )
    axes[0].errorbar(
        method_means,
        positions,
        xerr=np.vstack((
            method_means - method_intervals[:, 0],
            method_intervals[:, 1] - method_means,
        )),
        fmt="none",
        ecolor="#222222",
        elinewidth=0.8,
        capsize=2.2,
    )
    axes[0].set_yticks(
        positions, [label.replace("\n", " ") for _, label, _ in METHODS])
    axes[0].set_xlabel("Final classification error (%)")
    axes[0].set_title("(a) Nonlinear CLLN with component errors")

    sweep_specs = (
        ("pedestal_mismatch", (0.01, 0.025, 0.05, 0.1, 0.25, 0.5),
         "Pedestal mismatch", "#D55E00", "o"),
        ("gain_mismatch", (0.001, 0.005, 0.01, 0.025, 0.05),
         "Gain mismatch", "#0072B2", "D"),
    )
    for prefix, values, label, color, marker in sweep_specs:
        summaries = [
            analysis["sampler_conditions"][f"{prefix}_{value:g}"]
            for value in values
        ]
        x = np.asarray([
            summary["median_applied_rate_rmse_v_per_s"]
            for summary in summaries
        ])
        errorbar(
            axes[1], x, summaries, color=color, marker=marker,
            markersize=4.5, linewidth=1.35, label=label,
        )
    axes[1].set_xscale("log")
    axes[1].set_xlabel("Residual sampling error (V/s, RMSE)")
    axes[1].set_ylabel("Final classification error (%)")
    axes[1].set_title("(b) Local sample-path mismatch")
    axes[1].legend(frameon=False, loc="upper left")

    combined = [
        analysis["sampler_conditions"][condition]
        for condition, _ in COMBINED_CONDITIONS
    ]
    combined_means = np.asarray([
        summary["mean"] for summary in combined]) * 100.0
    combined_intervals = np.asarray([
        summary["task_bootstrap_95ci"] for summary in combined
    ]) * 100.0
    positions = np.arange(len(COMBINED_CONDITIONS))[::-1]
    axes[2].barh(
        positions,
        combined_means,
        color="#0072B2",
        height=0.70,
        edgecolor="white",
        linewidth=0.4,
    )
    axes[2].errorbar(
        combined_means,
        positions,
        xerr=np.vstack((
            combined_means - combined_intervals[:, 0],
            combined_intervals[:, 1] - combined_means,
        )),
        fmt="none",
        ecolor="#222222",
        elinewidth=0.8,
        capsize=2.2,
    )
    axes[2].set_yticks(
        positions,
        [label.replace("\n", " ") for _, label in COMBINED_CONDITIONS],
    )
    axes[2].set_xlabel("Final classification error (%)")
    axes[2].set_title("(c) Nonideal local sampling")

    axes[0].set_xlim(-0.8, 32.0)
    axes[1].set_ylim(-0.8, 32.0)
    axes[2].set_xlim(-0.8, 32.0)
    for index, axis in enumerate(axes):
        axis.grid(
            axis="y" if index == 1 else "x",
            color="#D9D9D9", linewidth=0.55, alpha=0.8)
        axis.tick_params(length=3)
    figure.text(
        0.995,
        0.005,
        "40 tasks × 4 device draws; error bars are task-bootstrap 95% intervals",
        ha="right",
        va="bottom",
        fontsize=7,
        color="#666666",
    )
    figure.tight_layout(rect=(0.0, 0.065, 1.0, 1.0), w_pad=2.0)
    path.parent.mkdir(parents=True, exist_ok=True)
    figure.savefig(path.with_suffix(".svg"), bbox_inches="tight")
    figure.savefig(path.with_suffix(".pdf"), bbox_inches="tight")
    figure.savefig(path.with_suffix(".png"), dpi=240, bbox_inches="tight")
    plt.close(figure)


def main() -> None:
    args = parse_args()
    physical = json.loads(args.physical.read_text())
    sampler = json.loads(args.sampler.read_text())
    spice = json.loads(args.spice.read_text())
    analysis = build_analysis(
        physical,
        sampler,
        spice,
        replicates=args.bootstrap_replicates,
        seed=args.bootstrap_seed,
    )
    analysis["sources"] = [
        str(args.physical), str(args.sampler), str(args.spice)]
    args.output_analysis.parent.mkdir(parents=True, exist_ok=True)
    args.output_analysis.write_text(json.dumps(analysis, indent=2) + "\n")
    write_csv(args.output_csv, analysis)
    plot(args.output_figure, analysis)
    print(json.dumps({
        "physical_methods": analysis["physical_methods"],
        "combined_strong": analysis["sampler_conditions"]["combined_strong"],
    }, indent=2))
    print(f"wrote {args.output_analysis}")
    print(f"wrote {args.output_csv}")
    print(f"wrote {args.output_figure}.svg/.pdf/.png")


if __name__ == "__main__":
    main()