"""Publication figures (matplotlib, vector PDF). Palette validated with the dataviz checker (light surface)."""
from __future__ import annotations

from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402

SERIES = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#6250d6", "#008300", "#e34948"]
INK, INK2, GRID = "#0b0b0b", "#52514e", "#e6e5e1"

plt.rcParams.update({
    "font.family": "serif", "font.size": 9, "axes.labelsize": 9, "axes.titlesize": 9.5, "legend.fontsize": 8,
    "xtick.labelsize": 8, "ytick.labelsize": 8, "axes.edgecolor": INK2, "axes.linewidth": 0.6, "axes.grid": True,
    "grid.color": GRID, "grid.linewidth": 0.5, "axes.spines.top": False, "axes.spines.right": False,
    "lines.linewidth": 1.6, "figure.dpi": 150, "savefig.bbox": "tight", "pdf.fonttype": 42, "text.color": INK,
    "axes.labelcolor": INK, "xtick.color": INK2, "ytick.color": INK2, "legend.frameon": False,
})


def save(fig, path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path)
    fig.savefig(path.with_suffix(".png"), dpi=200)
    plt.close(fig)


def price_paths(runs: list[tuple[str, pd.DataFrame]], shock_t: int | None = None, title: str = "") -> plt.Figure:
    """Mid vs latent value for several runs (one panel per run)."""
    n = len(runs)
    fig, axes = plt.subplots(1, n, figsize=(2.9 * n, 2.4), sharey=True)
    axes = np.atleast_1d(axes)
    for ax, (label, tk), c in zip(axes, runs, SERIES):
        d = tk[tk.asset == 0]
        ax.plot(d.t, d.full_value / 100, color=INK2, lw=1.0, ls="--", label="latent value")
        ax.plot(d.t, d.public_value / 100, color=INK2, lw=0.7, ls=":", label="public value")
        ax.plot(d.t, d.mid / 100, color=c, lw=1.2, label="mid")
        if shock_t:
            ax.axvline(shock_t, color=GRID, lw=1.0)
        ax.set_title(label)
        ax.set_xlabel("tick")
    axes[0].set_ylabel("price")
    axes[0].legend(loc="upper left")
    if title:
        fig.suptitle(title, fontsize=9.5)
    return fig


def dose_response(eff: dict[str, pd.DataFrame], ylabels: dict[str, str], x_col: str = "level", ncols: int = 3) -> plt.Figure:
    """eff: outcome -> DataFrame with columns level, diff, ci_lo, ci_hi (paired effects vs control)."""
    keys = list(eff)
    nrows = int(np.ceil(len(keys) / ncols))
    fig, axes = plt.subplots(nrows, ncols, figsize=(2.6 * ncols, 2.1 * nrows))
    axes = np.atleast_1d(axes).ravel()
    for i, (ax, k) in enumerate(zip(axes, keys)):
        d = eff[k]
        if "ci_lo" not in d or d["diff"].isna().all():
            ax.set_title(ylabels.get(k, k) + " (n/a)")
            continue
        x = d[x_col].astype(float).to_numpy()
        sc = 1e4 if k == "msd_log" else 1.0
        ax.axhline(0, color=INK2, lw=0.7)
        ax.fill_between(x, sc * d.ci_lo, sc * d.ci_hi, color=SERIES[0], alpha=0.18, lw=0)
        ax.plot(x, sc * d["diff"], color=SERIES[0], marker="o", ms=4)
        ax.set_title(ylabels.get(k, k))
        if i >= len(keys) - ncols:
            ax.set_xlabel("LLM share of trader slots")
    for ax in axes[len(keys):]:
        ax.axis("off")
    fig.tight_layout()
    return fig


def cell_bars(df: pd.DataFrame, outcome: str, cells: list[str], labels: dict[str, str] | None = None, ylabel: str = "",
              control: str | None = None) -> plt.Figure:
    g = df[df.cell.isin(cells)].groupby("cell")[outcome]
    m, se = g.mean().reindex(cells), (g.std(ddof=1) / np.sqrt(g.size())).reindex(cells)
    fig, ax = plt.subplots(figsize=(0.55 * len(cells) + 1.5, 2.3))
    cols = [SERIES[0] if c != control else INK2 for c in cells]
    ax.bar(range(len(cells)), m, yerr=1.96 * se, color=cols, width=0.62, error_kw={"lw": 0.8, "capsize": 2})
    ax.set_xticks(range(len(cells)))
    ax.set_xticklabels([(labels or {}).get(c, c) for c in cells], rotation=30, ha="right")
    ax.set_ylabel(ylabel or outcome)
    ax.grid(axis="x", visible=False)
    return fig


def local_projection_plot(lp: dict[str, pd.DataFrame], title: str = "") -> plt.Figure:
    fig, ax = plt.subplots(figsize=(3.6, 2.4))
    for (label, d), c in zip(lp.items(), SERIES):
        ax.plot(d.h, 100 * d.beta, color=c, label=label)
        ax.fill_between(d.h, 100 * (d.beta - 1.96 * d.se), 100 * (d.beta + 1.96 * d.se), color=c, alpha=0.15, lw=0)
    ax.axhline(0, color=INK2, lw=0.7)
    ax.set_xlabel("ticks after the fundamental shock")
    ax.set_ylabel("treatment effect on log(P/V) gap, %")
    ax.legend()
    if title:
        ax.set_title(title)
    return fig


def scatter_mechanism(df: pd.DataFrame, x: str, y: str, hue: str, xlabel: str, ylabel: str) -> plt.Figure:
    fig, ax = plt.subplots(figsize=(3.4, 2.6))
    for (k, d), c in zip(df.groupby(hue), SERIES):
        ax.scatter(d[x], d[y], s=14, color=c, label=str(k), alpha=0.85, edgecolor="white", linewidth=0.4)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.legend(title=hue)
    return fig


def wealth_and_inventory(agents: pd.DataFrame, ticks: pd.DataFrame | None = None) -> plt.Figure:
    fig, axes = plt.subplots(1, 2, figsize=(6.0, 2.4))
    ax = axes[0]
    for (pol, d), c in zip(agents.groupby("policy"), SERIES):
        ax.hist(d["pnl_pct"], bins=25, alpha=0.6, color=c, label=pol)
    ax.set_xlabel("P&L, % of initial wealth")
    ax.set_ylabel("traders")
    ax.legend()
    ax = axes[1]
    if ticks is not None:
        d = ticks[ticks.asset == 0]
        ax.plot(d.t, d.depth_bid5, color=SERIES[0], label="bid depth (5 levels)")
        ax.plot(d.t, d.depth_ask5, color=SERIES[1], label="ask depth")
        ax.set_xlabel("tick")
        ax.set_ylabel("lots")
        ax.legend()
    return fig


def cost_quality(df: pd.DataFrame, xcol: str, ycol: str, label_col: str, xlabel: str, ylabel: str) -> plt.Figure:
    fig, ax = plt.subplots(figsize=(3.4, 2.6))
    for (k, d), c in zip(df.groupby(label_col), SERIES):
        ax.errorbar(d[xcol].mean(), d[ycol].mean(), yerr=1.96 * d[ycol].std(ddof=1) / np.sqrt(len(d)) if len(d) > 1 else 0,
                    fmt="o", color=c, ms=5, capsize=2, label=str(k))
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.legend()
    return fig