#!/usr/bin/env python3
"""Analyze the paired overclamped CLLN scaling confirmation."""

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

from analyze_coupled_ladder_scaling import (
    bootstrap_mean,
    merge_reports,
    task_cluster_values,
    task_size_matrix,
)


METHOD_ORDER = ("overclamp_clean", "overclamp", "overclamp_sdil")
DISPLAY = {
    "overclamp_clean": "Clean overclamp",
    "overclamp": "Imperfect overclamp",
    "overclamp_sdil": "Overclamp + SDIL",
}
STYLE = {
    "overclamp_clean": dict(color="#222222", marker="^", linestyle=":"),
    "overclamp": dict(color="#E69F00", marker="s", linestyle="--"),
    "overclamp_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,
}
FROZEN_OVERCLAMP_TIMES = {
    4: 0.25,
    8: 0.25,
    12: 0.25,
    16: 2.5,
    24: 2.5,
    32: 2.5,
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sources", type=Path, nargs="+", required=True)
    parser.add_argument(
        "--output-analysis",
        type=Path,
        default=Path("results/coupled_ladder/p3_overclamp_scaling_analysis.json"),
    )
    parser.add_argument(
        "--output-csv",
        type=Path,
        default=Path("results/coupled_ladder/p3_overclamp_scaling_source.csv"),
    )
    parser.add_argument(
        "--output-figure",
        type=Path,
        default=Path("results/figs/figure_clln_overclamp_scaling_confirmation"),
    )
    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")
    return parser.parse_args()


def validate_confirmation(reports: list[dict], records: list[dict]) -> None:
    expected = {
        (side, task, seed)
        for side in FROZEN_SIZES
        for task in range(40)
        for seed in FROZEN_DEVICE_SEEDS
    }
    actual = {
        (record["side"], record["task_index"], record["device_seed"])
        for record in records
    }
    if actual != expected:
        raise ValueError(
            f"incomplete confirmation: {len(expected - actual)} missing, "
            f"{len(actual - expected)} extra cells"
        )
    required = set(METHOD_ORDER)
    observed_sizes = set()
    for record in records:
        if set(record["methods"]) != required:
            raise ValueError("every cell must contain the three overclamp methods")
        if any(value["status"] != "completed" for value in record["methods"].values()):
            raise ValueError("all confirmatory methods must complete")
    for report in reports:
        protocol = report["protocol"]
        if not report.get("confirmatory", False):
            raise ValueError("every source must be labeled confirmatory")
        if protocol["rotations_per_input_diameter"] != 8:
            raise ValueError("confirmation requires all eight task rotations")
        if protocol["epochs"] != 600:
            raise ValueError("confirmation requires 600 epochs")
        if tuple(protocol["device_seeds"]) != FROZEN_DEVICE_SEEDS:
            raise ValueError("confirmation device seeds do not match")
        if set(protocol["methods"]) != required:
            raise ValueError("confirmation methods do not match")
        for side in protocol["sizes"]:
            observed_sizes.add(side)
            learning_time = float(
                protocol["learning_time_seconds_by_side"][str(side)]
            )
            overclamp_time = float(
                protocol["overclamp_time_seconds_per_v_by_side"][str(side)]
            )
            if abs(learning_time - FROZEN_LEARNING_TIMES[side]) > 1e-15:
                raise ValueError(f"side {side} uses an unfrozen learning exposure")
            if abs(overclamp_time - FROZEN_OVERCLAMP_TIMES[side]) > 1e-15:
                raise ValueError(f"side {side} uses an unfrozen overclamp exposure")
    if observed_sizes != set(FROZEN_SIZES):
        raise ValueError("sources do not cover all frozen sizes")


def slope_bootstrap(
    records: list[dict],
    sizes: list[int],
    edges: np.ndarray,
    metric: str,
    replicates: int,
    seed: int,
) -> dict:
    matrices = {
        method: task_size_matrix(records, sizes, method, metric)[1]
        for method in METHOD_ORDER
    }
    if metric == "reached_stable_zero_error":
        matrices = {method: 1.0 - values for method, values in matrices.items()}
    x = np.log10(edges)

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

    raw_excess = matrices["overclamp"] - matrices["overclamp_clean"]
    sdil_excess = matrices["overclamp_sdil"] - matrices["overclamp_clean"]
    point_raw = slope(raw_excess)
    point_sdil = slope(sdil_excess)
    point_difference = point_raw - point_sdil
    point_reduction = (
        100.0 * point_difference / point_raw if point_raw > 0.0 else None
    )
    rng = np.random.default_rng(seed)
    raw_slopes = []
    sdil_slopes = []
    differences = []
    reductions = []
    for _ in range(replicates):
        sample = rng.integers(0, len(raw_excess), size=len(raw_excess))
        raw_slope = slope(raw_excess[sample])
        sdil_slope = slope(sdil_excess[sample])
        raw_slopes.append(raw_slope)
        sdil_slopes.append(sdil_slope)
        differences.append(raw_slope - sdil_slope)
        if raw_slope > 1e-12:
            reductions.append(100.0 * (raw_slope - sdil_slope) / raw_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",
        "metric": metric,
        "imperfect_overclamp_excess_slope": point_raw,
        "imperfect_overclamp_excess_slope_95ci": interval(raw_slopes),
        "overclamp_sdil_excess_slope": point_sdil,
        "overclamp_sdil_excess_slope_95ci": interval(sdil_slopes),
        "paired_slope_difference": point_difference,
        "paired_slope_difference_95ci": interval(differences),
        "relative_slope_reduction_percent": point_reduction,
        "relative_slope_reduction_95ci": interval(reductions) if reductions else None,
    }


def build_analysis(
    records: list[dict], replicates: int, seed: int, confirmatory: bool
) -> dict:
    sizes = sorted({record["side"] for record in records})
    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:
        methods = {}
        for method in METHOD_ORDER:
            methods[method] = {}
            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)
                methods[method][metric] = {
                    "mean": mean,
                    "task_bootstrap_95ci": interval,
                }
        clean = methods["overclamp_clean"]["classification_error"]["mean"]
        raw = methods["overclamp"]["classification_error"]["mean"]
        sdil = methods["overclamp_sdil"]["classification_error"]["mean"]
        summaries[str(side)] = {
            "learnable_edges": int(edges[sizes.index(side)]),
            "methods": methods,
            "raw_to_clean_gap_closed": (
                float((raw - sdil) / (raw - clean))
                if abs(raw - clean) > 1e-15
                else None
            ),
        }
    return {
        "analysis": "overclamped_digital_coupled_ladder_scaling",
        "confirmatory": confirmatory,
        "bootstrap": {
            "unit": "task; three component draws averaged within task",
            "task_clusters": len({record["task_index"] for record in records}),
            "component_draws_per_task_size": len({
                record["device_seed"] for record in records
            }),
            "replicates": replicates,
            "seed": seed,
            "interval": "percentile 95%",
        },
        "sizes": sizes,
        "summaries": summaries,
        "excess_error_scaling": slope_bootstrap(
            records, sizes, edges, "classification_error", replicates, seed + 1
        ),
        "excess_error_auc_scaling": slope_bootstrap(
            records, sizes, edges, "classification_error_auc", replicates, seed + 2
        ),
        "excess_stable_failure_scaling": slope_bootstrap(
            records,
            sizes,
            edges,
            "reached_stable_zero_error",
            replicates,
            seed + 3,
        ),
    }


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, summary in analysis["summaries"].items():
            for method in METHOD_ORDER:
                for metric, values in summary["methods"][method].items():
                    writer.writerow((
                        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": 8,
        "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
    ])
    edge_labels = [
        f"{value / 1000.0:.1f}k" if value >= 1000 else str(value)
        for value in edges
    ]
    panels = (
        ("classification_error", "Final classification error (%)", 100.0),
        ("classification_error_auc", "Classification-error AUC", 1.0),
        ("reached_stable_zero_error", "Stable failure fraction (%)", 100.0),
    )
    slope_keys = (
        "excess_error_scaling",
        "excess_error_auc_scaling",
        "excess_stable_failure_scaling",
    )
    figure, axes = plt.subplots(1, 3, figsize=(10.6, 3.05))
    for panel_index, (metric, ylabel, scale) in enumerate(panels):
        axis = axes[panel_index]
        for method in METHOD_ORDER:
            intervals = np.asarray([
                analysis["summaries"][str(side)]["methods"][method][metric][
                    "task_bootstrap_95ci"
                ]
                for side in sizes
            ])
            means = np.asarray([
                analysis["summaries"][str(side)]["methods"][method][metric]["mean"]
                for side in sizes
            ])
            if metric == "reached_stable_zero_error":
                means = 1.0 - means
                intervals = np.column_stack((
                    1.0 - intervals[:, 1], 1.0 - intervals[:, 0]
                ))
            means *= scale
            intervals *= scale
            axis.errorbar(
                edges,
                means,
                yerr=np.vstack((means - intervals[:, 0], intervals[:, 1] - means)),
                linewidth=1.7 if method == "overclamp_sdil" else 1.15,
                markersize=4.5,
                capsize=2.0,
                label=DISPLAY[method],
                **STYLE[method],
            )
        scaling = analysis[slope_keys[panel_index]]
        reduction = scaling["relative_slope_reduction_percent"]
        resolved = scaling["paired_slope_difference_95ci"][0] > 0.0
        title = f"({chr(97 + panel_index)}) {ylabel.split(' (')[0]}"
        if reduction is not None and resolved:
            title += f"\n{reduction:.0f}% lower growth slope"
        else:
            title += "\nNo resolved slope reduction"
        axis.set_title(title)
        axis.set_xscale("log", base=2)
        axis.set_xticks(edges, edge_labels)
        axis.set_xlabel("Learnable edges")
        axis.set_ylabel(ylabel)
        axis.grid(axis="y", color="#D9D9D9", linewidth=0.55, alpha=0.8)
    handles, labels = axes[0].get_legend_handles_labels()
    figure.legend(
        handles, labels, loc="upper center", ncol=3, frameon=False,
        bbox_to_anchor=(0.5, 1.02),
    )
    evidence = "Confirmation" if analysis["confirmatory"] else "Exploratory"
    figure.text(
        0.995,
        0.005,
        f"{evidence}: {analysis['bootstrap']['task_clusters']} tasks, "
        f"{analysis['bootstrap']['component_draws_per_task_size']} component draws",
        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()
    reports = [json.loads(path.read_text()) for path in args.sources]
    records = merge_reports(reports)
    if args.confirmatory:
        validate_confirmation(reports, records)
    analysis = build_analysis(
        records, args.bootstrap_replicates, args.bootstrap_seed, args.confirmatory
    )
    analysis["sources"] = [str(path) for path in args.sources]
    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({key: analysis[key] for key in (
        "excess_error_scaling",
        "excess_error_auc_scaling",
        "excess_stable_failure_scaling",
    )}, indent=2))


if __name__ == "__main__":
    main()