from __future__ import annotations

from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy import stats


PROJECT = Path(__file__).resolve().parent.parent
FIGURES = PROJECT / "results" / "figures"
TABLES = PROJECT / "results" / "tables"
STUDY1 = PROJECT / "results" / "study1"
STUDY2 = PROJECT / "results" / "study2"

COLORS = {
    "U3": "#009E73",
    "U5": "#0072B2",
    "U9": "#D55E00",
    "gaussian": "#0072B2",
    "selected": "#009E73",
    "ensemble": "#E69F00",
    "acceptable_set": "#D55E00",
}
MARKERS = {"U3": "^", "U5": "s", "U9": "o"}
LINESTYLES = {"U3": ":", "U5": "--", "U9": "-"}


def style() -> None:
    mpl.rcParams.update(
        {
            "font.family": "sans-serif",
            "font.sans-serif": ["Arial", "Helvetica", "DejaVu Sans"],
            "font.size": 8,
            "axes.labelsize": 8,
            "axes.titlesize": 9,
            "xtick.labelsize": 7,
            "ytick.labelsize": 7,
            "legend.fontsize": 7,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )


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


def figure1_pipeline() -> None:
    fig, ax = plt.subplots(figsize=(7.2, 2.6))
    ax.axis("off")
    boxes = [
        (0.02, "Declared model\nuniverse", "#D9EAF7"),
        (0.22, "Independent\nvalidation evidence", "#D9EAF7"),
        (0.42, "Empirically\nacceptable set", "#FCE8D5"),
        (0.62, "Decision\n diameter", "#FCE8D5"),
        (0.82, "Single report,\nrange, or policy", "#E2F0D9"),
    ]
    for x, label, color in boxes:
        rectangle = mpl.patches.FancyBboxPatch(
            (x, 0.40),
            0.15,
            0.30,
            boxstyle="round,pad=0.02",
            linewidth=1,
            edgecolor="#333333",
            facecolor=color,
        )
        ax.add_patch(rectangle)
        ax.text(x + 0.075, 0.55, label, ha="center", va="center", fontsize=8)
    for left in (0.17, 0.37, 0.57, 0.77):
        ax.annotate(
            "",
            xy=(left + 0.045, 0.55),
            xytext=(left, 0.55),
            arrowprops={"arrowstyle": "->", "lw": 1, "color": "#333333"},
        )
    ax.text(
        0.12,
        0.18,
        "Flexibility creates\nfit options",
        ha="center",
        va="center",
        color=COLORS["U5"],
        fontweight="bold",
        fontsize=7.5,
    )
    ax.text(
        0.52,
        0.18,
        "Evidence must separate\nconsequential alternatives",
        ha="center",
        va="center",
        color=COLORS["U9"],
        fontweight="bold",
        fontsize=7.5,
    )
    ax.text(
        0.88,
        0.18,
        "Action still requires\na loss function",
        ha="center",
        va="center",
        color="#333333",
        fontweight="bold",
        fontsize=7.5,
    )
    ax.plot([0.02, 0.97], [0.31, 0.31], color="#777777", lw=0.6)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    save(fig, "figure1_analytics_pipeline")


def mean_ci(group: pd.Series) -> tuple[float, float]:
    mean = float(group.mean())
    half = float(stats.t.ppf(0.975, len(group) - 1) * stats.sem(group))
    return mean, half


def figure2_interaction(cells: pd.DataFrame) -> None:
    stable = cells[
        cells.dgp.isin(
            [
                "stable_gaussian",
                "stable_student_t",
                "stable_skew_normal",
                "stable_mixture",
            ]
        )
        & (~cells.fit_failed)
    ]
    by_rep = (
        stable.groupby(["rep", "universe", "evidence_n"])[
            ["var_diameter", "acceptable_count"]
        ]
        .mean()
        .reset_index()
    )
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8))
    for ax, outcome, ylabel in (
        (axes[0], "var_diameter", "5% VaR decision diameter (SD units)"),
        (axes[1], "acceptable_count", "Acceptable models (count)"),
    ):
        for universe in ("U3", "U5", "U9"):
            subset = by_rep[by_rep.universe == universe]
            summary = (
                subset.groupby("evidence_n")[outcome]
                .apply(lambda values: pd.Series(mean_ci(values), index=["mean", "half"]))
                .unstack()
                .reset_index()
            )
            ax.errorbar(
                summary.evidence_n,
                summary["mean"],
                yerr=summary["half"],
                color=COLORS[universe],
                linestyle=LINESTYLES[universe],
                marker=MARKERS[universe],
                markersize=4,
                capsize=2,
                lw=1.2,
                label=universe,
            )
        ax.set_xscale("log")
        ax.set_xticks([20, 60, 120, 250, 500])
        ax.get_xaxis().set_major_formatter(mpl.ticker.ScalarFormatter())
        ax.set_xlabel("Validation observations, m")
        ax.set_ylabel(ylabel)
        ax.grid(axis="y", color="#DDDDDD", linewidth=0.5)
    axes[0].text(-0.17, 1.05, "A", transform=axes[0].transAxes, fontweight="bold")
    axes[1].text(-0.17, 1.05, "B", transform=axes[1].transAxes, fontweight="bold")
    axes[0].legend(frameon=False, title="Model universe")
    fig.tight_layout(w_pad=2.2)
    save(fig, "figure2_evidence_universe_interaction")


def figure3_dgp_contrasts(contrasts: pd.DataFrame) -> None:
    data = contrasts[
        (contrasts.outcome == "var_diameter")
        & (contrasts.scope != "stable_dgps_average")
    ].copy()
    order = [
        "stable_gaussian",
        "stable_student_t",
        "stable_skew_normal",
        "stable_mixture",
        "garch_scale_adjusted",
        "single_break",
        "markov_switching",
        "smooth_drift",
    ]
    labels = {
        "stable_gaussian": "Stable Gaussian",
        "stable_student_t": "Stable Student-t",
        "stable_skew_normal": "Stable skew-normal",
        "stable_mixture": "Stable mixture",
        "garch_scale_adjusted": "GARCH, scale adjusted",
        "single_break": "Single break",
        "markov_switching": "Markov switching",
        "smooth_drift": "Smooth drift",
    }
    data["position"] = data.scope.map({name: i for i, name in enumerate(order)})
    data = data.sort_values("position", ascending=False)
    fig, ax = plt.subplots(figsize=(5.4, 3.6))
    colors = [
        COLORS["U9"] if low > 0 else "#777777"
        for low in data.ci_low
    ]
    for (_, row), color in zip(data.iterrows(), colors):
        ax.errorbar(
            row["mean"],
            row.position,
            xerr=np.asarray(
                [[row["mean"] - row.ci_low], [row.ci_high - row["mean"]]]
            ),
            fmt="o",
            color=color,
            ecolor=color,
            markersize=4,
            capsize=2,
            lw=1.2,
            zorder=3,
        )
    ax.axvline(0, color="#333333", linewidth=0.8)
    ax.set_yticks(range(len(order)))
    ax.set_yticklabels([labels[name] for name in order])
    ax.set_xlabel("Evidence × universe interaction in VaR diameter")
    ax.grid(axis="x", color="#DDDDDD", linewidth=0.5)
    save(fig, "figure3_dgp_interactions")


def figure4_financial(summary: pd.DataFrame) -> None:
    policy_means = (
        summary.groupby("policy")[["coverage_error", "mean_capital"]]
        .mean()
        .loc[list(POLICIES)]
    )
    wide = summary.pivot(
        index="asset", columns="policy", values="coverage_error"
    )
    assets = wide.index.tolist()
    differences = pd.DataFrame(
        {
            "Gaussian": wide["acceptable_set"] - wide["gaussian"],
            "Selected": wide["acceptable_set"] - wide["selected"],
        }
    )
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.3))
    y = np.arange(len(assets))
    axes[0].axvline(0, color="#333333", lw=0.8)
    axes[0].scatter(
        differences["Gaussian"],
        y + 0.10,
        color=COLORS["gaussian"],
        marker="s",
        s=18,
        label="vs Gaussian",
    )
    axes[0].scatter(
        differences["Selected"],
        y - 0.10,
        color=COLORS["selected"],
        marker="o",
        s=18,
        label="vs selected",
    )
    axes[0].set_yticks(y)
    axes[0].set_yticklabels(assets)
    axes[0].set_xlabel("Acceptable-set minus comparator\nabsolute coverage error")
    axes[0].legend(frameon=False)
    axes[0].grid(axis="x", color="#DDDDDD", linewidth=0.5)
    for policy in POLICIES:
        subset = summary[summary.policy == policy]
        axes[1].scatter(
            subset.mean_capital,
            subset.coverage_error,
            color=COLORS[policy],
            alpha=0.18,
            s=12,
        )
        axes[1].scatter(
            policy_means.loc[policy, "mean_capital"],
            policy_means.loc[policy, "coverage_error"],
            color=COLORS[policy],
            edgecolor="#222222",
            linewidth=0.5,
            s=48,
            label=policy.replace("_", " ").title(),
            zorder=4,
        )
    axes[1].set_xlabel("Mean standardized capital")
    axes[1].set_ylabel("Absolute error from 5% violation target")
    axes[1].grid(color="#DDDDDD", linewidth=0.5)
    axes[1].legend(frameon=False, loc="upper left")
    axes[0].text(-0.18, 1.05, "A", transform=axes[0].transAxes, fontweight="bold")
    axes[1].text(-0.17, 1.05, "B", transform=axes[1].transAxes, fontweight="bold")
    fig.tight_layout(w_pad=2.0)
    save(fig, "figure4_coverage_capital")


POLICIES = ("gaussian", "selected", "ensemble", "acceptable_set")


def make_tables(
    cells: pd.DataFrame,
    contrasts: pd.DataFrame,
    asset_policy: pd.DataFrame,
    inference: pd.DataFrame,
) -> None:
    TABLES.mkdir(parents=True, exist_ok=True)
    stable = cells[
        cells.dgp.isin(
            [
                "stable_gaussian",
                "stable_student_t",
                "stable_skew_normal",
                "stable_mixture",
            ]
        )
        & (~cells.fit_failed)
    ]
    table2 = (
        stable.groupby(["universe", "evidence_n"])[
            [
                "best_validation_score",
                "acceptable_count",
                "var_diameter",
                "es_diameter",
                "winner_reversal",
                "sealed_test_regret",
            ]
        ]
        .mean()
        .reset_index()
    )
    table2.to_csv(TABLES / "table2_study1_levels.csv", index=False)
    contrasts.to_csv(TABLES / "table3_study1_interactions.csv", index=False)
    asset_policy.to_csv(TABLES / "table4_asset_policy.csv", index=False)
    inference.to_csv(TABLES / "table5_financial_inference.csv", index=False)
    summary = pd.DataFrame(
        [
            {
                "study": "Study 1",
                "unit": "replication × DGP × universe × evidence",
                "units": len(cells),
                "failures": int(cells.fit_failed.sum()),
                "primary_result": "VaR interaction 0.0526 [0.0440, 0.0612]",
            },
            {
                "study": "Study 2",
                "unit": "non-overlapping asset block",
                "units": 2710,
                "failures": 0,
                "primary_result": "Conservative policy worsens coverage",
            },
        ]
    )
    summary.to_csv(TABLES / "table1_evidence_overview.csv", index=False)


def main() -> None:
    style()
    cells = pd.read_csv(STUDY1 / "study1_cells.csv")
    contrasts = pd.read_csv(STUDY1 / "study1_contrasts.csv")
    asset_policy = pd.read_csv(STUDY2 / "study2_asset_policy.csv")
    inference = pd.read_csv(STUDY2 / "study2_inference.csv")
    figure1_pipeline()
    figure2_interaction(cells)
    figure3_dgp_contrasts(contrasts)
    figure4_financial(asset_policy)
    make_tables(cells, contrasts, asset_policy, inference)


if __name__ == "__main__":
    main()
