#!/usr/bin/env python3
"""Two-panel teacher-task figure: train gap (soft ramp) vs test gap (sign flip).

Reads width_summary.csv from a teacher-task sweep produced by
downstream_capacity_sweep.py and plots, side by side:
  left  - FA/BP train gap vs width (log y): the optimization cost, soft ramp;
  right - FA/BP test gap vs width (linear y, error bars): task-dependent,
          crosses zero (FA can generalize better at intermediate width).
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

import matplotlib

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


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Teacher-task train vs test gap figure.")
    p.add_argument("--run-dir", type=Path,
                   default=Path("outputs/teacher_test_gap_sgd_T8000"))
    return p.parse_args()


def main() -> None:
    args = parse_args()
    rows = []
    with (args.run_dir / "width_summary.csv").open() as fh:
        for row in csv.DictReader(fh):
            rows.append({k: float(v) for k, v in row.items()})
    rows.sort(key=lambda r: r["width"])

    widths = np.array([r["width"] for r in rows])
    n_fa = np.array([r["runs_fa"] for r in rows])
    train_gap = np.array([r["train_gap_mean"] for r in rows])
    train_se = np.array([r["train_gap_std"] for r in rows]) / np.sqrt(n_fa)
    test_gap = np.array([r["gap_mean"] for r in rows])
    test_se = np.array([r["gap_std"] for r in rows]) / np.sqrt(n_fa)

    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.4), dpi=180)
    axes[0].errorbar(widths, train_gap, yerr=train_se, fmt="o-", color="#c65f16",
                     capsize=2.5)
    axes[0].set_xscale("log", base=2)
    axes[0].set_yscale("log")
    axes[0].set_xlabel("hidden width")
    axes[0].set_ylabel("FA train MSE - BP train MSE")
    axes[0].set_title("Optimization gap: soft, monotone, never zero")
    axes[0].grid(alpha=0.18, which="both")

    axes[1].errorbar(widths, test_gap, yerr=test_se, fmt="o-", color="#174f91",
                     capsize=2.5)
    axes[1].axhline(0.0, color="black", linewidth=1)
    axes[1].set_xscale("log", base=2)
    axes[1].set_xlabel("hidden width")
    axes[1].set_ylabel("FA test MSE - BP test MSE")
    axes[1].set_title("Generalization gap: task-dependent, changes sign")
    axes[1].grid(alpha=0.18)

    fig.suptitle("Teacher task (MLP teacher, SGD, T=8000): optimization cost vs test effect",
                 y=1.02)
    fig.tight_layout()
    out = args.run_dir / "teacher_train_vs_test_gap.png"
    fig.savefig(out, bbox_inches="tight")
    print(f"plot: {out}")


if __name__ == "__main__":
    main()