#!/usr/bin/env python3
"""Plot trajectory-level capacity transition at two training horizons."""

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


RUNS = [
    (
        "T=10000, 1024 FA trajectories",
        Path("outputs/phase_transition_zoom_T10000_1024traj"),
        "#c65f16",
    ),
    (
        "T=30000, 256 FA trajectories",
        Path("outputs/phase_transition_zoom_T30000_256traj"),
        "#174f91",
    ),
]


def summarize_gaps(runs: pd.DataFrame) -> pd.DataFrame:
    fa = runs[runs["run_type"] == "fa"].copy()
    records = []
    for margin, group in fa.groupby("fa_capacity_margin"):
        gaps = group["train_gap_to_bp"].to_numpy(dtype=np.float64)
        records.append(
            {
                "margin": float(margin),
                "mean": float(gaps.mean()),
                "q05": float(np.quantile(gaps, 0.05)),
                "q25": float(np.quantile(gaps, 0.25)),
                "q75": float(np.quantile(gaps, 0.75)),
                "q95": float(np.quantile(gaps, 0.95)),
            }
        )
    return pd.DataFrame(records).sort_values("margin", ascending=False)


def add_summary(ax: plt.Axes, runs: pd.DataFrame, label: str, color: str) -> pd.DataFrame:
    summary = summarize_gaps(runs)
    x = summary["margin"].to_numpy(dtype=np.float64)
    ax.fill_between(
        x,
        summary["q05"],
        summary["q95"],
        color=color,
        alpha=0.13,
        linewidth=0,
    )
    ax.fill_between(
        x,
        summary["q25"],
        summary["q75"],
        color=color,
        alpha=0.24,
        linewidth=0,
    )
    ax.plot(
        x,
        summary["mean"],
        color=color,
        marker="o",
        linewidth=2.3,
        label=label,
    )
    return summary


def add_trajectories(ax: plt.Axes, runs: pd.DataFrame, color: str) -> None:
    fa = runs[runs["run_type"] == "fa"].copy()
    fa["_tid"] = fa["init_seed"].astype(str) + ":" + fa["feedback_seed"].astype(str)
    for _tid, group in fa.groupby("_tid"):
        group = group.sort_values("fa_capacity_margin", ascending=False)
        ax.plot(
            group["fa_capacity_margin"],
            group["train_gap_to_bp"],
            color=color,
            alpha=0.11,
            linewidth=0.85,
        )


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

    fig, axes = plt.subplots(1, 2, figsize=(13.2, 5.2), dpi=180)
    for ax, (label, run_dir, color) in zip(axes, RUNS, strict=True):
        runs = pd.read_csv(run_dir / "runs.csv")
        add_trajectories(ax, runs, color)
        summary = add_summary(ax, runs, label, color)
        ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
        ax.axhline(0, color="black", linewidth=1.0)
        ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
        ax.set_title(label)
        ax.grid(alpha=0.18)
        ax.legend(frameon=True)
        fa = runs[runs["run_type"] == "fa"]
        ymax = max(
            0.01,
            float(summary["q95"].max()) * 1.18,
            float(fa["train_gap_to_bp"].max()) * 1.03,
        )
        ax.set_ylim(-0.004 * ymax, ymax)
    axes[0].set_xlim(285, -280)
    axes[1].set_xlim(90, -150)
    for ax in axes:
        ax.set_ylabel("FA train MSE - BP train MSE")
    fig.suptitle("More trajectories separate finite-time ramp from capacity-limited gap")
    fig.tight_layout()
    path = outdir / "trajectory_gap_clouds_T10000_vs_T30000.png"
    fig.savefig(path)
    plt.close(fig)

    fig2, ax = plt.subplots(figsize=(8.8, 5.3), dpi=180)
    for label, run_dir, color in RUNS:
        runs = pd.read_csv(run_dir / "runs.csv")
        add_summary(ax, runs, label, color)
    ax.axvline(0, color="black", linestyle="--", linewidth=1.1, label="hard margin=0")
    ax.axhline(0, color="black", linewidth=1.0)
    ax.set_xlim(285, -280)
    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
    ax.set_ylabel("FA train MSE - BP train MSE")
    ax.set_title("Training horizon compresses the apparent transition")
    ax.grid(alpha=0.18)
    ax.legend(frameon=True)
    fig2.tight_layout()
    path2 = outdir / "trajectory_gap_quantiles_T10000_vs_T30000.png"
    fig2.savefig(path2)
    plt.close(fig2)

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


if __name__ == "__main__":
    main()