#!/usr/bin/env python3
"""Compare capacity-transition curves at different SGD horizons."""

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd


RUNS = [
    ("T=3000, 256 traj", Path("outputs/phase_transition_sgd_256_lr001_T3000")),
    ("T=10000 probe", Path("outputs/phase_transition_sgd_long_lr001_T10000_probe")),
    ("T=30000 probe", Path("outputs/phase_transition_sgd_verylong_lr001_T30000_probe")),
]


def main() -> None:
    outdir = Path("outputs/phase_transition_training_time_comparison")
    outdir.mkdir(parents=True, exist_ok=True)

    fig, ax = plt.subplots(figsize=(9.2, 5.7), dpi=180)
    for label, run_dir in RUNS:
        df = pd.read_csv(run_dir / "width_summary.csv").sort_values(
            "fa_capacity_margin", ascending=False
        )
        ax.plot(
            df["fa_capacity_margin"],
            df["train_gap_mean"],
            marker="o",
            linewidth=2.0,
            label=label,
        )
    ax.axvline(0, color="black", linestyle="--", linewidth=1.1, label="hard FA margin=0")
    ax.axhline(0, color="black", linewidth=1.0)
    ax.set_xlim(1100, -430)
    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out (capacity decreases ->)")
    ax.set_ylabel("FA train MSE - BP train MSE")
    ax.set_title("Training time shifts the apparent FA capacity transition")
    ax.legend(frameon=True)
    ax.grid(alpha=0.18)
    fig.tight_layout()
    path = outdir / "train_gap_vs_margin_by_training_time.png"
    fig.savefig(path)
    plt.close(fig)

    fig2, ax2 = plt.subplots(figsize=(9.2, 5.7), dpi=180)
    for label, run_dir in RUNS:
        df = pd.read_csv(run_dir / "width_summary.csv").sort_values(
            "fa_capacity_margin", ascending=False
        )
        ax2.plot(
            df["fa_capacity_margin"],
            df["bp_train_mean"],
            marker="o",
            linewidth=1.8,
            linestyle="--",
            label=f"BP {label}",
        )
        ax2.plot(
            df["fa_capacity_margin"],
            df["fa_train_mean"],
            marker="o",
            linewidth=2.0,
            label=f"FA {label}",
        )
    ax2.axvline(0, color="black", linestyle="--", linewidth=1.1)
    ax2.set_xlim(1100, -430)
    ax2.set_yscale("log")
    ax2.set_xlabel("hard FA capacity margin P - K_FA - N*out (capacity decreases ->)")
    ax2.set_ylabel("train MSE")
    ax2.set_title("BP/FA train losses across training time")
    ax2.legend(frameon=True, fontsize=8, ncol=2)
    ax2.grid(alpha=0.18)
    fig2.tight_layout()
    path2 = outdir / "bp_fa_train_loss_by_training_time.png"
    fig2.savefig(path2)
    plt.close(fig2)

    print(f"plot: {path}")
    print(f"plot: {path2}")


if __name__ == "__main__":
    main()