#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 12 — Supplementary figures (not referenced by the paper).

The working paper is tables-only by design; these figures are reproducible
visual companions built from the result CSVs, useful for talks and quick
inspection. Colors follow a CVD-validated categorical palette; every series
is direct-labeled so identity never relies on color alone.

Inputs : results/rolling_r2.csv, results/irf_results.csv,
         results/portfolio_sort_results.csv, results/subperiod_results.csv
Outputs: figures/fig_rolling_r2.(png|pdf), figures/fig_irf.(png|pdf),
         figures/fig_portfolio_sorts.(png|pdf), figures/fig_subperiod.(png|pdf)
"""

import warnings

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import pandas as pd

import _bootstrap  # noqa: F401
from wp7 import config

warnings.filterwarnings('ignore')

# CVD-validated categorical palette (see AUDIT trail: adjacent-pair ΔE ≥ 8)
BLUE, ORANGE, AQUA, YELLOW = "#2a78d6", "#eb6834", "#1baf7a", "#eda100"
INK, INK_2 = "#0b0b0b", "#52514e"

plt.rcParams.update({
    "figure.dpi": 150,
    "font.size": 9,
    "axes.edgecolor": INK_2,
    "axes.labelcolor": INK,
    "axes.titlecolor": INK,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.grid": True,
    "grid.color": "#e5e4e0",
    "grid.linewidth": 0.6,
    "xtick.color": INK_2,
    "ytick.color": INK_2,
    "savefig.facecolor": "white",
    "axes.facecolor": "white",
})


def _save(fig, name: str) -> None:
    for ext in ("png", "pdf"):
        fig.savefig(config.FIGURES_DIR / f"{name}.{ext}", bbox_inches="tight")
    plt.close(fig)
    print(f"  saved figures/{name}.png|.pdf")


def fig_rolling_r2() -> None:
    """252-day rolling explanatory power of option-implied information."""
    df = pd.read_csv(config.RESULTS_DIR / "rolling_r2.csv", parse_dates=["date"])
    fig, ax = plt.subplots(figsize=(7.2, 3.4))
    series = [("1D_RV", "1-day RV forecasting (HAR+IV)", BLUE),
              ("5D_Return", "5-day return predictability", ORANGE)]
    for key, label, color in series:
        sub = df[df["target"] == key].sort_values("date")
        ax.plot(sub["date"], sub["r2"], color=color, lw=2)
        ax.annotate(label, xy=(sub["date"].iloc[-1], sub["r2"].iloc[-1]),
                    xytext=(6, 0), textcoords="offset points",
                    va="center", color=INK, fontsize=8.5)
    ax.set_ylabel("Rolling $R^2$ (252-day window)")
    ax.set_title("Option-implied information content through time", loc="left")
    ax.margins(x=0.02)
    fig.subplots_adjust(right=0.72)
    _save(fig, "fig_rolling_r2")


def _dodge(values: list, min_gap: float) -> list:
    """Push label anchor positions apart until every pair clears min_gap."""
    order = sorted(range(len(values)), key=lambda i: values[i])
    adjusted = list(values)
    for prev, cur in zip(order, order[1:]):
        if adjusted[cur] - adjusted[prev] < min_gap:
            adjusted[cur] = adjusted[prev] + min_gap
    return adjusted


def fig_irf() -> None:
    """Average VAR(5) impulse responses (20 tickers, standardized units)."""
    df = pd.read_csv(config.RESULTS_DIR / "irf_results.csv")
    avg = df.groupby("horizon")[["iv_to_iv", "rv_to_iv", "iv_to_rv", "rv_to_rv"]].mean()
    fig, ax = plt.subplots(figsize=(6.4, 3.6))
    series = [("iv_to_rv", "RV response to IV shock", BLUE),
              ("rv_to_rv", "RV response to RV shock", ORANGE),
              ("iv_to_iv", "IV response to IV shock", AQUA),
              ("rv_to_iv", "IV response to RV shock", YELLOW)]
    ends = _dodge([avg[key].iloc[-1] for key, _, _ in series], min_gap=0.07)
    for (key, label, color), label_y in zip(series, ends):
        ax.plot(avg.index, avg[key], color=color, lw=2)
        ax.annotate(label, xy=(avg.index[-1], label_y),
                    xytext=(6, 0), textcoords="offset points",
                    va="center", color=INK, fontsize=8.5)
    ax.axhline(0, color=INK_2, lw=0.8)
    ax.set_xlabel("Horizon (days)")
    ax.set_ylabel("Response (SD units)")
    ax.set_title("Average impulse responses, bivariate VAR(5): ATM IV ↔ RV",
                 loc="left")
    ax.margins(x=0.02)
    fig.subplots_adjust(right=0.70)
    _save(fig, "fig_irf")


def fig_portfolio_sorts() -> None:
    """Annualized long-short (Q5−Q1) returns of the 5-day sorts."""
    df = pd.read_csv(config.RESULTS_DIR / "portfolio_sort_results.csv")
    ls = df[(df["quintile"] == "L/S(5-1)") & (df["return_horizon"] == "5-Day")]
    ls = ls.sort_values("annualized_return_pct")
    colors = [BLUE if v >= 0 else ORANGE for v in ls["annualized_return_pct"]]

    fig, ax = plt.subplots(figsize=(6.4, 3.2))
    bars = ax.barh(ls["sort_variable"], ls["annualized_return_pct"],
                   color=colors, height=0.62)
    for bar, val, sharpe in zip(bars, ls["annualized_return_pct"], ls["sharpe_ratio"]):
        ax.annotate(f"{val:+.1f}%  (SR {sharpe:.2f})",
                    xy=(val, bar.get_y() + bar.get_height() / 2),
                    xytext=(5 if val >= 0 else -5, 0), textcoords="offset points",
                    va="center", ha="left" if val >= 0 else "right",
                    color=INK, fontsize=8)
    ax.axvline(0, color=INK_2, lw=0.8)
    ax.set_xlabel("Annualized L/S return (%)")
    ax.set_title("Long-short quintile portfolios, 5-day returns (Q5 − Q1)",
                 loc="left")
    ax.grid(axis="y", visible=False)
    xmin, xmax = ax.get_xlim()
    ax.set_xlim(xmin - 0.42 * (xmax - xmin), xmax + 0.18 * (xmax - xmin))
    _save(fig, "fig_portfolio_sorts")


def fig_subperiod() -> None:
    """HAR-RV vs HAR+IV in-sample R² across the nine subperiods."""
    df = pd.read_csv(config.RESULTS_DIR / "subperiod_results.csv")
    rv = df[df["model"].isin(["HAR-RV", "HAR+IV"])]
    pivot = rv.pivot_table(index="subperiod", columns="model", values="r2")
    pivot = pivot.reindex(list(config.SUBPERIODS))

    fig, ax = plt.subplots(figsize=(7.2, 3.6))
    x = range(len(pivot))
    w = 0.38
    ax.bar([i - w / 2 for i in x], pivot["HAR-RV"], width=w - 0.04,
           color=ORANGE, label="HAR-RV")
    ax.bar([i + w / 2 for i in x], pivot["HAR+IV"], width=w - 0.04,
           color=BLUE, label="HAR-RV + IV surface")
    ax.set_xticks(list(x))
    ax.set_xticklabels([s.split(" (")[0] for s in pivot.index],
                       rotation=30, ha="right", fontsize=8)
    ax.set_ylabel("In-sample $R^2$ (1-day RV)")
    ax.set_title("Stability of the IV-surface improvement across subperiods",
                 loc="left")
    ax.grid(axis="x", visible=False)
    ax.legend(frameon=False, fontsize=8.5, loc="upper left")
    _save(fig, "fig_subperiod")


def main():
    print("=" * 70)
    print("SUPPLEMENTARY FIGURES")
    print("=" * 70)
    config.ensure_output_dirs()
    fig_rolling_r2()
    fig_irf()
    fig_portfolio_sorts()
    fig_subperiod()
    print("\nFIGURES COMPLETE.")


if __name__ == "__main__":
    main()