"""Publication-quality figures for the Polymarket Arrow-Debreu paper."""

from __future__ import annotations

from datetime import datetime

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

from src.config import ELECTION_KEY_DATES, JOF_COLORS
from src.plots.style import COLORS, save_figure, set_jof_style


def fig_state_prices(state_prices: pd.DataFrame, title: str = "") -> plt.Figure:
    """Figure: Arrow-Debreu state prices over time (Fed rate decisions)."""
    set_jof_style()
    fig, ax = plt.subplots(figsize=(7, 4.5))
    for i, col in enumerate(state_prices.columns):
        ax.plot(state_prices.index, state_prices[col], label=col,
                color=COLORS[i % len(COLORS)], linewidth=1.3)
    ax.set_xlabel("Date")
    ax.set_ylabel("State Price")
    ax.set_title(title or "Arrow--Debreu State Prices: Fed Rate Decision")
    ax.legend(loc="best", framealpha=0.9)
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
    fig.autofmt_xdate(rotation=45)
    fig.tight_layout()
    return fig


def fig_completeness_test(daily_sums: pd.Series) -> plt.Figure:
    """Figure: Sum of state prices over time (should ≈ 1.0)."""
    set_jof_style()
    fig, ax = plt.subplots(figsize=(7, 3.5))
    ax.plot(daily_sums.index, daily_sums.values, color=JOF_COLORS["primary"], linewidth=1.2)
    ax.axhline(y=1.0, color=JOF_COLORS["secondary"], linestyle="--", linewidth=0.8, label="Complete market (1.0)")
    ax.fill_between(daily_sums.index, 0.98, 1.02, alpha=0.1, color=JOF_COLORS["gray"])
    ax.set_xlabel("Date")
    ax.set_ylabel("Sum of State Prices")
    ax.set_title("Arrow--Debreu Completeness Test")
    ax.legend(loc="best")
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    fig.autofmt_xdate(rotation=45)
    fig.tight_layout()
    return fig


def fig_risk_neutral_density(
    rn_density: pd.DataFrame,
    state_values: list[float],
    snapshot_dates: list[str] | None = None,
) -> plt.Figure:
    """Figure: Risk-neutral density snapshots at key dates."""
    set_jof_style()
    fig, ax = plt.subplots(figsize=(7, 4.5))
    if snapshot_dates is None:
        idx = rn_density.index
        n = len(idx)
        positions = [0, n // 4, n // 2, 3 * n // 4, n - 1]
        snapshot_dates = [idx[p] for p in positions if p < n]

    for i, d in enumerate(snapshot_dates):
        if d in rn_density.index:
            row = rn_density.loc[d]
        else:
            nearest = rn_density.index[rn_density.index.get_indexer([d], method="nearest")]
            row = rn_density.loc[nearest[0]]
            d = nearest[0]
        label = d.strftime("%b %d, %Y") if hasattr(d, "strftime") else str(d)
        ax.plot(state_values, row.values, "o-", color=COLORS[i % len(COLORS)],
                label=label, linewidth=1.3, markersize=5)

    ax.set_xlabel("Rate Change (basis points)")
    ax.set_ylabel("Risk-Neutral Probability")
    ax.set_title("Risk-Neutral Density: Fed Rate Decision")
    ax.legend(loc="best")
    fig.tight_layout()
    return fig


def fig_event_study_cars(
    cars_by_type: dict[str, pd.DataFrame],
    window: tuple[int, int] = (-5, 5),
) -> plt.Figure:
    """Figure: Cumulative abnormal returns around events, paneled by event type."""
    set_jof_style()
    n_types = len(cars_by_type)
    fig, axes = plt.subplots(1, n_types, figsize=(4 * n_types, 4), sharey=True)
    if n_types == 1:
        axes = [axes]

    for ax, (event_type, cars) in zip(axes, cars_by_type.items()):
        days = range(window[0], window[1] + 1)
        if "mean_CAR" in cars.columns:
            mean_car = cars["mean_CAR"].values
            if "std_CAR" in cars.columns:
                std_car = cars["std_CAR"].values
                ax.fill_between(days[:len(mean_car)],
                                mean_car - 1.96 * std_car,
                                mean_car + 1.96 * std_car,
                                alpha=0.2, color=JOF_COLORS["primary"])
            ax.plot(days[:len(mean_car)], mean_car, color=JOF_COLORS["primary"], linewidth=1.5)
        ax.axhline(0, color="black", linewidth=0.5)
        ax.axvline(0, color=JOF_COLORS["secondary"], linewidth=0.8, linestyle="--")
        ax.set_xlabel("Days Relative to Event")
        ax.set_title(event_type)

    axes[0].set_ylabel("Cumulative Abnormal Return")
    fig.suptitle("Event Study: Cumulative Abnormal Returns", y=1.02)
    fig.tight_layout()
    return fig


def fig_granger_heatmap(granger_pvalues: pd.DataFrame) -> plt.Figure:
    """Figure: Heatmap of Granger causality p-values."""
    set_jof_style()
    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(granger_pvalues.values, cmap="RdYlGn_r", vmin=0, vmax=0.10, aspect="auto")
    ax.set_xticks(range(len(granger_pvalues.columns)))
    ax.set_xticklabels(granger_pvalues.columns, rotation=45, ha="right")
    ax.set_yticks(range(len(granger_pvalues.index)))
    ax.set_yticklabels(granger_pvalues.index)
    ax.set_xlabel("Causing Variable")
    ax.set_ylabel("Caused Variable")
    ax.set_title("Granger Causality Test ($p$-values)")
    for i in range(len(granger_pvalues.index)):
        for j in range(len(granger_pvalues.columns)):
            val = granger_pvalues.iloc[i, j]
            color = "white" if val < 0.05 else "black"
            ax.text(j, i, f"{val:.3f}", ha="center", va="center", fontsize=8, color=color)
    fig.colorbar(im, ax=ax, label="$p$-value")
    fig.tight_layout()
    return fig


def fig_impulse_response(irf_data: dict, periods: int = 20) -> plt.Figure:
    """Figure: VAR impulse response functions."""
    set_jof_style()
    shock_vars = list({k[0] for k in irf_data.keys()})
    response_vars = list({k[1] for k in irf_data.keys()})
    n_shocks = len(shock_vars)
    n_responses = len(response_vars)
    fig, axes = plt.subplots(n_shocks, n_responses, figsize=(4 * n_responses, 3 * n_shocks),
                             sharex=True)
    if n_shocks == 1 and n_responses == 1:
        axes = np.array([[axes]])
    elif n_shocks == 1:
        axes = axes.reshape(1, -1)
    elif n_responses == 1:
        axes = axes.reshape(-1, 1)

    for i, shock in enumerate(shock_vars):
        for j, resp in enumerate(response_vars):
            ax = axes[i, j]
            key = (shock, resp)
            if key in irf_data:
                response = irf_data[key][:periods]
                ax.plot(range(len(response)), response, color=JOF_COLORS["primary"], linewidth=1.2)
                ax.axhline(0, color="black", linewidth=0.5)
            if i == 0:
                ax.set_title(f"Response: {resp}", fontsize=9)
            if j == 0:
                ax.set_ylabel(f"Shock: {shock}", fontsize=9)
            if i == n_shocks - 1:
                ax.set_xlabel("Periods")

    fig.suptitle("Impulse Response Functions", y=1.02)
    fig.tight_layout()
    return fig


def fig_volatility(
    rv: pd.Series,
    cond_vol: pd.Series | None = None,
    fitted_har: pd.Series | None = None,
) -> plt.Figure:
    """Figure: Volatility analysis panels."""
    set_jof_style()
    n_panels = 1 + (cond_vol is not None) + (fitted_har is not None)
    fig, axes = plt.subplots(n_panels, 1, figsize=(7, 3 * n_panels), sharex=True)
    if n_panels == 1:
        axes = [axes]

    panel = 0
    axes[panel].plot(rv.index, rv.values, color=JOF_COLORS["primary"], linewidth=0.8)
    axes[panel].set_ylabel("Realized Volatility")
    axes[panel].set_title("Panel A: Realized Volatility")
    panel += 1

    if cond_vol is not None:
        axes[panel].plot(cond_vol.index, cond_vol.values, color=JOF_COLORS["secondary"], linewidth=0.8)
        axes[panel].set_ylabel("Conditional Volatility")
        axes[panel].set_title("Panel B: GARCH Conditional Volatility")
        panel += 1

    if fitted_har is not None:
        axes[panel].plot(rv.index, rv.values, color=JOF_COLORS["gray"], linewidth=0.6, alpha=0.5, label="Actual")
        axes[panel].plot(fitted_har.index, fitted_har.values, color=JOF_COLORS["tertiary"], linewidth=0.8, label="HAR-RV Fit")
        axes[panel].set_ylabel("Realized Volatility")
        axes[panel].set_title("Panel C: HAR-RV Model Fit")
        axes[panel].legend()

    axes[-1].set_xlabel("Date")
    axes[-1].xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    fig.autofmt_xdate(rotation=45)
    fig.tight_layout()
    return fig


def fig_liquidity(volume: pd.Series, spreads: pd.Series | None = None) -> plt.Figure:
    """Figure: Liquidity dynamics."""
    set_jof_style()
    n_panels = 1 + (spreads is not None)
    fig, axes = plt.subplots(n_panels, 1, figsize=(7, 3 * n_panels), sharex=True)
    if n_panels == 1:
        axes = [axes]

    axes[0].bar(volume.index, volume.values, color=JOF_COLORS["primary"], alpha=0.6, width=1)
    axes[0].set_ylabel("Trading Volume (\\$)")
    axes[0].set_title("Panel A: Daily Trading Volume")

    if spreads is not None:
        axes[1].plot(spreads.index, spreads.values, color=JOF_COLORS["secondary"], linewidth=0.8)
        axes[1].set_ylabel("Estimated Spread")
        axes[1].set_title("Panel B: Bid--Ask Spread (Roll Estimator)")

    axes[-1].set_xlabel("Date")
    axes[-1].xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    fig.autofmt_xdate(rotation=45)
    fig.tight_layout()
    return fig


def fig_election_case_study(
    trump_prices: pd.Series,
    spy_returns: pd.Series | None = None,
) -> plt.Figure:
    """Figure: 2024 Presidential Election case study with event annotations."""
    set_jof_style()
    has_spy = spy_returns is not None
    fig, ax1 = plt.subplots(figsize=(8, 4.5))

    ax1.plot(trump_prices.index, trump_prices.values, color=JOF_COLORS["primary"],
             linewidth=1.2, label="Trump Win Probability")
    ax1.set_xlabel("Date")
    ax1.set_ylabel("Probability", color=JOF_COLORS["primary"])

    if has_spy:
        ax2 = ax1.twinx()
        cum_ret = (1 + spy_returns).cumprod() - 1
        common_idx = cum_ret.index.intersection(trump_prices.index)
        ax2.plot(common_idx, cum_ret.loc[common_idx].values, color=JOF_COLORS["secondary"],
                 linewidth=0.8, alpha=0.7, label="SPY Cumulative Return")
        ax2.set_ylabel("SPY Cumulative Return", color=JOF_COLORS["secondary"])

    for label, d in ELECTION_KEY_DATES.items():
        dt = pd.Timestamp(d)
        if dt in trump_prices.index or (trump_prices.index.min() <= dt <= trump_prices.index.max()):
            ax1.axvline(dt, color=JOF_COLORS["gray"], linewidth=0.8, linestyle=":")
            ypos = ax1.get_ylim()[1] * 0.95
            ax1.annotate(label, xy=(dt, ypos), fontsize=7, rotation=90,
                         ha="right", va="top")

    ax1.set_title("2024 US Presidential Election: Price Discovery on Polymarket")
    ax1.legend(loc="upper left")
    if has_spy:
        ax2.legend(loc="upper right")
    ax1.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    fig.autofmt_xdate(rotation=45)
    fig.tight_layout()
    return fig


def fig_arbitrage_violations(violations: pd.DataFrame) -> plt.Figure:
    """Figure: Arbitrage violations over time."""
    set_jof_style()
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))

    if "violation_size" in violations.columns:
        ax1.plot(violations.index, violations["violation_size"], color=JOF_COLORS["primary"],
                 linewidth=0.8)
        ax1.axhline(0, color="black", linewidth=0.5)
        ax1.set_xlabel("Date")
        ax1.set_ylabel("State Price Sum $-$ 1.0")
        ax1.set_title("Panel A: Completeness Deviation Over Time")
        ax1.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))

        ax2.hist(violations["violation_size"].dropna(), bins=30,
                 color=JOF_COLORS["primary"], alpha=0.7, edgecolor="white")
        ax2.axvline(0, color=JOF_COLORS["secondary"], linewidth=0.8, linestyle="--")
        ax2.set_xlabel("Deviation from Completeness")
        ax2.set_ylabel("Frequency")
        ax2.set_title("Panel B: Distribution of Violations")

    fig.autofmt_xdate(rotation=45)
    fig.tight_layout()
    return fig