"""Chart generation for the MAX strategy (saved as PNGs under results/)."""
from __future__ import annotations

from pathlib import Path

import matplotlib

matplotlib.use("Agg")  # headless / reproducible
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from .backtest import equity_curve


def plot_quintile_returns(
    quintile_returns: pd.DataFrame, out_path: Path, *, title_suffix: str = ""
) -> None:
    """Bar chart of annualised mean return by MAX quintile (Q1 low ... Qn high)."""
    ann = quintile_returns.mean() * 12.0
    fig, ax = plt.subplots(figsize=(8, 5))
    colors = ["#2c7fb8"] * len(ann)
    colors[0] = "#238b45"          # long leg (low MAX)
    colors[-1] = "#cb181d"         # short leg (high MAX)
    ax.bar(ann.index, ann.values * 100, color=colors, edgecolor="black", linewidth=0.6)
    ax.axhline(0, color="black", linewidth=0.8)
    ax.set_ylabel("Annualised mean return (%)")
    ax.set_xlabel("MAX quintile  (Q1 = lowest MAX / long,  Q5 = highest MAX / short)")
    ax.set_title(f"Returns by MAX quintile{title_suffix}")
    for i, v in enumerate(ann.values * 100):
        ax.text(i, v + (0.3 if v >= 0 else -0.6), f"{v:.1f}%", ha="center", fontsize=9)
    fig.tight_layout()
    fig.savefig(out_path, dpi=130)
    plt.close(fig)


def plot_long_short_equity(
    gross: pd.Series, net: pd.Series, out_path: Path, *, oos_start=None
) -> None:
    """Cumulative growth of $1 for the long-short strategy (gross vs net)."""
    fig, ax = plt.subplots(figsize=(10, 5.5))
    eq_gross = equity_curve(gross)
    eq_net = equity_curve(net)
    ax.plot(eq_gross.index, eq_gross.values, label="Long-short (gross)", color="#74a9cf", lw=1.6)
    ax.plot(eq_net.index, eq_net.values, label="Long-short (net of costs)", color="#08519c", lw=1.8)
    ax.axhline(1.0, color="grey", lw=0.8, ls="--")
    if oos_start is not None:
        ax.axvline(pd.Timestamp(oos_start), color="firebrick", lw=1.2, ls=":")
        ax.text(pd.Timestamp(oos_start), ax.get_ylim()[1], "  OOS start",
                color="firebrick", va="top", fontsize=9)
    ax.set_ylabel("Growth of $1")
    ax.set_xlabel("Date")
    ax.set_title("Low-MAX minus High-MAX long-short equity curve")
    ax.legend(loc="upper left")
    fig.tight_layout()
    fig.savefig(out_path, dpi=130)
    plt.close(fig)


def plot_max_ivol_scatter(
    max_panel: pd.DataFrame, ivol_panel: pd.DataFrame, out_path: Path
) -> None:
    """Hexbin of pooled MAX vs idiosyncratic volatility."""
    common_idx = max_panel.index.intersection(ivol_panel.index)
    common_cols = max_panel.columns.intersection(ivol_panel.columns)
    m = max_panel.loc[common_idx, common_cols].to_numpy().ravel()
    v = ivol_panel.loc[common_idx, common_cols].to_numpy().ravel()
    mask = np.isfinite(m) & np.isfinite(v)
    m, v = m[mask] * 100, v[mask] * 100

    fig, ax = plt.subplots(figsize=(7.5, 6))
    hb = ax.hexbin(v, m, gridsize=45, cmap="viridis", mincnt=1, bins="log")
    fig.colorbar(hb, ax=ax, label="log10(count)")
    if len(m) > 2:
        b1, b0 = np.polyfit(v, m, 1)
        xs = np.linspace(v.min(), v.max(), 50)
        ax.plot(xs, b0 + b1 * xs, color="red", lw=1.8,
                label=f"OLS slope = {b1:.2f}")
        ax.legend(loc="upper left")
    ax.set_xlabel("Idiosyncratic volatility (daily, %)")
    ax.set_ylabel("MAX (avg of top-N daily returns, %)")
    ax.set_title("MAX is largely a repackaging of idiosyncratic volatility")
    fig.tight_layout()
    fig.savefig(out_path, dpi=130)
    plt.close(fig)


def plot_factor_loadings(reg_summary: dict, out_path: Path) -> None:
    """Bar chart of factor loadings with the alpha highlighted."""
    table = reg_summary["table"]
    fig, ax = plt.subplots(figsize=(8, 5))
    coefs = table["coef"]
    colors = ["#08519c" if idx == "alpha" else "#969696" for idx in coefs.index]
    ax.bar(coefs.index, coefs.values, color=colors, edgecolor="black", linewidth=0.6)
    ax.axhline(0, color="black", lw=0.8)
    ax.set_ylabel("Coefficient (monthly)")
    ax.set_title("Factor loadings of the long-short strategy "
                 f"(alpha t = {reg_summary['alpha_tstat']:.2f})")
    for i, (idx, v) in enumerate(coefs.items()):
        t = table.loc[idx, "t_stat"]
        ax.text(i, v + (0.0005 if v >= 0 else -0.0010),
                f"t={t:.1f}", ha="center", fontsize=8)
    fig.tight_layout()
    fig.savefig(out_path, dpi=130)
    plt.close(fig)
