#!/usr/bin/env python3
"""Summarize finite-time operator-predictor stress tests."""

from __future__ import annotations

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

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


STRESS_RE = re.compile(
    r"stress_depth(?P<depth>\d+)_w(?P<width>\d+)_T(?P<T>\d+)_s(?P<s>\d+)_N(?P<N>\d+)"
)


@dataclass(frozen=True)
class SettingMetrics:
    setting: str
    hidden_layers: int
    width: int
    target_steps: int
    early_steps: int
    rows: int
    fixed_mae: float
    fixed_bias: float
    compressed_mae: float
    compressed_bias: float
    linear_mae: float
    linear_bias: float
    retangent_mae: float
    retangent_bias: float
    linear_corr: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-dim", type=int, default=16)
    parser.add_argument("--output-dim", type=int, default=4)
    parser.add_argument("--include-baseline", action="store_true")
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/operator_stress_grid_summary"),
    )
    return parser.parse_args()


def hard_margin(
    input_dim: int,
    output_dim: int,
    width: np.ndarray,
    hidden_layers: np.ndarray,
    train_samples: np.ndarray,
) -> np.ndarray:
    parameter_count = (
        input_dim * width
        + np.maximum(hidden_layers - 1, 0) * width * width
        + width * output_dim
    )
    constraint_rank = (
        np.maximum(hidden_layers - 1, 0) * np.maximum(width * width - 1, 0)
        + np.maximum(width * output_dim - 1, 0)
    )
    return parameter_count - constraint_rank - output_dim * train_samples


def load_stress_rows(input_dim: int, output_dim: int, include_baseline: bool) -> pd.DataFrame:
    frames: list[pd.DataFrame] = []
    for path in sorted(Path("outputs").glob("stress_*/compressed_operator_rows.csv")):
        match = STRESS_RE.fullmatch(path.parent.name)
        if not match:
            continue
        meta = {key: int(value) for key, value in match.groupdict().items()}
        frame = pd.read_csv(path)
        frame["setting"] = (
            f"d{meta['depth']}_w{meta['width']}_T{meta['T']}_s{meta['s']}"
        )
        frame["hidden_layers"] = meta["depth"]
        frame["width"] = meta["width"]
        frame["target_steps"] = meta["T"]
        frame["early_steps"] = meta["s"]
        frame["train_samples"] = meta["N"]
        frames.append(frame)

    if include_baseline:
        for path in sorted(
            Path("outputs").glob(
                "compressed_operator_s20_256traj_T50_width64_N*/compressed_operator_rows.csv"
            )
        ):
            train_samples = int(path.parent.name.rsplit("N", 1)[1])
            frame = pd.read_csv(path)
            frame["setting"] = "d2_w64_T50_s20_256traj"
            frame["hidden_layers"] = 2
            frame["width"] = 64
            frame["target_steps"] = 50
            frame["early_steps"] = 20
            frame["train_samples"] = train_samples
            frames.append(frame)

    if not frames:
        raise FileNotFoundError("No stress rows found.")

    df = pd.concat(frames, ignore_index=True)
    df["hard_fa_capacity_margin"] = hard_margin(
        input_dim,
        output_dim,
        df["width"].to_numpy(dtype=np.float64),
        df["hidden_layers"].to_numpy(dtype=np.float64),
        df["train_samples"].to_numpy(dtype=np.float64),
    )
    return df


def summarize(df: pd.DataFrame) -> list[SettingMetrics]:
    rows: list[SettingMetrics] = []
    for setting, group in df.groupby("setting"):
        empirical = group["empirical_gap"].to_numpy()

        def error_stats(column: str) -> tuple[float, float]:
            error = group[column].to_numpy() - empirical
            return float(np.mean(np.abs(error))), float(np.mean(error))

        fixed_mae, fixed_bias = error_stats("fixed_gap")
        compressed_mae, compressed_bias = error_stats("compressed_gap")
        linear_mae, linear_bias = error_stats("linear_gap")
        retangent_mae, retangent_bias = error_stats("retangent_gap")
        linear_corr = float(np.corrcoef(group["linear_gap"], empirical)[0, 1])
        rows.append(
            SettingMetrics(
                setting=setting,
                hidden_layers=int(group["hidden_layers"].iloc[0]),
                width=int(group["width"].iloc[0]),
                target_steps=int(group["target_steps"].iloc[0]),
                early_steps=int(group["early_steps"].iloc[0]),
                rows=len(group),
                fixed_mae=fixed_mae,
                fixed_bias=fixed_bias,
                compressed_mae=compressed_mae,
                compressed_bias=compressed_bias,
                linear_mae=linear_mae,
                linear_bias=linear_bias,
                retangent_mae=retangent_mae,
                retangent_bias=retangent_bias,
                linear_corr=linear_corr,
            )
        )
    return sorted(rows, key=lambda row: row.setting)


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


def plot_mae_bars(metrics: list[SettingMetrics], outdir: Path) -> Path:
    labels = [row.setting for row in metrics]
    x = np.arange(len(labels))
    width = 0.22
    fig, ax = plt.subplots(figsize=(13.5, 6.4), dpi=180)
    ax.bar(x - width, [row.fixed_mae for row in metrics], width, label="fixed K(0)", color="#777777")
    ax.bar(x, [row.linear_mae for row in metrics], width, label="linear velocity", color="#1f5a9d")
    ax.bar(
        x + width,
        [row.retangent_mae for row in metrics],
        width,
        label="early re-tangent",
        color="#2b8a3e",
    )
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=35, ha="right")
    ax.set_ylabel("gap prediction MAE")
    ax.set_title("Operator-predictor stress grid")
    ax.legend(frameon=True)
    ax.grid(axis="y", alpha=0.18)
    fig.tight_layout()
    path = outdir / "stress_grid_gap_mae_by_setting.png"
    fig.savefig(path)
    plt.close(fig)
    return path


def plot_linear_error_by_margin(df: pd.DataFrame, outdir: Path) -> Path:
    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)
    for setting, group in df.groupby("setting"):
        records = []
        for margin, by_margin in group.groupby("hard_fa_capacity_margin"):
            error = by_margin["linear_gap"].to_numpy() - by_margin["empirical_gap"].to_numpy()
            records.append((margin, float(np.mean(error)), float(np.mean(np.abs(error)))))
        records = sorted(records)
        ax.plot(
            [record[0] for record in records],
            [record[1] for record in records],
            marker="o",
            linewidth=1.7,
            label=setting,
        )
    ax.axhline(0, color="black", linewidth=1.0)
    ax.axvline(0, color="black", linestyle="--", linewidth=1.0)
    ax.set_xlabel("hard FA capacity margin")
    ax.set_ylabel("linear velocity prediction bias")
    ax.set_title("Linear-velocity bias across capacity margins")
    ax.legend(frameon=True, fontsize=8, ncol=2)
    ax.grid(alpha=0.18)
    fig.tight_layout()
    path = outdir / "stress_grid_linear_bias_by_margin.png"
    fig.savefig(path)
    plt.close(fig)
    return path


def plot_linear_scatter(df: pd.DataFrame, outdir: Path) -> Path:
    fig, ax = plt.subplots(figsize=(7.2, 6.4), dpi=180)
    settings = sorted(df["setting"].unique())
    cmap = plt.get_cmap("tab10")
    for index, setting in enumerate(settings):
        group = df[df["setting"] == setting]
        ax.scatter(
            group["linear_gap"],
            group["empirical_gap"],
            s=22,
            alpha=0.62,
            color=cmap(index % 10),
            label=setting,
            edgecolors="none",
        )
    lo = float(min(df["linear_gap"].min(), df["empirical_gap"].min()))
    hi = float(max(df["linear_gap"].max(), df["empirical_gap"].max()))
    ax.plot([lo, hi], [lo, hi], color="black", linewidth=1.0)
    ax.set_xlabel("linear velocity predicted gap")
    ax.set_ylabel("empirical gap")
    ax.set_title("Stress-grid predicted vs empirical gaps")
    ax.legend(frameon=True, fontsize=8, ncol=2)
    ax.grid(alpha=0.18)
    fig.tight_layout()
    path = outdir / "stress_grid_linear_prediction_scatter.png"
    fig.savefig(path)
    plt.close(fig)
    return path


def main() -> None:
    args = parse_args()
    args.outdir.mkdir(parents=True, exist_ok=True)
    df = load_stress_rows(args.input_dim, args.output_dim, args.include_baseline)
    df.to_csv(args.outdir / "stress_grid_rows.csv", index=False)
    metrics = summarize(df)
    write_metrics(args.outdir / "stress_grid_metrics.csv", metrics)
    paths = [
        plot_mae_bars(metrics, args.outdir),
        plot_linear_error_by_margin(df, args.outdir),
        plot_linear_scatter(df, args.outdir),
    ]
    print(f"rows: {args.outdir / 'stress_grid_rows.csv'}")
    print(f"metrics: {args.outdir / 'stress_grid_metrics.csv'}")
    for row in metrics:
        print(
            f"{row.setting}: rows={row.rows}, fixed={row.fixed_mae:.6g}, "
            f"linear={row.linear_mae:.6g}, retangent={row.retangent_mae:.6g}, "
            f"linear_bias={row.linear_bias:.6g}, corr={row.linear_corr:.6g}"
        )
    for path in paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()