#!/usr/bin/env python3
"""Render the two three-panel composite figures for the focused AAAI story."""

from __future__ import annotations

import argparse
import csv
import math
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--depth-dir", type=Path, default=Path("outputs/aaai_depth_experiments"))
    parser.add_argument(
        "--cnn-dir", type=Path, default=Path("outputs/cnn_initialization_validation")
    )
    parser.add_argument(
        "--mnist-dir", type=Path, default=Path("outputs/real_data_validation_mnist")
    )
    parser.add_argument(
        "--null-dir", type=Path, default=Path("outputs/soft_erosion_theory_vs_empirical")
    )
    parser.add_argument("--outdir", type=Path, default=Path("outputs/aaai_main_figures"))
    return parser.parse_args()


def read_csv(path: Path) -> list[dict[str, str]]:
    if not path.exists():
        raise FileNotFoundError(path)
    with path.open() as handle:
        return list(csv.DictReader(handle))


def panel_label(ax: plt.Axes, label: str) -> None:
    ax.text(-0.13, 1.06, label, transform=ax.transAxes, fontsize=12, fontweight="bold")


def draw_schematic(ax: plt.Axes) -> None:
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis("off")
    rows = [(0.82, "BP", "#2f6f9f"), (0.50, "FA", "#c65f16"), (0.18, "DFA", "#2b8a3e")]
    x_positions = [0.18, 0.42, 0.66, 0.90]
    labels = ["$h_1$", "$h_2$", "$h_3$", "error"]
    for y, name, color in rows:
        ax.text(0.01, y, name, va="center", ha="left", fontweight="bold", color=color)
        for x, node in zip(x_positions, labels):
            ax.add_patch(plt.Circle((x, y), 0.035, facecolor="white", edgecolor=color, lw=1.5))
            ax.text(x, y - 0.075, node, ha="center", va="top", fontsize=8)
        if name in {"BP", "FA"}:
            arrow_label = r"$W^\top$" if name == "BP" else "$B$"
            for right, left in zip(x_positions[:0:-1], x_positions[-2::-1]):
                ax.annotate(
                    "",
                    xy=(left + 0.04, y),
                    xytext=(right - 0.04, y),
                    arrowprops=dict(arrowstyle="->", color=color, lw=1.5),
                )
                ax.text((left + right) / 2, y + 0.035, arrow_label, ha="center", fontsize=8)
        else:
            for left in x_positions[:-1]:
                ax.annotate(
                    "",
                    xy=(left + 0.04, y),
                    xytext=(x_positions[-1] - 0.04, y),
                    arrowprops=dict(arrowstyle="->", color=color, lw=1.2, alpha=0.9),
                )
            ax.text(0.56, y + 0.055, "direct random maps", ha="center", fontsize=8)
    ax.set_title("Fixed random feedback rules", fontsize=10)


def figure_one(args: argparse.Namespace) -> Path:
    init_rows = read_csv(args.depth_dir / "initialization_rows.csv")
    cnn_rows = read_csv(args.cnn_dir / "cnn_initialization_rows.csv")
    mnist_rows = read_csv(args.mnist_dir / "e0_mnist_rows.csv")
    theory_rows = read_csv(args.null_dir / "theory_gap_summary.csv")
    empirical_rows = read_csv(args.null_dir / "empirical_gap_summary.csv")

    fig, axes = plt.subplots(1, 3, figsize=(14.2, 4.1), dpi=220)
    draw_schematic(axes[0])
    panel_label(axes[0], "a")

    ax = axes[1]
    theory_by_width = {int(row["width"]): float(row["mean"]) for row in theory_rows}
    empirical_by_width = {int(row["width"]): float(row["mean"]) for row in empirical_rows}
    widths = sorted(set(theory_by_width) & set(empirical_by_width))
    ax.plot(
        widths,
        [theory_by_width[width] for width in widths],
        "o-",
        color="#8c6d31",
        label="random-direction deletion",
    )
    ax.plot(
        widths,
        [empirical_by_width[width] for width in widths],
        "s-",
        color="#c65f16",
        label="measured FA cost",
    )
    ax.set_yscale("log")
    ax.set_xlabel("width")
    ax.set_ylabel("finite-time loss gap")
    ax.grid(alpha=0.16, which="both")
    ax2 = ax.twinx()
    ax2.plot(widths, [1.0 / (width * width) for width in widths], "k--", lw=1.4, label="$1/D$")
    ax2.set_yscale("log")
    ax2.set_ylabel("best guaranteed squared alignment")
    handles1, labels1 = ax.get_legend_handles_labels()
    handles2, labels2 = ax2.get_legend_handles_labels()
    ax.legend(handles1 + handles2, labels1 + labels2, fontsize=7, loc="upper right")
    worst_ratio = max(theory_by_width[width] / empirical_by_width[width] for width in widths)
    ax.text(0.04, 0.05, f"null overestimate: up to {worst_ratio:.0f}×", transform=ax.transAxes, fontsize=8)
    ax.set_title("Matrix mismatch is not optimization cost", fontsize=10)
    panel_label(ax, "b")

    ax = axes[2]
    depths = sorted({int(row["depth"]) for row in init_rows})
    cmap = plt.cm.viridis
    depth_color = {depth: cmap(index / max(len(depths) - 1, 1)) for index, depth in enumerate(depths)}
    for row in init_rows:
        depth = int(row["depth"])
        marker = "o" if row["rule"] == "FA" else "s"
        ax.errorbar(
            float(row["predicted_deficit"]),
            float(row["empirical_deficit"]),
            yerr=2 * float(row["empirical_stderr"]),
            fmt=marker,
            ms=3.8,
            color=depth_color[depth],
            alpha=0.68,
            lw=0.7,
            capsize=1,
        )
    for row in mnist_rows:
        ax.plot(float(row["hidden_share"]), float(row["empirical_mean"]), "^", color="#7b6fa6", ms=4, alpha=0.65)
    for row in cnn_rows:
        marker = "P" if row["rule"] == "FA" else "X"
        ax.errorbar(
            float(row["predicted_deficit"]),
            float(row["empirical_deficit"]),
            yerr=2 * float(row["empirical_stderr"]),
            fmt=marker,
            color="#b23a48",
            ms=5,
            capsize=1,
        )
    values = [
        float(row[key])
        for row in init_rows
        for key in ("predicted_deficit", "empirical_deficit")
    ]
    values += [
        float(row[key])
        for row in cnn_rows
        for key in ("predicted_deficit", "empirical_deficit")
    ]
    values += [float(row[key]) for row in mnist_rows for key in ("hidden_share", "empirical_mean")]
    lo, hi = min(values), max(values)
    pad = 0.04 * (hi - lo + 1e-12)
    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], "k-", lw=1)
    legend_handles = [
        plt.Line2D([], [], marker="o", ls="", color="0.35", label="synthetic FA"),
        plt.Line2D([], [], marker="s", ls="", color="0.35", label="synthetic DFA"),
        plt.Line2D([], [], marker="^", ls="", color="#7b6fa6", label="MNIST MLP"),
        plt.Line2D([], [], marker="P", ls="", color="#b23a48", label="CE CNN FA/DFA"),
    ]
    shape_legend = ax.legend(handles=legend_handles, fontsize=7, loc="lower right")
    ax.add_artist(shape_legend)
    depth_handles = [
        plt.Line2D([], [], marker="o", ls="", color=depth_color[depth], label=f"depth {depth}")
        for depth in depths
    ]
    ax.legend(handles=depth_handles, fontsize=6, ncols=2, loc="upper left")
    ax.set_xlabel("exact expected initial deficit")
    ax.set_ylabel("measured mean deficit")
    ax.set_title("One calibration across rule, depth, and data", fontsize=10)
    ax.grid(alpha=0.16)
    panel_label(ax, "c")

    fig.tight_layout(w_pad=2.0)
    path = args.outdir / "figure1_mismatch_vs_initial_cost.png"
    fig.savefig(path, bbox_inches="tight")
    fig.savefig(path.with_suffix(".pdf"), bbox_inches="tight")
    plt.close(fig)
    return path


def scatter_predictor(
    ax: plt.Axes,
    rows: list[dict[str, str]],
    fields: list[tuple[str, str, str]],
    title: str,
) -> None:
    depths = sorted({int(row["depth"]) for row in rows})
    colors = {depth: plt.cm.viridis(i / max(len(depths) - 1, 1)) for i, depth in enumerate(depths)}
    values: list[float] = []
    for field, marker, label in fields:
        for row in rows:
            x = float(row[field])
            y = float(row["empirical_gap"])
            values.extend([x, y])
            ax.scatter(x, y, s=13, marker=marker, color=colors[int(row["depth"])], alpha=0.60)
        ax.scatter([], [], marker=marker, color="0.35", label=label)
    lo, hi = min(values), max(values)
    pad = 0.04 * (hi - lo + 1e-12)
    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], "k-", lw=1)
    ax.set_xlabel("predicted gap")
    ax.set_ylabel("measured gap")
    ax.set_title(title, fontsize=10)
    ax.grid(alpha=0.16)
    ax.legend(fontsize=7)


def figure_two(args: argparse.Namespace) -> Path:
    rows = read_csv(args.depth_dir / "finite_time_rows.csv")
    fig, axes = plt.subplots(1, 3, figsize=(14.0, 4.1), dpi=220)

    ax = axes[0]
    widths = sorted({int(row["width"]) for row in rows})
    markers = ["o", "s", "^"]
    for width, marker in zip(widths, markers):
        depths = sorted({int(row["depth"]) for row in rows if int(row["width"]) == width})
        means = []
        sems = []
        for depth in depths:
            cell = [
                row
                for row in rows
                if int(row["width"]) == width and int(row["depth"]) == depth
            ]
            values = np.asarray([float(row["empirical_gap"]) for row in cell])
            init_means = np.asarray(
                [
                    np.mean(
                        [
                            float(row["empirical_gap"])
                            for row in cell
                            if int(row["init_seed"]) == init_seed
                        ]
                    )
                    for init_seed in sorted({int(row["init_seed"]) for row in cell})
                ]
            )
            means.append(float(values.mean()))
            sems.append(float(init_means.std(ddof=1) / math.sqrt(len(init_means))))
        ax.errorbar(depths, means, yerr=sems, marker=marker, capsize=2, label=f"width {width}")
    ax.set_xlabel("hidden-layer depth")
    ax.set_ylabel("FA loss - BP loss")
    ax.set_title("Later cost changes with training dynamics", fontsize=10)
    ax.grid(alpha=0.16)
    ax.legend(fontsize=7)
    panel_label(ax, "a")

    scatter_predictor(axes[1], rows, [("fixed_gap", "^", "frozen $K(0)$")], "Frozen initialization misses drift")
    panel_label(axes[1], "b")

    scatter_predictor(
        axes[2],
        rows,
        [
            ("velocity_gap", "o", "linear velocity"),
            ("retangent_gap", "s", "early retangent"),
        ],
        "Two snapshots recover the later cost",
    )
    panel_label(axes[2], "c")

    predictor_legend = axes[2].get_legend()
    if predictor_legend is not None:
        axes[2].add_artist(predictor_legend)

    depth_handles = [
        plt.Line2D(
            [],
            [],
            marker="o",
            ls="",
            color=plt.cm.viridis(i / max(len(sorted({int(row['depth']) for row in rows})) - 1, 1)),
            label=f"depth {depth}",
        )
        for i, depth in enumerate(sorted({int(row["depth"]) for row in rows}))
    ]
    axes[2].legend(handles=depth_handles, fontsize=6, ncols=2, loc="lower right")

    fig.tight_layout(w_pad=2.0)
    path = args.outdir / "figure2_training_dynamics.png"
    fig.savefig(path, bbox_inches="tight")
    fig.savefig(path.with_suffix(".pdf"), bbox_inches="tight")
    plt.close(fig)
    return path


def main() -> None:
    args = parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)
    path1 = figure_one(args)
    path2 = figure_two(args)
    print(f"figure: {path1}")
    print(f"figure: {path2}")


if __name__ == "__main__":
    main()