#!/usr/bin/env python3
from __future__ import annotations

import argparse
from pathlib import Path

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


ROOT = Path(__file__).resolve().parents[1]
COLORS = {"low": "#3B82F6", "high": "#F97316"}


def ci95(values: pd.Series) -> tuple[float, float, float]:
    x = values.dropna().to_numpy(float)
    mean = float(x.mean())
    if len(x) < 2:
        return mean, mean, mean
    rng = np.random.default_rng(20260717 + len(x))
    samples = rng.choice(x, size=(4000, len(x)), replace=True).mean(axis=1)
    lo, hi = np.quantile(samples, [0.025, 0.975])
    return mean, float(lo), float(hi)


def grouped_ci(frame: pd.DataFrame, by: list[str], metric: str) -> pd.DataFrame:
    records = []
    for keys, values in frame.groupby(by, sort=False)[metric]:
        keys = keys if isinstance(keys, tuple) else (keys,)
        mean, lo, hi = ci95(values)
        records.append({**dict(zip(by, keys)), "mean": mean, "lo": lo, "hi": hi})
    return pd.DataFrame(records)


def style() -> None:
    plt.rcParams.update(
        {
            "font.size": 8.5,
            "axes.titlesize": 9.5,
            "axes.labelsize": 9,
            "legend.fontsize": 8,
            "figure.dpi": 160,
            "savefig.bbox": "tight",
            "axes.spines.top": False,
            "axes.spines.right": False,
        }
    )


def save(fig: plt.Figure, directory: Path, name: str) -> None:
    directory.mkdir(parents=True, exist_ok=True)
    fig.savefig(directory / f"{name}.pdf")
    fig.savefig(directory / f"{name}.png", dpi=220)
    plt.close(fig)


def factorial_figure(exact: pd.DataFrame, directory: Path) -> None:
    fig, axes = plt.subplots(1, 3, figsize=(10.0, 2.7))
    metrics = [
        ("mean_root_event_entropy", "Event entropy", "Transition intervention"),
        ("root_gap", r"$V_{gen}-V_{ctrl}$", "Readout gap"),
        ("generative_policy_regret", "True control regret", "Generated policy"),
    ]
    x = np.array([0, 1])
    for ax, (metric, ylabel, title) in zip(axes, metrics):
        summary = grouped_ci(exact, ["entropy", "dispersion"], metric)
        for dispersion, offset, marker in [("low", -0.04, "o"), ("high", 0.04, "s")]:
            part = summary[summary.dispersion == dispersion].set_index("entropy").loc[["low", "high"]]
            means = part["mean"].to_numpy()
            errors = np.vstack([means - part["lo"].to_numpy(), part["hi"].to_numpy() - means])
            ax.errorbar(
                x + offset,
                means,
                yerr=errors,
                color=COLORS[dispersion],
                marker=marker,
                capsize=3,
                linewidth=1.7,
                label=f"{dispersion} dispersion",
            )
        ax.set_xticks(x, ["low", "high"])
        ax.set_xlabel("Transition entropy")
        ax.set_ylabel(ylabel)
        ax.set_title(title)
        ax.grid(axis="y", alpha=0.2)
    axes[0].legend(frameon=False, loc="best")
    fig.tight_layout()
    save(fig, directory, "factorial_gap")


def tabular_figure(tabular: pd.DataFrame, directory: Path) -> None:
    base = tabular[tabular.threshold == tabular.threshold.min()]
    fig, axes = plt.subplots(1, 3, figsize=(10.0, 2.7))
    for ax, metric, ylabel in [
        (axes[0], "trajectory_kl", "Exact trajectory KL"),
        (axes[1], "weighted_tv", "Occupancy-weighted event TV"),
    ]:
        summary = grouped_ci(base, ["n"], metric).sort_values("n")
        ax.plot(summary.n, summary["mean"], color="#2563EB", marker="o", linewidth=1.7)
        ax.fill_between(summary.n, summary.lo, summary.hi, color="#93C5FD", alpha=0.4)
        ax.set_xscale("log")
        ax.set_xlabel("Training trajectories")
        ax.set_ylabel(ylabel)
        ax.grid(alpha=0.2)
    for metric, label, color in [
        ("control_value_error", "control readout", "#2563EB"),
        ("generative_value_error", "generative readout", "#F97316"),
    ]:
        summary = grouped_ci(base, ["n"], metric).sort_values("n")
        axes[2].plot(
            summary.n, summary["mean"], color=color, marker="o",
            linewidth=1.7, label=label
        )
        axes[2].fill_between(summary.n, summary.lo, summary.hi, color=color, alpha=0.16)
    axes[2].set_xscale("log")
    axes[2].set_xlabel("Training trajectories")
    axes[2].set_ylabel("Absolute value error")
    axes[2].legend(frameon=False)
    axes[2].grid(alpha=0.2)
    fig.tight_layout()
    save(fig, directory, "tabular_learning")


def support_figure(tabular: pd.DataFrame, directory: Path) -> None:
    n_max = int(tabular.n.max())
    data = tabular[tabular.n == n_max]
    fig, axes = plt.subplots(1, 3, figsize=(10.0, 2.7))
    for ax, metric, ylabel, color in [
        (axes[0], "chance_coverage", "Supported chance mass", "#059669"),
        (axes[1], "max_supported_tv", "Max TV on supported nodes", "#7C3AED"),
        (axes[2], "total_gated_regret", "Total gated regret", "#DC2626"),
    ]:
        summary = grouped_ci(data, ["threshold"], metric).sort_values("threshold")
        ax.plot(summary.threshold, summary["mean"], color=color, marker="o", linewidth=1.7)
        ax.fill_between(summary.threshold, summary.lo, summary.hi, color=color, alpha=0.16)
        ax.set_xlabel("Visit-count threshold")
        ax.set_ylabel(ylabel)
        ax.grid(alpha=0.2)
    fig.tight_layout()
    save(fig, directory, "support_tradeoff")


def neural_figure(neural: pd.DataFrame, directory: Path) -> None:
    order = [("low", "low"), ("high", "low"), ("low", "high"), ("high", "high")]
    labels = ["L/L", "H/L", "L/H", "H/H"]
    x = np.arange(len(order))
    fig, axes = plt.subplots(1, 3, figsize=(10.0, 2.7))
    comparisons = [
        ("control_value_error", "generative_value_error", "Absolute value error"),
        ("control_policy_regret", "generative_policy_regret", "True policy regret"),
    ]
    for ax, (control_metric, gen_metric, ylabel) in zip(axes[:2], comparisons):
        for metric, label, color, offset in [
            (control_metric, "control readout", "#2563EB", -0.16),
            (gen_metric, "generative readout", "#F97316", 0.16),
        ]:
            means, los, his = [], [], []
            for entropy, dispersion in order:
                values = neural[(neural.entropy == entropy) & (neural.dispersion == dispersion)][metric]
                mean, lo, hi = ci95(values)
                means.append(mean); los.append(lo); his.append(hi)
            means = np.array(means); los = np.array(los); his = np.array(his)
            ax.bar(x + offset, means, width=0.30, color=color, label=label)
            ax.errorbar(x + offset, means, yerr=np.vstack([means-los, his-means]), fmt="none", ecolor="#111827", capsize=2, linewidth=0.8)
        ax.set_xticks(x, labels)
        ax.set_xlabel("Entropy / dispersion")
        ax.set_ylabel(ylabel)
        ax.grid(axis="y", alpha=0.2)
    axes[0].legend(frameon=False)
    summary = grouped_ci(neural, ["entropy"], "trajectory_kl").set_index("entropy").loc[["low", "high"]]
    axes[2].bar([0, 1], summary["mean"], color=[COLORS["low"], COLORS["high"]])
    axes[2].errorbar([0, 1], summary["mean"], yerr=np.vstack([summary["mean"]-summary.lo, summary.hi-summary["mean"]]), fmt="none", ecolor="#111827", capsize=3)
    axes[2].set_xticks([0, 1], ["low", "high"])
    axes[2].set_xlabel("Transition entropy")
    axes[2].set_ylabel("Exact trajectory KL")
    axes[2].grid(axis="y", alpha=0.2)
    fig.tight_layout()
    save(fig, directory, "neural_readouts")


def residual_figure(exact: pd.DataFrame, directory: Path) -> None:
    columns = [
        ("factorization_residual", "factor"),
        ("max_gap_recursion_residual", "gap"),
        ("relative_audit_residual", "audit"),
        ("path_decomposition_residual", "path"),
        ("event_mean_residual", "event mean"),
        ("kl_regret_identity_residual", "KL regret"),
    ]
    values = [max(float(exact[c].max()), 1e-18) for c, _ in columns]
    fig, ax = plt.subplots(figsize=(6.6, 2.65))
    ax.bar(np.arange(len(values)), values, color="#334155")
    ax.set_yscale("log")
    ax.axhline(1e-10, color="#DC2626", linestyle="--", linewidth=1.2, label="pre-specified rejection threshold")
    ax.set_xticks(np.arange(len(values)), [label for _, label in columns])
    ax.set_ylabel("Maximum absolute residual")
    ax.legend(frameon=False)
    ax.grid(axis="y", alpha=0.2, which="both")
    fig.tight_layout()
    save(fig, directory, "identity_residuals")


def macros(exact: pd.DataFrame, tabular: pd.DataFrame, neural: pd.DataFrame, output: Path) -> None:
    def fmt(x: float, digits: int = 3) -> str:
        return f"{x:.{digits}f}"

    ex = exact.groupby(["entropy", "dispersion"], sort=False).mean(numeric_only=True)
    tb = tabular[tabular.threshold == tabular.threshold.min()].groupby("n").mean(numeric_only=True)
    nn = neural.groupby(["entropy", "dispersion"], sort=False).mean(numeric_only=True)
    support = tabular[tabular.n == tabular.n.max()].groupby("threshold").mean(numeric_only=True)
    max_residual = exact[
        [
            "max_gap_recursion_residual", "factorization_residual", "normalization_residual",
            "relative_audit_residual", "counting_audit_residual", "path_decomposition_residual",
            "event_mean_residual", "kl_regret_identity_residual",
        ]
    ].to_numpy().max()
    lines = [
        f"\\newcommand{{\\ExactSeeds}}{{{exact.seed.nunique()}}}",
        f"\\newcommand{{\\ExactMaxResidual}}{{{max_residual:.2e}}}",
        f"\\newcommand{{\\GapLowLow}}{{{fmt(ex.loc[('low', 'low'), 'root_gap'])}}}",
        f"\\newcommand{{\\GapHighLow}}{{{fmt(ex.loc[('high', 'low'), 'root_gap'])}}}",
        f"\\newcommand{{\\GapLowHigh}}{{{fmt(ex.loc[('low', 'high'), 'root_gap'])}}}",
        f"\\newcommand{{\\GapHighHigh}}{{{fmt(ex.loc[('high', 'high'), 'root_gap'])}}}",
        f"\\newcommand{{\\RegretLowHigh}}{{{fmt(ex.loc[('low', 'high'), 'generative_policy_regret'])}}}",
        f"\\newcommand{{\\RegretHighHigh}}{{{fmt(ex.loc[('high', 'high'), 'generative_policy_regret'])}}}",
        f"\\newcommand{{\\TabularKLFirst}}{{{fmt(tb.loc[tb.index.min(), 'trajectory_kl'])}}}",
        f"\\newcommand{{\\TabularKLLast}}{{{fmt(tb.loc[tb.index.max(), 'trajectory_kl'])}}}",
        f"\\newcommand{{\\TabularTVFirst}}{{{fmt(tb.loc[tb.index.min(), 'weighted_tv'])}}}",
        f"\\newcommand{{\\TabularTVLast}}{{{fmt(tb.loc[tb.index.max(), 'weighted_tv'])}}}",
        f"\\newcommand{{\\TabularCtrlErrorFirst}}{{{fmt(tb.loc[tb.index.min(), 'control_value_error'])}}}",
        f"\\newcommand{{\\TabularCtrlErrorLast}}{{{fmt(tb.loc[tb.index.max(), 'control_value_error'])}}}",
        f"\\newcommand{{\\TabularGenErrorFirst}}{{{fmt(tb.loc[tb.index.min(), 'generative_value_error'])}}}",
        f"\\newcommand{{\\TabularGenErrorLast}}{{{fmt(tb.loc[tb.index.max(), 'generative_value_error'])}}}",
        f"\\newcommand{{\\TabularNFirst}}{{{int(tb.index.min())}}}",
        f"\\newcommand{{\\TabularNLast}}{{{int(tb.index.max()):,}}}",
        f"\\newcommand{{\\SupportAtLow}}{{{fmt(support.loc[support.index.min(), 'chance_coverage'])}}}",
        f"\\newcommand{{\\SupportAtHigh}}{{{fmt(support.loc[support.index.max(), 'chance_coverage'])}}}",
        f"\\newcommand{{\\NeuralKLHigh}}{{{fmt(neural[neural.entropy == 'high'].trajectory_kl.mean())}}}",
        f"\\newcommand{{\\NeuralKLlow}}{{{fmt(neural[neural.entropy == 'low'].trajectory_kl.mean())}}}",
        f"\\newcommand{{\\NeuralCtrlErrorHH}}{{{fmt(nn.loc[('high', 'high'), 'control_value_error'])}}}",
        f"\\newcommand{{\\NeuralGenErrorHH}}{{{fmt(nn.loc[('high', 'high'), 'generative_value_error'])}}}",
        f"\\newcommand{{\\NeuralCtrlRegretHH}}{{{fmt(nn.loc[('high', 'high'), 'control_policy_regret'])}}}",
        f"\\newcommand{{\\NeuralGenRegretHH}}{{{fmt(nn.loc[('high', 'high'), 'generative_policy_regret'])}}}",
    ]
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--exact", default=str(ROOT / "results" / "exact_factorial.csv"))
    parser.add_argument("--tabular", default=str(ROOT / "results" / "tabular.csv"))
    parser.add_argument("--neural", default=str(ROOT / "results" / "neural.csv"))
    parser.add_argument("--figures", default=str(ROOT / "paper" / "figures"))
    parser.add_argument("--macros", default=str(ROOT / "paper" / "generated" / "results_macros.tex"))
    args = parser.parse_args()
    style()
    exact = pd.read_csv(args.exact)
    tabular = pd.read_csv(args.tabular)
    neural = pd.read_csv(args.neural)
    directory = Path(args.figures)
    factorial_figure(exact, directory)
    tabular_figure(tabular, directory)
    support_figure(tabular, directory)
    neural_figure(neural, directory)
    residual_figure(exact, directory)
    macros(exact, tabular, neural, Path(args.macros))
    print(f"wrote figures to {directory}")


if __name__ == "__main__":
    main()
