"""
Augment the base research ``panel`` (from ``build_panel`` in
``research_literature_theta_strategies``) with **additional columns** used by the diverse
strategy catalog.

Rationale: the legacy panel already carries ``close``, ``rv21``, ``rv5``, ``sma_50`` (and
``sma_200`` / VIX columns from ``normalize_spy_df``). This module adds **cross-cutting**
features (rolling highs, Bollinger bandwidth proxy, VIX path features, calendar flags)
so individual strategies do not each recompute rolling windows in ad hoc ways.

**Important:** this is *not* strategy logic — it is deterministic feature engineering on
the SPY/VIX time series aligned to Theta session dates.
"""
from __future__ import annotations

from datetime import date

import numpy as np
import pandas as pd


def _third_friday(year: int, month: int) -> date:
    """Return the calendar date of the equity options-expiry Friday (3rd Friday)."""
    c = date(year, month, 1)
    # weekday: Mon=0 ... Sun=6; we need first Friday then +14 days
    days_until_fri = (4 - c.weekday()) % 7
    first_fri = date(year, month, 1 + days_until_fri)
    return date(year, month, first_fri.day + 14)


def _is_op_exp_week(d: pd.Timestamp) -> bool:
    """True if ``d`` lies in the Monday–Friday week containing the 3rd Friday of that month."""
    y, m = int(d.year), int(d.month)
    exp = _third_friday(y, m)
    dd = d.date()
    # Monday of expiry week
    exp_wd = exp.weekday()
    mon = date(y, m, exp.day - exp_wd)
    fri = date(mon.year, mon.month, mon.day + 4)
    return mon <= dd <= fri


def augment_research_panel(panel: pd.DataFrame) -> pd.DataFrame:
    """
    Return a copy of ``panel`` with extra float/bool columns. Existing columns are preserved.

    Added columns (non-exhaustive documentation of intent):

    - ``sma_20``: 20-session simple moving average of ``close``.
    - ``roll_high_20`` / ``roll_low_20``: trailing max/min of **close** over 20 sessions
      (close-only proxy for Donchian-style bounds; not true high/low bars).
    - ``roll_high_55`` / ``roll_low_55``: 55-session window (intermediate trend channel).
    - ``roll_high_252`` / ``roll_low_252``: ~1y trading window on **close** extrema.
    - ``dist_roll_high_252``: ``close / roll_high_252 - 1`` (negative if off highs).
    - ``bb_mid_20``, ``bb_up_20``, ``bb_lo_20``: classic Bollinger on ``close``, k=2.
    - ``bb_width_20``: ``(bb_up_20 - bb_lo_20) / bb_mid_20`` with guard for tiny mid.
    - ``bb_width_pct_rank_252``: percentile rank of ``bb_width_20`` vs its prior 252 rows.
    - ``ret_20``: 20-session simple return of ``close``.
    - ``vix_chg_5``: ``vix_close - vix_close.shift(5)`` (5-session change, not calendar days).
    - ``vix_roll_max_20``: trailing max of ``vix_close`` over 20 sessions.
    - ``vix_pct_rank_252``: percentile rank of ``vix_close`` in trailing 252 sessions.
    - ``rv_ratio_5_21``: ``rv5 / rv21`` when ``rv21 > 0``.
    - ``is_op_exp_week``: boolean, week of monthly listed equity expiry (3rd Friday rule).
    - ``day_of_month``, ``month``: calendar decomposition for seasonality-style sleeves.
    - ``weekday``: ``0=Monday`` … ``6=Sunday`` (NY session calendar day from the index).
    """
    o = panel.copy()
    c = pd.to_numeric(o["close"], errors="coerce").astype(float)

    o["sma_20"] = c.rolling(20, min_periods=1).mean()
    o["roll_high_20"] = c.rolling(20, min_periods=5).max()
    o["roll_low_20"] = c.rolling(20, min_periods=5).min()
    o["roll_high_55"] = c.rolling(55, min_periods=20).max()
    o["roll_low_55"] = c.rolling(55, min_periods=20).min()
    o["roll_high_252"] = c.rolling(252, min_periods=50).max()
    o["roll_low_252"] = c.rolling(252, min_periods=50).min()
    o["dist_roll_high_252"] = np.where(
        o["roll_high_252"] > 0, c / o["roll_high_252"] - 1.0, np.nan
    )

    std20 = c.rolling(20, min_periods=5).std(ddof=0)
    mid = o["sma_20"]
    o["bb_mid_20"] = mid
    o["bb_up_20"] = mid + 2.0 * std20
    o["bb_lo_20"] = mid - 2.0 * std20
    denom = mid.replace(0, np.nan).abs()
    o["bb_width_20"] = (o["bb_up_20"] - o["bb_lo_20"]) / denom

    bw = pd.to_numeric(o["bb_width_20"], errors="coerce")
    o["bb_width_pct_rank_252"] = bw.rolling(252, min_periods=50).apply(
        lambda s: float(np.mean(s[:-1] <= s.iloc[-1])) if len(s) > 1 and pd.notna(s.iloc[-1]) else np.nan,
        raw=False,
    )

    o["ret_20"] = c / c.shift(20) - 1.0

    vx = pd.to_numeric(o["vix_close"], errors="coerce").astype(float)
    o["vix_chg_5"] = vx - vx.shift(5)
    o["vix_roll_max_20"] = vx.rolling(20, min_periods=5).max()
    o["vix_pct_rank_252"] = vx.rolling(252, min_periods=50).apply(
        lambda s: float(np.mean(s[:-1] <= s.iloc[-1])) if len(s) > 1 and pd.notna(s.iloc[-1]) else np.nan,
        raw=False,
    )

    rv5 = pd.to_numeric(o["rv5"], errors="coerce")
    rv21 = pd.to_numeric(o["rv21"], errors="coerce")
    o["rv_ratio_5_21"] = np.where(rv21 > 1e-12, rv5 / rv21, np.nan)

    idx = pd.DatetimeIndex(pd.to_datetime(o.index).normalize())
    o["is_op_exp_week"] = [bool(_is_op_exp_week(ts)) for ts in idx]
    o["day_of_month"] = [int(ts.day) for ts in idx]
    o["month"] = [int(ts.month) for ts in idx]
    o["weekday"] = [int(ts.weekday()) for ts in idx]  # Mon=0 ... Sun=6

    return o