#!/usr/bin/env python3
"""Plot early-kernel finite-time FA/BP gap predictions."""

from __future__ import annotations

import argparse
import csv
from dataclasses import dataclass
from pathlib import Path

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


@dataclass(frozen=True)
class MetricsRow:
    predictor: str
    rows: int
    mae: float
    bias: float
    corr: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--glob",
        default="outputs/compressed_operator_s20_256traj_T50_width64_N*/compressed_operator_rows.csv",
    )
    parser.add_argument("--input-dim", type=int, default=16)
    parser.add_argument("--output-dim", type=int, default=4)
    parser.add_argument("--width", type=int, default=64)
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/compressed_operator_s20_256traj_T50_width64_plots"),
    )
    return parser.parse_args()


def read_rows(pattern: str, input_dim: int, output_dim: int, width: int) -> pd.DataFrame:
    frames = []
    for path in sorted(Path().glob(pattern)):
        dirname = path.parent.name
        try:
            train_samples = int(dirname.rsplit("N", 1)[1])
        except (IndexError, ValueError) as exc:
            raise ValueError(f"Cannot infer train sample count from {dirname}") from exc
        frame = pd.read_csv(path)
        frame["train_samples"] = train_samples
        frames.append(frame)
    if not frames:
        raise FileNotFoundError(f"No rows found for pattern: {pattern}")
    df = pd.concat(frames, ignore_index=True)
    if "width" in df.columns:
        row_width = df["width"].astype(int)
    else:
        row_width = pd.Series(width, index=df.index)
    if "hidden_layers" in df.columns:
        hidden_layers = df["hidden_layers"].astype(int)
    else:
        hidden_layers = pd.Series(2, index=df.index)
    parameter_count = (
        input_dim * row_width
        + (hidden_layers - 1).clip(lower=0) * row_width * row_width
        + row_width * output_dim
    )
    constraint_rank = (
        (hidden_layers - 1).clip(lower=0) * np.maximum(row_width * row_width - 1, 0)
        + np.maximum(row_width * output_dim - 1, 0)
    )
    df["hard_fa_capacity_margin"] = (
        parameter_count - constraint_rank - output_dim * df["train_samples"]
    )
    df["trajectory_id"] = (
        df["init_seed"].astype(str) + ":" + df["feedback_seed"].astype(str)
    )
    return df.sort_values(["hard_fa_capacity_margin", "trajectory_id"]).reset_index(drop=True)


def predictor_metrics(df: pd.DataFrame) -> list[MetricsRow]:
    rows: list[MetricsRow] = []
    for predictor, column in [
        ("fixed K(0)", "fixed_gap"),
        ("FA compressed", "compressed_gap"),
        ("linear velocity", "linear_gap"),
        ("early re-tangent", "retangent_gap"),
    ]:
        error = df[column].to_numpy() - df["empirical_gap"].to_numpy()
        corr = float(np.corrcoef(df[column], df["empirical_gap"])[0, 1])
        rows.append(
            MetricsRow(
                predictor=predictor,
                rows=len(df),
                mae=float(np.mean(np.abs(error))),
                bias=float(np.mean(error)),
                corr=corr,
            )
        )
    return rows


def write_metrics(path: Path, rows: list[MetricsRow]) -> None:
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(MetricsRow.__annotations__.keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(row.__dict__)


def summarize_by_margin(df: pd.DataFrame, column: str) -> pd.DataFrame:
    records = []
    for margin, group in df.groupby("hard_fa_capacity_margin"):
        values = group[column].to_numpy()
        records.append(
            {
                "margin": 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 plot_theory_band_with_empirical(df: pd.DataFrame, outdir: Path) -> Path:
    theory = summarize_by_margin(df, "linear_gap")
    empirical = summarize_by_margin(df, "empirical_gap")
    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)

    for _tid, group in df.groupby("trajectory_id"):
        group = group.sort_values("hard_fa_capacity_margin")
        ax.plot(
            group["hard_fa_capacity_margin"],
            group["empirical_gap"],
            color="#c65f16",
            alpha=0.16,
            linewidth=1.0,
        )

    ax.fill_between(
        theory["margin"],
        theory["q05"],
        theory["q95"],
        color="#1f5a9d",
        alpha=0.16,
        linewidth=0,
        label="velocity theory 5-95%",
    )
    ax.fill_between(
        theory["margin"],
        theory["q25"],
        theory["q75"],
        color="#1f5a9d",
        alpha=0.28,
        linewidth=0,
        label="velocity theory 25-75%",
    )
    ax.plot(
        theory["margin"],
        theory["mean"],
        color="#174f91",
        linewidth=2.3,
        marker="o",
        markersize=4.5,
        label="velocity theory mean",
    )
    ax.plot(
        empirical["margin"],
        empirical["mean"],
        color="#a84708",
        linewidth=2.1,
        marker="o",
        markersize=4.2,
        label="empirical mean",
    )
    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
    ax.set_title("Early-kernel velocity prediction vs empirical FA/BP gap")
    ax.set_xlabel("hard FA capacity margin")
    ax.set_ylabel("FA train MSE - BP train MSE")
    ax.legend(loc="upper left", frameon=True)
    ax.grid(alpha=0.18)
    fig.tight_layout()
    outpath = outdir / "T50_s20_velocity_theory_band_empirical_trajectories.png"
    fig.savefig(outpath)
    plt.close(fig)
    return outpath


def plot_empirical_only(df: pd.DataFrame, outdir: Path) -> Path:
    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)
    for _tid, group in df.groupby("trajectory_id"):
        group = group.sort_values("hard_fa_capacity_margin")
        ax.plot(
            group["hard_fa_capacity_margin"],
            group["empirical_gap"],
            color="#c65f16",
            alpha=0.20,
            linewidth=1.05,
        )
    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
    ax.set_title("Empirical FA/BP gap trajectories")
    ax.set_xlabel("hard FA capacity margin")
    ax.set_ylabel("FA train MSE - BP train MSE")
    ax.grid(alpha=0.18)
    fig.tight_layout()
    outpath = outdir / "T50_s20_empirical_trajectories_only.png"
    fig.savefig(outpath)
    plt.close(fig)
    return outpath


def plot_prediction_scatter(df: pd.DataFrame, outdir: Path) -> Path:
    fig, ax = plt.subplots(figsize=(7.0, 6.2), dpi=180)
    empirical = df["empirical_gap"].to_numpy()
    predictors = [
        ("fixed K(0)", "fixed_gap", "#7f7f7f", 0.38),
        ("linear velocity", "linear_gap", "#1f5a9d", 0.58),
        ("early re-tangent", "retangent_gap", "#2b8a3e", 0.45),
    ]
    max_value = float(
        max(
            empirical.max(),
            *(df[column].max() for _label, column, _color, _alpha in predictors),
        )
    )
    min_value = float(
        min(
            empirical.min(),
            *(df[column].min() for _label, column, _color, _alpha in predictors),
        )
    )
    for label, column, color, alpha in predictors:
        ax.scatter(
            df[column],
            empirical,
            s=19,
            alpha=alpha,
            color=color,
            label=label,
            edgecolors="none",
        )
    ax.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1.0)
    ax.set_title("Predicted vs empirical FA/BP train-gap")
    ax.set_xlabel("predicted gap")
    ax.set_ylabel("empirical gap")
    ax.legend(loc="upper left", frameon=True)
    ax.grid(alpha=0.18)
    fig.tight_layout()
    outpath = outdir / "T50_s20_prediction_scatter.png"
    fig.savefig(outpath)
    plt.close(fig)
    return outpath


def plot_error_by_margin(df: pd.DataFrame, outdir: Path) -> Path:
    fig, ax = plt.subplots(figsize=(10.5, 5.8), dpi=180)
    for label, column, color in [
        ("fixed K(0)", "fixed_gap", "#7f7f7f"),
        ("linear velocity", "linear_gap", "#1f5a9d"),
        ("early re-tangent", "retangent_gap", "#2b8a3e"),
    ]:
        records = []
        for margin, group in df.groupby("hard_fa_capacity_margin"):
            error = group[column].to_numpy() - group["empirical_gap"].to_numpy()
            records.append(
                {
                    "margin": margin,
                    "mean": float(error.mean()),
                    "q25": float(np.quantile(error, 0.25)),
                    "q75": float(np.quantile(error, 0.75)),
                }
            )
        summary = pd.DataFrame(records).sort_values("margin")
        ax.plot(summary["margin"], summary["mean"], color=color, marker="o", label=label)
        ax.fill_between(
            summary["margin"],
            summary["q25"],
            summary["q75"],
            color=color,
            alpha=0.15,
            linewidth=0,
        )
    ax.axhline(0, color="black", linewidth=1.0)
    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
    ax.set_title("Prediction error by capacity margin")
    ax.set_xlabel("hard FA capacity margin")
    ax.set_ylabel("predicted gap - empirical gap")
    ax.legend(loc="upper left", frameon=True)
    ax.grid(alpha=0.18)
    fig.tight_layout()
    outpath = outdir / "T50_s20_prediction_error_by_margin.png"
    fig.savefig(outpath)
    plt.close(fig)
    return outpath


def main() -> None:
    args = parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)
    df = read_rows(args.glob, args.input_dim, args.output_dim, args.width)
    df.to_csv(args.outdir / "combined_rows.csv", index=False)
    metrics = predictor_metrics(df)
    write_metrics(args.outdir / "predictor_metrics.csv", metrics)
    paths = [
        plot_theory_band_with_empirical(df, args.outdir),
        plot_empirical_only(df, args.outdir),
        plot_prediction_scatter(df, args.outdir),
        plot_error_by_margin(df, args.outdir),
    ]
    print(f"combined rows: {args.outdir / 'combined_rows.csv'}")
    print(f"metrics: {args.outdir / 'predictor_metrics.csv'}")
    for row in metrics:
        print(
            f"{row.predictor}: rows={row.rows}, MAE={row.mae:.6g}, "
            f"bias={row.bias:.6g}, corr={row.corr:.6g}"
        )
    for path in paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()