#!/usr/bin/env python3
"""Analyze and plot the paired digital CLLN scaling ladder."""

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


METHOD_ORDER = (
    "clean",
    "matched_noise",
    "raw",
    "constant",
    "sdil",
)
DISPLAY = {
    "clean": "Clean",
    "matched_noise": "Same-RMS noise",
    "raw": "Raw imperfection",
    "constant": "Static calibration",
    "sdil": "SDIL",
}
STYLE = {
    "clean": dict(color="#222222", marker="^", linestyle=":"),
    "matched_noise": dict(color="#CC79A7", marker="D", linestyle="--"),
    "raw": dict(color="#D55E00", marker="X", linestyle="-"),
    "constant": dict(color="#666666", marker="s", linestyle="--"),
    "sdil": dict(color="#0072B2", marker="o", linestyle="-"),
}

FROZEN_SIZES = (4, 8, 12, 16, 24, 32)
FROZEN_DEVICE_SEEDS = (20260830, 20260831, 20260832)
FROZEN_LEARNING_TIMES = {
    4: 0.01,
    8: 0.01,
    12: 0.03,
    16: 0.03,
    24: 0.03,
    32: 0.01,
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--core", type=Path, nargs="+",
        default=[Path(
            "results/coupled_ladder/p1_imperfection_pilot.json")])
    parser.add_argument(
        "--baselines", type=Path, nargs="*",
        default=[Path(
            "results/coupled_ladder/p1_bias_baseline_pilot.json")])
    parser.add_argument(
        "--output-analysis", type=Path,
        default=Path("results/coupled_ladder/p1_scaling_analysis.json"))
    parser.add_argument(
        "--output-csv", type=Path,
        default=Path("results/coupled_ladder/p1_scaling_source.csv"))
    parser.add_argument(
        "--output-figure", type=Path,
        default=Path("results/figs/figure_clln_scaling_pilot"))
    parser.add_argument("--bootstrap-replicates", type=int, default=20000)
    parser.add_argument("--bootstrap-seed", type=int, default=20260829)
    parser.add_argument(
        "--confirmatory",
        action="store_true",
        help="label the output as confirmation after validating its sources",
    )
    return parser.parse_args()


def merge_reports(reports: list[dict]) -> list[dict]:
    records = {}
    for report in reports:
        for record in report["records"]:
            key = (
                record["side"],
                record["task_index"],
                record["device_seed"],
            )
            if key not in records:
                records[key] = {
                    key: value for key, value in record.items()
                    if key != "methods"
                }
                records[key]["methods"] = {}
            overlap = set(records[key]["methods"]) & set(record["methods"])
            if overlap:
                raise ValueError(f"duplicate method records: {sorted(overlap)}")
            records[key]["methods"].update(record["methods"])
    merged = list(records.values())
    merged.sort(key=lambda record: (
        record["side"], record["task_index"], record["device_seed"]))
    return merged


def validate_core_confirmation(reports: list[dict], records: list[dict]) -> None:
    """Reject a confirmatory label unless the frozen core panel is complete."""
    expected_cells = {
        (side, task, seed)
        for side in FROZEN_SIZES
        for task in range(40)
        for seed in FROZEN_DEVICE_SEEDS
    }
    actual_cells = {
        (record["side"], record["task_index"], record["device_seed"])
        for record in records
    }
    if actual_cells != expected_cells:
        missing = len(expected_cells - actual_cells)
        extra = len(actual_cells - expected_cells)
        raise ValueError(
            f"incomplete frozen confirmation: {missing} missing, {extra} extra cells")
    required_methods = set(METHOD_ORDER)
    for record in records:
        if set(record["methods"]) != required_methods:
            raise ValueError("confirmatory cells must contain all five core methods")
        if any(
            method["status"] != "completed"
            for method in record["methods"].values()
        ):
            raise ValueError("confirmatory cells must complete without failures")
    observed_sizes = set()
    for report in reports:
        protocol = report["protocol"]
        if protocol["rotations_per_input_diameter"] != 8:
            raise ValueError("confirmation requires all eight task rotations")
        if protocol["epochs"] != 600:
            raise ValueError("confirmation requires the frozen 600 epochs")
        if tuple(protocol["device_seeds"]) != FROZEN_DEVICE_SEEDS:
            raise ValueError("confirmation device seeds do not match the frozen panel")
        if set(protocol["methods"]) != required_methods:
            raise ValueError("confirmation methods do not match the frozen core panel")
        for side in protocol["sizes"]:
            observed_sizes.add(side)
            observed = float(protocol["learning_time_seconds_by_side"][str(side)])
            if abs(observed - FROZEN_LEARNING_TIMES[side]) > 1e-15:
                raise ValueError(f"side {side} uses an unfrozen learning exposure")
    if observed_sizes != set(FROZEN_SIZES):
        raise ValueError("confirmation sources do not cover all frozen sizes")


def task_cluster_values(
    records: list[dict], side: int, method: str, metric: str
) -> tuple[np.ndarray, np.ndarray]:
    task_indices = sorted({
        record["task_index"] for record in records
        if record["side"] == side and method in record["methods"]
    })
    values = []
    for task_index in task_indices:
        device_values = [
            record["methods"][method][metric]
            for record in records
            if record["side"] == side
            and record["task_index"] == task_index
            and method in record["methods"]
        ]
        values.append(float(np.mean(device_values)))
    return np.asarray(task_indices), np.asarray(values)


def bootstrap_mean(
    values: np.ndarray,
    rng: np.random.Generator,
    replicates: int,
) -> tuple[float, list[float]]:
    indices = rng.integers(0, len(values), size=(replicates, len(values)))
    means = np.mean(values[indices], axis=1)
    return float(np.mean(values)), [
        float(value) for value in np.percentile(means, (2.5, 97.5))
    ]


def task_size_matrix(
    records: list[dict], sizes: list[int], method: str, metric: str
) -> tuple[np.ndarray, np.ndarray]:
    common_tasks = None
    by_size = {}
    for side in sizes:
        tasks, values = task_cluster_values(records, side, method, metric)
        by_size[side] = dict(zip(tasks.tolist(), values.tolist()))
        task_set = set(tasks.tolist())
        common_tasks = task_set if common_tasks is None else (
            common_tasks & task_set)
    task_indices = np.asarray(sorted(common_tasks), dtype=int)
    matrix = np.asarray([
        [by_size[side][int(task)] for side in sizes]
        for task in task_indices
    ])
    return task_indices, matrix


def slope_bootstrap(
    records: list[dict],
    sizes: list[int],
    edges: np.ndarray,
    *,
    replicates: int,
    seed: int,
) -> dict:
    _, clean = task_size_matrix(
        records, sizes, "clean", "classification_error")
    _, constant = task_size_matrix(
        records, sizes, "constant", "classification_error")
    _, sdil = task_size_matrix(
        records, sizes, "sdil", "classification_error")
    x = np.log10(edges)

    def slope(matrix: np.ndarray) -> float:
        return float(np.polyfit(x, np.mean(matrix, axis=0), 1)[0])

    constant_excess = constant - clean
    sdil_excess = sdil - clean
    point_constant = slope(constant_excess)
    point_sdil = slope(sdil_excess)
    point_difference = point_constant - point_sdil
    point_reduction = (
        100.0 * point_difference / point_constant
        if point_constant > 0.0 else None)

    rng = np.random.default_rng(seed)
    constant_slopes = []
    sdil_slopes = []
    slope_differences = []
    relative_reductions = []
    for _ in range(replicates):
        sample = rng.integers(0, len(clean), size=len(clean))
        constant_slope = slope(constant_excess[sample])
        sdil_slope = slope(sdil_excess[sample])
        constant_slopes.append(constant_slope)
        sdil_slopes.append(sdil_slope)
        slope_differences.append(constant_slope - sdil_slope)
        if constant_slope > 1e-12:
            relative_reductions.append(
                100.0 * (constant_slope - sdil_slope) / constant_slope)

    def interval(values: list[float]) -> list[float]:
        return [
            float(value)
            for value in np.percentile(values, (2.5, 97.5))
        ]

    return {
        "x_axis": "log10 learnable edges",
        "y_axis": "classification error excess over clean",
        "static_calibration_slope": point_constant,
        "static_calibration_slope_95ci": interval(constant_slopes),
        "sdil_slope": point_sdil,
        "sdil_slope_95ci": interval(sdil_slopes),
        "paired_slope_difference": point_difference,
        "paired_slope_difference_95ci": interval(slope_differences),
        "relative_slope_reduction_percent": point_reduction,
        "relative_slope_reduction_95ci": (
            interval(relative_reductions) if relative_reductions else None),
        "relative_reduction_valid_bootstrap_fraction": float(
            len(relative_reductions) / replicates),
    }


def trace_summary(
    records: list[dict], side: int, method: str
) -> list[dict]:
    source_records = [
        record
        for record in records
        if record["side"] == side and method in record["methods"]
    ]
    task_indices = sorted({
        record["task_index"] for record in source_records})
    epochs = [
        record["epoch"]
        for record in source_records[0]["methods"][method]["trace"]
    ]
    output = []
    for index, epoch in enumerate(epochs):
        values = np.asarray([
            np.mean([
                record["methods"][method]["trace"][index][
                    "classification_error"
                ]
                for record in source_records
                if record["task_index"] == task_index
            ])
            for task_index in task_indices
        ])
        output.append({
            "epoch": epoch,
            "mean_classification_error": float(np.mean(values)),
            "standard_error": float(
                np.std(values, ddof=1) / np.sqrt(len(values))
                if len(values) > 1 else 0.0),
        })
    return output


def build_analysis(
    records: list[dict], *, replicates: int, seed: int,
    confirmatory: bool = False,
) -> dict:
    sizes = sorted({record["side"] for record in records})
    task_size_draw_counts = {
        len({
            record["device_seed"] for record in records
            if record["side"] == side
            and record["task_index"] == task_index
        })
        for side in sizes
        for task_index in {
            record["task_index"] for record in records
            if record["side"] == side
        }
    }
    if len(task_size_draw_counts) != 1:
        raise ValueError(
            "every task-size cell must contain the same number of device draws")
    component_draws_per_task_size = task_size_draw_counts.pop()
    edges = np.asarray([
        next(
            record["learnable_edges"] for record in records
            if record["side"] == side
        ) for side in sizes
    ])
    rng = np.random.default_rng(seed)
    summaries = {}
    for side in sizes:
        side_summary = {}
        for method in METHOD_ORDER:
            metrics = {}
            for metric in (
                "classification_error",
                "classification_error_auc",
                "reached_stable_zero_error",
            ):
                _, values = task_cluster_values(
                    records, side, method, metric)
                mean, interval = bootstrap_mean(
                    values, rng, replicates)
                metrics[metric] = {
                    "mean": mean,
                    "task_bootstrap_95ci": interval,
                }
            side_summary[method] = metrics
        clean = side_summary["clean"]["classification_error"]["mean"]
        raw = side_summary["raw"]["classification_error"]["mean"]
        sdil = side_summary["sdil"]["classification_error"]["mean"]
        side_summary["sdil_raw_to_clean_gap_closed"] = float(
            (raw - sdil) / (raw - clean))
        summaries[str(side)] = {
            "learnable_edges": int(edges[sizes.index(side)]),
            "methods": side_summary,
        }
    return {
        "analysis": "digital_coupled_ladder_scaling_pilot_analysis",
        "confirmatory": confirmatory,
        "bootstrap": {
            "unit": "task; component draws averaged within task",
            "task_clusters": len({
                record["task_index"] for record in records}),
            "component_draws_per_task_size": component_draws_per_task_size,
            "replicates": replicates,
            "seed": seed,
            "interval": "percentile 95%",
        },
        "sizes": sizes,
        "summaries": summaries,
        "excess_error_scaling": slope_bootstrap(
            records,
            sizes,
            edges,
            replicates=replicates,
            seed=seed + 1,
        ),
        "largest_grid_learning_curves": {
            method: trace_summary(records, sizes[-1], method)
            for method in METHOD_ORDER
        },
    }


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((
            "side", "learnable_edges", "method", "metric", "mean",
            "ci_low", "ci_high",
        ))
        for side, side_summary in analysis["summaries"].items():
            for method in METHOD_ORDER:
                for metric, values in side_summary["methods"][method].items():
                    writer.writerow((
                        side,
                        side_summary["learnable_edges"],
                        method,
                        metric,
                        values["mean"],
                        values["task_bootstrap_95ci"][0],
                        values["task_bootstrap_95ci"][1],
                    ))


def plot_figure(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": 8,
        "ytick.labelsize": 8,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "svg.fonttype": "none",
        "pdf.fonttype": 42,
        "figure.facecolor": "white",
        "axes.facecolor": "white",
    })
    sizes = analysis["sizes"]
    edges = np.asarray([
        analysis["summaries"][str(side)]["learnable_edges"]
        for side in sizes
    ])
    figure, axes = plt.subplots(1, 3, figsize=(10.6, 3.05))
    edge_labels = [
        f"{value / 1000.0:.1f}k" if value >= 1000 else str(value)
        for value in edges
    ]

    for method in METHOD_ORDER:
        means = np.asarray([
            analysis["summaries"][str(side)]["methods"][method][
                "classification_error"
            ]["mean"] * 100.0
            for side in sizes
        ])
        intervals = np.asarray([
            analysis["summaries"][str(side)]["methods"][method][
                "classification_error"
            ]["task_bootstrap_95ci"]
            for side in sizes
        ]) * 100.0
        axes[0].errorbar(
            edges,
            means,
            yerr=np.vstack((means - intervals[:, 0], intervals[:, 1] - means)),
            linewidth=1.7 if method == "sdil" else 1.15,
            markersize=4.5,
            capsize=2.0,
            label=DISPLAY[method],
            **STYLE[method],
        )
    axes[0].set_xscale("log", base=2)
    axes[0].set_xticks(edges, edge_labels)
    axes[0].set_ylim(-3, 70)
    axes[0].set_xlabel("Learnable edges")
    axes[0].set_ylabel("Final classification error (%)")
    reduction = analysis["excess_error_scaling"][
        "relative_slope_reduction_percent"]
    axes[0].set_title(
        f"(a) Error-growth slope reduced {reduction:.0f}% vs static calibration")

    for method in METHOD_ORDER:
        means = np.asarray([
            analysis["summaries"][str(side)]["methods"][method][
                "reached_stable_zero_error"
            ]["mean"] * 100.0
            for side in sizes
        ])
        axes[1].plot(
            edges,
            means,
            linewidth=1.7 if method == "sdil" else 1.15,
            markersize=4.5,
            label=DISPLAY[method],
            **STYLE[method],
        )
    axes[1].set_xscale("log", base=2)
    axes[1].set_xticks(edges, edge_labels)
    axes[1].set_ylim(-5, 105)
    axes[1].set_xlabel("Learnable edges")
    axes[1].set_ylabel("Stable zero-error runs (%)")
    axes[1].set_title("(b) Recovery remains reliable")

    for method in METHOD_ORDER:
        trace = analysis["largest_grid_learning_curves"][method]
        epochs = np.asarray([record["epoch"] for record in trace])
        means = np.asarray([
            record["mean_classification_error"] * 100.0
            for record in trace
        ])
        standard_errors = np.asarray([
            record["standard_error"] * 100.0 for record in trace
        ])
        axes[2].plot(
            epochs,
            means,
            linewidth=1.7 if method == "sdil" else 1.15,
            label=DISPLAY[method],
            **{key: value for key, value in STYLE[method].items()
               if key != "marker"},
        )
        axes[2].fill_between(
            epochs,
            np.maximum(0.0, means - standard_errors),
            means + standard_errors,
            color=STYLE[method]["color"],
            alpha=0.08,
            linewidth=0,
        )
    axes[2].set_xlim(0, 600)
    axes[2].set_ylim(-3, 70)
    axes[2].set_xlabel("Training epoch")
    axes[2].set_ylabel("Classification error (%)")
    axes[2].set_title("(c) Learning at 2,048 edges")

    for axis in axes:
        axis.grid(axis="y", color="#D9D9D9", linewidth=0.55, alpha=0.8)
        axis.tick_params(length=3)
    handles, labels = axes[0].get_legend_handles_labels()
    figure.legend(
        handles,
        labels,
        loc="upper center",
        ncol=5,
        frameon=False,
        bbox_to_anchor=(0.5, 1.02),
    )
    task_count = analysis["bootstrap"]["task_clusters"]
    component_draws = analysis["bootstrap"]["component_draws_per_task_size"]
    evidence_label = "Confirmation" if analysis["confirmatory"] else "Exploratory pilot"
    figure.text(
        0.995,
        0.005,
        f"{evidence_label}: {task_count} tasks, {component_draws} component "
        "draws per task and size",
        ha="right",
        va="bottom",
        fontsize=7,
        color="#666666",
    )
    figure.tight_layout(rect=(0.0, 0.04, 1.0, 0.92), 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()
    source_paths = list(args.core) + list(args.baselines)
    reports = [json.loads(path.read_text()) for path in source_paths]
    records = merge_reports(reports)
    if args.confirmatory:
        validate_core_confirmation(reports, records)
    analysis = build_analysis(
        records,
        replicates=args.bootstrap_replicates,
        seed=args.bootstrap_seed,
        confirmatory=args.confirmatory,
    )
    analysis["sources"] = [str(path) for path in source_paths]
    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_figure(args.output_figure, analysis)
    print(json.dumps({
        "excess_error_scaling": analysis["excess_error_scaling"],
        "largest_side": analysis["sizes"][-1],
        "largest_side_methods": analysis["summaries"][
            str(analysis["sizes"][-1])
        ]["methods"],
    }, 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()