#!/usr/bin/env python3
"""Cluster bootstrap metrics and the supplementary finite-time figure.

Both outputs are derived from outputs/aaai_depth_experiments/finite_time_rows.csv.

1. A LaTeX fragment with MAE, relative MAE, and correlation for the three
   predictors, each with a 95% cluster bootstrap confidence interval that
   resamples the four forward initializations within every architecture cell
   (feedback draws stay nested inside their forward initialization).
2. A six-panel supplementary figure: measured cost against depth, the three
   predictor calibration scatters, and signed residuals against depth and width.
"""

from __future__ import annotations

import argparse
import csv
from collections import defaultdict
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

PREDICTORS = [
    ("fixed_gap", "Frozen initialization operator", "frozen"),
    ("velocity_gap", "Linear extrapolation", "linear"),
    ("retangent_gap", "Relinearization", "relinearized"),
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--rows", type=Path, default=Path("outputs/aaai_depth_experiments/finite_time_rows.csv"))
    parser.add_argument("--table-outdir", type=Path, default=Path("outputs/tables"))
    parser.add_argument("--figure-outdir", type=Path, default=Path("outputs/aaai_supplementary"))
    parser.add_argument("--bootstrap-samples", type=int, default=10_000)
    parser.add_argument("--bootstrap-seed", type=int, default=2027)
    return parser.parse_args()


def load(path: Path) -> list[dict[str, str]]:
    with path.open() as handle:
        return list(csv.DictReader(handle))


def metrics(measured: np.ndarray, predicted: np.ndarray) -> tuple[float, float, float]:
    mae = float(np.mean(np.abs(predicted - measured)))
    rmae = mae / float(np.mean(measured))
    corr = float(np.corrcoef(predicted, measured)[0, 1])
    return mae, rmae, corr


def bootstrap_table(rows: list[dict[str, str]], args: argparse.Namespace) -> None:
    # cell -> init -> array of shape (8, 1 + predictors)
    cells: dict[tuple[int, int], dict[int, list[list[float]]]] = defaultdict(lambda: defaultdict(list))
    for row in rows:
        cells[(int(row["depth"]), int(row["width"]))][int(row["init_seed"])].append(
            [float(row["empirical_gap"])] + [float(row[f]) for f, _, _ in PREDICTORS]
        )
    stacked = [
        np.asarray([by_init[i] for i in sorted(by_init)])  # (4, 8, 4)
        for _, by_init in sorted(cells.items())
    ]
    full = np.concatenate([cell.reshape(-1, cell.shape[-1]) for cell in stacked])
    point = [metrics(full[:, 0], full[:, 1 + k]) for k in range(len(PREDICTORS))]

    rng = np.random.default_rng(args.bootstrap_seed)
    draws = np.empty((args.bootstrap_samples, len(PREDICTORS), 3))
    n_init = stacked[0].shape[0]
    for b in range(args.bootstrap_samples):
        sample = np.concatenate(
            [
                cell[rng.integers(0, n_init, size=n_init)].reshape(-1, cell.shape[-1])
                for cell in stacked
            ]
        )
        for k in range(len(PREDICTORS)):
            draws[b, k] = metrics(sample[:, 0], sample[:, 1 + k])
    lo, hi = np.percentile(draws, [2.5, 97.5], axis=0)

    lines = [
        r"\begin{tabular}{@{}>{\raggedright\arraybackslash}p{2.4cm}ccc@{}}",
        r"\toprule",
        r"Predictor & MAE & Relative MAE & Correlation \\",
        r"\midrule",
    ]
    for k, (_, label, _) in enumerate(PREDICTORS):
        mae, rmae, corr = point[k]
        lines.append(
            f"{label} & {mae:.5f} [{lo[k,0]:.5f}, {hi[k,0]:.5f}] & "
            f"{100*rmae:.2f}\\% [{100*lo[k,1]:.2f}, {100*hi[k,1]:.2f}] & "
            f"{corr:.5f} [{lo[k,2]:.5f}, {hi[k,2]:.5f}] \\\\"
        )
    lines += [r"\bottomrule", r"\end{tabular}"]
    args.table_outdir.mkdir(parents=True, exist_ok=True)
    out = args.table_outdir / "finite_time_metrics_ci.tex"
    out.write_text("\n".join(lines) + "\n")
    print(f"wrote {out}")


def supplementary_figure(rows: list[dict[str, str]], args: argparse.Namespace) -> None:
    depths = sorted({int(r["depth"]) for r in rows})
    widths = sorted({int(r["width"]) for r in rows})
    depth_color = {d: plt.cm.viridis(i / (len(depths) - 1)) for i, d in enumerate(depths)}

    fig, axes = plt.subplots(2, 3, figsize=(14.0, 8.2), dpi=220)

    ax = axes[0, 0]
    markers = ["o", "s", "^"]
    for width, marker in zip(widths, markers):
        means, sems = [], []
        for depth in depths:
            cell = [r for r in rows if int(r["width"]) == width and int(r["depth"]) == depth]
            init_means = [
                np.mean([float(r["empirical_gap"]) for r in cell if int(r["init_seed"]) == seed])
                for seed in sorted({int(r["init_seed"]) for r in cell})
            ]
            means.append(np.mean(init_means))
            sems.append(np.std(init_means, ddof=1) / np.sqrt(len(init_means)))
        ax.errorbar(depths, means, yerr=sems, marker=marker, capsize=2, label=f"width {width}")
    ax.set_xlabel("hidden-layer depth")
    ax.set_ylabel("FA loss - BP loss")
    ax.set_title("Measured cost by depth and width", fontsize=10)
    ax.grid(alpha=0.16)
    ax.legend(fontsize=7)

    scatter_axes = [axes[0, 1], axes[0, 2], axes[1, 0]]
    for (field, label, short), ax in zip(PREDICTORS, scatter_axes):
        xs = np.asarray([float(r[field]) for r in rows])
        ys = np.asarray([float(r["empirical_gap"]) for r in rows])
        colors = [depth_color[int(r["depth"])] for r in rows]
        ax.scatter(xs, ys, s=12, c=colors, alpha=0.6)
        lo = min(xs.min(), ys.min())
        hi = max(xs.max(), ys.max())
        pad = 0.04 * (hi - lo)
        ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], "k-", lw=1)
        ax.set_xlabel("predicted gap")
        ax.set_ylabel("measured gap")
        ax.set_title(label, fontsize=10)
        ax.grid(alpha=0.16)

    residual_specs = [(axes[1, 1], "depth", "hidden-layer depth"), (axes[1, 2], "width", "hidden width")]
    predictor_colors = ["#2f6f9f", "#c65f16", "#2b8a3e"]
    offsets = [-0.18, 0.0, 0.18]
    for ax, key, xlabel in residual_specs:
        levels = depths if key == "depth" else widths
        spacing = min(np.diff(levels)) if len(levels) > 1 else 1.0
        for (field, label, short), color, off in zip(PREDICTORS, predictor_colors, offsets):
            xs = np.asarray([float(r[key]) for r in rows]) + off * spacing
            resid = np.asarray([float(r[field]) - float(r["empirical_gap"]) for r in rows])
            ax.scatter(xs, resid, s=8, color=color, alpha=0.45, label=label)
        ax.axhline(0.0, color="0.3", lw=1)
        ax.set_xlabel(xlabel)
        ax.set_ylabel("predicted minus measured")
        ax.set_title(f"Signed residual by {key}", fontsize=10)
        ax.grid(alpha=0.16)
        if key == "depth":
            ax.legend(fontsize=7)

    for ax, letter in zip(axes.flat, "abcdef"):
        ax.text(-0.14, 1.06, letter, transform=ax.transAxes, fontsize=12, fontweight="bold")

    depth_handles = [
        plt.Line2D([], [], marker="o", ls="", color=depth_color[d], label=f"depth {d}")
        for d in depths
    ]
    axes[0, 2].legend(handles=depth_handles, fontsize=6, ncols=2, loc="upper left")

    fig.tight_layout(w_pad=2.0, h_pad=2.4)
    args.figure_outdir.mkdir(parents=True, exist_ok=True)
    path = args.figure_outdir / "finite_time_supplement.png"
    fig.savefig(path, bbox_inches="tight")
    fig.savefig(path.with_suffix(".pdf"), bbox_inches="tight")
    plt.close(fig)
    print(f"figure: {path}")


def main() -> None:
    args = parse_args()
    rows = load(args.rows)
    bootstrap_table(rows, args)
    supplementary_figure(rows, args)


if __name__ == "__main__":
    main()