#!/usr/bin/env python3
"""Plot the P5 physical-grid classification crossover from committed JSON."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


METHOD_STYLE = {
    "raw": {
        "label": "Raw local update",
        "color": "#666666",
        "marker": "x",
        "linestyle": ":",
    },
    "constant": {
        "label": "Constant calibration",
        "color": "#E69F00",
        "marker": "^",
        "linestyle": "--",
    },
    "sdil": {
        "label": "SDIL",
        "color": "#0072B2",
        "marker": "o",
        "linestyle": "-",
    },
    "overclamp": {
        "label": "Overclamping",
        "color": "#222222",
        "marker": "D",
        "linestyle": "--",
    },
    "overclamp_sdil": {
        "label": "Overclamping + SDIL",
        "color": "#009E73",
        "marker": "s",
        "linestyle": "-",
    },
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input", type=Path,
        default=Path("results/physical_bias/p5_grid_bias_crossover.json"))
    parser.add_argument(
        "--output-prefix", type=Path,
        default=Path("results/figs/physical_grid_p5"))
    return parser.parse_args()


def values_by_diameter(records: list[dict], method: str) -> tuple[np.ndarray, list[np.ndarray]]:
    diameters = np.asarray(sorted({
        record["input_diameter_v"] for record in records
    }))
    values = []
    for diameter in diameters:
        values.append(np.asarray([
            100.0 * record["methods"][method]["classification_error"]
            for record in records
            if record["input_diameter_v"] == diameter
        ]))
    return 1_000.0 * diameters, values


def draw_panel(
    axis: plt.Axes,
    records: list[dict],
    methods: tuple[str, ...],
    title: str,
) -> None:
    offsets = np.linspace(-3.0, 3.0, len(methods))
    for method_index, method in enumerate(methods):
        diameters, values = values_by_diameter(records, method)
        style = METHOD_STYLE[method]
        for diameter, trial_values in zip(diameters, values):
            trial_jitter = np.linspace(-1.5, 1.5, len(trial_values))
            axis.scatter(
                diameter + offsets[method_index] + trial_jitter,
                trial_values,
                color=style["color"],
                marker=style["marker"],
                s=18,
                alpha=0.28,
                linewidths=0.8,
                zorder=2,
            )
        means = np.asarray([np.mean(trials) for trials in values])
        axis.plot(
            diameters,
            means,
            label=style["label"],
            color=style["color"],
            marker=style["marker"],
            linestyle=style["linestyle"],
            linewidth=1.8,
            markersize=5.5,
            markerfacecolor=(
                "white" if method in {"sdil", "overclamp_sdil"}
                else style["color"]),
            markeredgewidth=1.2,
            zorder=3,
        )
    axis.set_title(title, loc="left", fontsize=10, fontweight="bold")
    axis.set_xlabel("Input diameter (mV)")
    axis.set_xticks(diameters)
    axis.set_ylim(-2.0, 80.0)
    axis.set_yticks((0, 20, 40, 60, 80))
    axis.grid(axis="y", color="#D9D9D9", linewidth=0.7, zorder=0)
    axis.spines[["top", "right"]].set_visible(False)
    axis.legend(frameon=False, fontsize=8.5, loc="upper right")


def main() -> None:
    args = parse_args()
    report = json.loads(args.input.read_text())
    records = report["records"]
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 9,
        "axes.labelsize": 9,
        "xtick.labelsize": 8,
        "ytick.labelsize": 8,
        "svg.fonttype": "none",
        "pdf.fonttype": 42,
    })
    figure, axes = plt.subplots(
        1, 2, figsize=(7.1, 2.75), sharey=True, constrained_layout=True)
    draw_panel(
        axes[0], records, ("raw", "constant", "sdil"),
        "(a) Standard local learning")
    draw_panel(
        axes[1], records, ("overclamp", "overclamp_sdil"),
        "(b) Composition with overclamping")
    axes[0].set_ylabel("Classification error (%)")
    axes[1].text(
        0.98, 0.82, "8 trials per input diameter",
        transform=axes[1].transAxes, ha="right", va="top",
        fontsize=8, color="#666666")
    args.output_prefix.parent.mkdir(parents=True, exist_ok=True)
    figure.savefig(args.output_prefix.with_suffix(".svg"))
    figure.savefig(args.output_prefix.with_suffix(".pdf"))
    figure.savefig(args.output_prefix.with_suffix(".png"), dpi=300)
    plt.close(figure)
    print(f"wrote {args.output_prefix}.{{svg,pdf,png}}")


if __name__ == "__main__":
    main()