#!/usr/bin/env python3
"""Plot theory/empirical overlap panels for operator stress settings."""

from __future__ import annotations

import argparse
from pathlib import Path

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


SETTING_ORDER = [
    "d1_w64_T50_s20",
    "d2_w32_T50_s20",
    "d2_w64_T25_s10",
    "d2_w64_T50_s20_256traj",
    "d2_w64_T100_s40",
    "d2_w96_T50_s20",
    "d3_w64_T50_s20",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--rows",
        type=Path,
        default=Path("outputs/operator_stress_grid_summary/stress_grid_rows.csv"),
    )
    parser.add_argument(
        "--metrics",
        type=Path,
        default=Path("outputs/operator_stress_grid_summary/stress_grid_metrics.csv"),
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/operator_stress_grid_summary"),
    )
    return parser.parse_args()


def summarize_by_margin(group: pd.DataFrame, column: str) -> pd.DataFrame:
    records = []
    for margin, by_margin in group.groupby("hard_fa_capacity_margin"):
        values = by_margin[column].to_numpy(dtype=np.float64)
        records.append(
            {
                "margin": float(margin),
                "mean": float(values.mean()),
                "q05": float(np.quantile(values, 0.05)),
                "q25": float(np.quantile(values, 0.25)),
                "q75": float(np.quantile(values, 0.75)),
                "q95": float(np.quantile(values, 0.95)),
            }
        )
    return pd.DataFrame(records).sort_values("margin")


def trajectory_id(frame: pd.DataFrame) -> pd.Series:
    return frame["init_seed"].astype(str) + ":" + frame["feedback_seed"].astype(str)


def pretty_setting(setting: str) -> str:
    if setting == "d2_w64_T50_s20_256traj":
        return "d2 w64 T50 s20 (256 traj)"
    return setting.replace("_", " ")


def panel_ylim(group: pd.DataFrame) -> tuple[float, float]:
    values = np.concatenate(
        [
            group["empirical_gap"].to_numpy(dtype=np.float64),
            group["linear_gap"].to_numpy(dtype=np.float64),
        ]
    )
    lo = float(values.min())
    hi = float(values.max())
    pad = max((hi - lo) * 0.12, 0.004)
    return max(0.0, lo - pad), hi + pad


def plot_panel(
    ax: plt.Axes,
    group: pd.DataFrame,
    metrics: pd.Series,
    show_legend: bool,
) -> None:
    group = group.sort_values(["hard_fa_capacity_margin", "init_seed", "feedback_seed"])
    theory = summarize_by_margin(group, "linear_gap")
    empirical = summarize_by_margin(group, "empirical_gap")
    group = group.assign(_trajectory_id=trajectory_id(group))

    for _tid, traj in group.groupby("_trajectory_id"):
        traj = traj.sort_values("hard_fa_capacity_margin")
        ax.plot(
            traj["hard_fa_capacity_margin"],
            traj["empirical_gap"],
            color="#c65f16",
            alpha=0.16 if len(group) > 80 else 0.26,
            linewidth=1.05,
        )

    ax.fill_between(
        theory["margin"],
        theory["q05"],
        theory["q95"],
        color="#1f5a9d",
        alpha=0.14,
        linewidth=0,
        label="theory 5-95%",
    )
    ax.fill_between(
        theory["margin"],
        theory["q25"],
        theory["q75"],
        color="#1f5a9d",
        alpha=0.26,
        linewidth=0,
        label="theory 25-75%",
    )
    ax.plot(
        theory["margin"],
        theory["mean"],
        color="#174f91",
        marker="o",
        markersize=3.3,
        linewidth=2.0,
        label="theory mean",
    )
    ax.plot(
        empirical["margin"],
        empirical["mean"],
        color="#a84708",
        marker="o",
        markersize=3.0,
        linewidth=1.8,
        label="empirical mean",
    )
    ax.axvline(0.0, color="black", linestyle="--", linewidth=0.9)
    ax.set_ylim(*panel_ylim(group))
    ax.grid(alpha=0.17)
    ax.set_title(
        f"{pretty_setting(str(metrics.name))}\n"
        f"MAE={metrics['linear_mae']:.4g}, "
        f"bias={metrics['linear_bias']:.4g}, "
        f"corr={metrics['linear_corr']:.4f}",
        fontsize=10,
    )
    if show_legend:
        ax.legend(loc="upper left", fontsize=8, frameon=True)


def main() -> None:
    args = parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)
    rows = pd.read_csv(args.rows)
    metrics = pd.read_csv(args.metrics).set_index("setting")
    settings = [setting for setting in SETTING_ORDER if setting in set(rows["setting"])]

    fig, axes = plt.subplots(4, 2, figsize=(15.5, 18.5), dpi=180)
    axes_flat = axes.ravel()
    for index, setting in enumerate(settings):
        group = rows[rows["setting"] == setting]
        plot_panel(
            axes_flat[index],
            group,
            metrics.loc[setting],
            show_legend=index == 0,
        )
        axes_flat[index].set_xlabel("hard FA capacity margin")
        axes_flat[index].set_ylabel("FA train MSE - BP train MSE")

    for index in range(len(settings), len(axes_flat)):
        axes_flat[index].axis("off")

    fig.suptitle(
        "Early-kernel velocity theory vs empirical trajectory distributions",
        fontsize=18,
        y=0.995,
    )
    fig.tight_layout(rect=(0, 0, 1, 0.985))
    outpath = args.outdir / "stress_grid_overlap_panels.png"
    fig.savefig(outpath)
    plt.close(fig)
    print(f"plot: {outpath}")


if __name__ == "__main__":
    main()