scripts/make_release_figures.py

"""Generate the minimal release figure set.

Figures are deterministic and use only existing generated CSV outputs. They do
not modify or rerun the AX-CPT experiment scripts.
"""

from __future__ import annotations

import csv
from collections import defaultdict
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.patches import FancyBboxPatch
import numpy as np


TRIAL_TYPES = ["AX", "AY", "BX", "BY"]
FIGURE_DIR = Path("figures")


def read_csv(path: Path) -> list[dict[str, str]]:
    with path.open("r", newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def write_caption(path: Path, text: str) -> None:
    path.write_text(text.strip() + "\n", encoding="utf-8")


def set_style() -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 9,
            "axes.titlesize": 11,
            "axes.labelsize": 9,
            "figure.dpi": 160,
            "savefig.dpi": 300,
            "axes.spines.top": False,
            "axes.spines.right": False,
        }
    )


def save_figure(fig: plt.Figure, name: str) -> None:
    FIGURE_DIR.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIGURE_DIR / f"{name}.png", bbox_inches="tight")
    fig.savefig(FIGURE_DIR / f"{name}.svg", bbox_inches="tight")
    plt.close(fig)


def figure_1_task_schematic() -> None:
    fig, ax = plt.subplots(figsize=(7.4, 4.8))
    ax.set_axis_off()
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)

    ax.text(0.06, 0.965, "AX-CPT task logic", fontsize=15, weight="bold", ha="left", va="top")
    ax.text(
        0.06,
        0.865,
        "Respond TARGET only when cue A is followed by probe X.",
        fontsize=10.5,
        ha="left",
        color="#333333",
    )

    x0, y0 = 0.09, 0.735
    row_h = 0.108
    cols = [0.0, 0.19, 0.38, 0.64]
    headers = ["Cue", "Probe", "Trial", "Correct response"]
    for idx, header in enumerate(headers):
        ax.text(x0 + cols[idx], y0 + 0.06, header, fontsize=9.5, weight="bold", ha="left", va="center")

    rows = [
        ("A", "X", "AX", "TARGET", True),
        ("A", "Y", "AY", "NONTARGET", False),
        ("B", "X", "BX", "NONTARGET", False),
        ("B", "Y", "BY", "NONTARGET", False),
    ]
    for row_idx, (cue, probe, trial, response, is_target) in enumerate(rows):
        y = y0 - (row_idx + 1) * row_h
        fill = "#e8f3ee" if is_target else "#f7f7f7"
        edge = "#26734d" if is_target else "#d0d0d0"
        rect = FancyBboxPatch(
            (x0 - 0.025, y - 0.038),
            0.80,
            0.076,
            boxstyle="round,pad=0.008,rounding_size=0.01",
            linewidth=1.0,
            edgecolor=edge,
            facecolor=fill,
        )
        ax.add_patch(rect)
        ax.text(x0 + cols[0], y, cue, fontsize=13, weight="bold", ha="left", va="center")
        ax.text(x0 + cols[1], y, probe, fontsize=13, weight="bold", ha="left", va="center")
        ax.text(x0 + cols[2], y, trial, fontsize=12, ha="left", va="center")
        ax.text(
            x0 + cols[3],
            y,
            response,
            fontsize=11,
            weight="bold" if is_target else "normal",
            color="#17643d" if is_target else "#333333",
            ha="left",
            va="center",
        )

    note = (
        "v5 implicit-learning runs use token remapping:\n"
        "A/B -> miv/pel, X/Y -> dor/kex,\n"
        "TARGET/NONTARGET -> J/F."
    )
    note_box = FancyBboxPatch(
        (0.59, 0.055),
        0.32,
        0.155,
        boxstyle="round,pad=0.018,rounding_size=0.012",
        linewidth=0.8,
        edgecolor="#c5c5c5",
        facecolor="#fbfbfb",
    )
    ax.add_patch(note_box)
    ax.text(0.61, 0.185, "Small v5 note", fontsize=8.0, weight="bold", ha="left", va="top", color="#555555")
    ax.text(0.61, 0.153, note, fontsize=7.4, ha="left", va="top", linespacing=1.25, color="#555555")

    ax.text(
        0.06,
        0.075,
        "Schematic only; summarizes the task rule.\nThe token mapping note is shown at right.",
        fontsize=8,
        color="#555555",
        ha="left",
        va="top",
        linespacing=1.25,
    )

    save_figure(fig, "figure_1_axcpt_task_schematic")
    write_caption(
        FIGURE_DIR / "figure_1_axcpt_task_schematic_caption.txt",
        "Figure 1. AX-CPT task schematic. The task rule is TARGET only for A-X and NONTARGET otherwise; the side note shows the v5 token remapping used in the implicit-learning runs.",
    )


def transition_group(condition: str, dataset: str) -> str:
    if dataset == "axcpt_v4b":
        return "v4b"
    if condition.endswith("_BASE"):
        return "v5 BASE"
    if condition.endswith("_DCM"):
        return "v5 DCM"
    return "other"


def mean(values: list[float]) -> float:
    return sum(values) / len(values) if values else float("nan")


def transition_matrix(rows: list[dict[str, str]], group_name: str) -> np.ndarray:
    buckets: dict[tuple[str, str], list[float]] = defaultdict(list)
    for row in rows:
        if group_name != "all" and row["plot_group"] != group_name:
            continue
        value = row["error_delta_vs_current_type_baseline"]
        if value == "":
            continue
        buckets[(row["previous_type"], row["current_type"])].append(float(value))

    matrix = np.full((len(TRIAL_TYPES), len(TRIAL_TYPES)), np.nan, dtype=float)
    for i, prev_type in enumerate(TRIAL_TYPES):
        for j, current_type in enumerate(TRIAL_TYPES):
            values = buckets[(prev_type, current_type)]
            if values:
                matrix[i, j] = mean(values)
    return matrix


def figure_2_transition_summary() -> None:
    rows = read_csv(Path("outputs/accuracy_transition_summary.csv"))
    for row in rows:
        row["plot_group"] = transition_group(row["condition"], row["dataset"])

    groups = ["v4b", "v5 BASE", "v5 DCM", "all"]
    matrices = [transition_matrix(rows, group) for group in groups]
    vmax = max(0.05, max(float(np.nanmax(np.abs(matrix))) for matrix in matrices))
    cmap = plt.get_cmap("RdBu_r").copy()
    cmap.set_bad("#f0f0f0")

    fig, axes = plt.subplots(1, 4, figsize=(9.8, 2.9), sharex=True, sharey=True)
    im = None
    for ax, group, matrix in zip(axes, groups, matrices):
        im = ax.imshow(np.ma.masked_invalid(matrix), cmap=cmap, vmin=-vmax, vmax=vmax)
        ax.set_title(group)
        ax.set_xticks(range(len(TRIAL_TYPES)), TRIAL_TYPES, rotation=0)
        ax.set_yticks(range(len(TRIAL_TYPES)), TRIAL_TYPES)
        ax.set_xlabel("Current")
        if ax is axes[0]:
            ax.set_ylabel("Previous")
        for i in range(len(TRIAL_TYPES)):
            for j in range(len(TRIAL_TYPES)):
                label = "—" if np.isnan(matrix[i, j]) else f"{matrix[i, j]:.2f}"
                ax.text(j, i, label, ha="center", va="center", fontsize=7)
        ax.tick_params(length=0)

    assert im is not None
    cbar = fig.colorbar(im, ax=axes, fraction=0.025, pad=0.025)
    cbar.set_label("Error-rate delta")
    fig.suptitle("Accuracy-based transition structure", y=1.04, fontsize=13, weight="bold")
    fig.text(
        0.5,
        -0.03,
        "Values are error-rate deltas from current-trial-type baselines. No reaction times are available.",
        ha="center",
        fontsize=8,
        color="#555555",
    )

    save_figure(fig, "figure_2_accuracy_transition_summary")
    write_caption(
        FIGURE_DIR / "figure_2_accuracy_transition_summary_caption.txt",
        "Figure 2. Accuracy-based transition structure. Heatmap values are error-rate deltas relative to the baseline for each current trial type, grouped compactly by v4b, v5 BASE, v5 DCM, and all conditions. This is not a reaction-time transition-cost plot.",
    )


def load_projection(path: Path, source: str) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    for row in read_csv(path):
        rows.append(
            {
                "source": source,
                "dataset": row["dataset"],
                "condition": row["condition"],
                "pc1": float(row["pc1"]),
                "pc2": float(row["pc2"]),
                "pc1_var": float(row["pc1_explained_variance_ratio"]),
                "pc2_var": float(row["pc2_explained_variance_ratio"]),
            }
        )
    return rows


def marker_for(condition: str) -> str:
    if condition.endswith("_BASE"):
        return "o"
    if condition.endswith("_DCM"):
        return "s"
    return "^"


def short_label(condition: str) -> str:
    return {
        "C1_classic_no_dist": "C1",
        "C2_classic_dist3": "C2",
        "C3_bxheavy_no_dist": "C3",
        "C4_bxheavy_dist3": "C4",
        "WINDOW_10_BASE": "10B",
        "WINDOW_10_DCM": "10D",
        "WINDOW_5_BASE": "5B",
        "WINDOW_5_DCM": "5D",
    }.get(condition, condition)


def label_offset(panel: str, condition: str) -> tuple[int, int]:
    base_offsets = {
        "C1_classic_no_dist": (-24, 9),
        "C2_classic_dist3": (-22, -15),
        "C3_bxheavy_no_dist": (8, 12),
        "C4_bxheavy_dist3": (8, -16),
        "WINDOW_10_BASE": (12, 10),
        "WINDOW_10_DCM": (12, -16),
        "WINDOW_5_BASE": (-30, 10),
        "WINDOW_5_DCM": (-30, -16),
    }
    strict_offsets = {
        "C1_classic_no_dist": (-28, 14),
        "C2_classic_dist3": (-28, -18),
        "C3_bxheavy_no_dist": (12, 16),
        "C4_bxheavy_dist3": (12, -20),
        "WINDOW_10_BASE": (16, 20),
        "WINDOW_10_DCM": (16, -24),
        "WINDOW_5_BASE": (-34, 20),
        "WINDOW_5_DCM": (-34, -24),
    }
    return (strict_offsets if panel == "Strict leakage-reduced" else base_offsets).get(condition, (8, 8))


def figure_3_embedding_comparison() -> None:
    initial = load_projection(Path("outputs/embedding_analysis/condition_embedding_projection_2d.csv"), "Initial")
    strict = load_projection(
        Path("outputs/embedding_analysis_strict_leakage_reduced/condition_embedding_projection_2d.csv"),
        "Strict leakage-reduced",
    )
    rows_by_panel = [initial, strict]

    x_values = [row["pc1"] for row in initial + strict]
    y_values = [row["pc2"] for row in initial + strict]
    x_pad = max(0.08, (max(x_values) - min(x_values)) * 0.18)
    y_pad = max(0.08, (max(y_values) - min(y_values)) * 0.22)
    xlim = (min(x_values) - x_pad, max(x_values) + x_pad)
    ylim = (min(y_values) - y_pad, max(y_values) + y_pad)

    colors = {"axcpt_v4b": "#2f6f9f", "axcpt_v5_dcm": "#c35a3a"}
    fig, axes = plt.subplots(1, 2, figsize=(9.8, 4.1), sharex=True, sharey=True)

    for ax, title, panel_rows in zip(axes, ["Initial embedding", "Strict leakage-reduced"], rows_by_panel):
        for row in panel_rows:
            color = colors[row["dataset"]]
            ax.scatter(
                row["pc1"],
                row["pc2"],
                s=56,
                marker=marker_for(row["condition"]),
                color=color,
                edgecolor="white",
                linewidth=0.8,
                zorder=3,
            )
            dx, dy = label_offset(title, row["condition"])
            ax.annotate(
                short_label(row["condition"]),
                xy=(row["pc1"], row["pc2"]),
                xytext=(dx, dy),
                textcoords="offset points",
                fontsize=7,
                ha="center",
                va="center",
                bbox={"boxstyle": "round,pad=0.15", "facecolor": "white", "edgecolor": "none", "alpha": 0.78},
                arrowprops={"arrowstyle": "-", "color": "#999999", "linewidth": 0.45, "shrinkA": 1, "shrinkB": 2},
                zorder=4,
            )
        ax.axhline(0, color="#dddddd", linewidth=0.8, zorder=1)
        ax.axvline(0, color="#dddddd", linewidth=0.8, zorder=1)
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        ax.set_title(title)
        ax.set_xlabel(f"PC1 ({panel_rows[0]['pc1_var']:.1%})")
        ax.set_ylabel(f"PC2 ({panel_rows[0]['pc2_var']:.1%})")

    legend_items = [
        Line2D([0], [0], marker="o", color="none", label="v5 BASE", markerfacecolor=colors["axcpt_v5_dcm"], markersize=7),
        Line2D([0], [0], marker="s", color="none", label="v5 DCM", markerfacecolor=colors["axcpt_v5_dcm"], markersize=7),
        Line2D([0], [0], marker="^", color="none", label="v4b", markerfacecolor=colors["axcpt_v4b"], markersize=7),
    ]
    axes[1].legend(handles=legend_items, loc="upper right", frameon=False, fontsize=8)

    fig.suptitle("Exploratory text-embedding projection", y=1.03, fontsize=13, weight="bold")
    fig.text(
        0.5,
        -0.025,
        "Strict masking removes dataset, condition, direct DCM fields, and context_window from embedded text. v4b/v5 separation remains; v5 BASE/DCM largely overlaps.",
        ha="center",
        fontsize=8,
        color="#555555",
    )

    save_figure(fig, "figure_3_embedding_initial_vs_strict")
    write_caption(
        FIGURE_DIR / "figure_3_embedding_initial_vs_strict_caption.txt",
        "Figure 3. Exploratory text-embedding projections before and after strict leakage reduction. The strict pass removes dataset labels, condition labels, direct DCM fields, and context_window from embedded text. v4b versus v5 separation remains visible in the remaining observable serialized fields, while v5 BASE and DCM points largely overlap. These are local hashed text embeddings, not latent model embeddings.",
    )


def main() -> int:
    set_style()
    figure_1_task_schematic()
    figure_2_transition_summary()
    figure_3_embedding_comparison()
    print(f"Wrote release figures to {FIGURE_DIR}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())