#!/usr/bin/env python3
"""
Download VIX futures term-structure and VIX spot history **directly from CBOE**
(via the ``vix_utils`` wrapper) and build a flat daily panel with contango /
backwardation signals.

Output::

    RenTech/data/vix_futures_cboe.parquet

Columns
-------
trade_date          datetime64[ns]   (index)
vix_spot            float64          VIX cash close
vx1_settle          float64          front-month VX settle
vx2_settle          float64          second-month VX settle
vx3_settle          float64          third-month VX settle (~3M tenor)
vx1_close           float64          front-month VX close
vx2_close           float64          second-month VX close
vx1_expiry          datetime64[ns]   front-month expiry
vx2_expiry          datetime64[ns]   second-month expiry
vx1_dte             int              calendar days to front-month expiry
vx2_dte             int              calendar days to second-month expiry
vx1_volume          float64          front-month total volume
vx2_volume          float64          second-month total volume
vx1_oi              float64          front-month open interest
vx2_oi              float64          second-month open interest
contango_ratio      float64          vx2/vx1 − 1  (>0 = contango)
contango_flag       int8             1 contango, −1 backwardation, 0 flat
basis_pct           float64          (vx1 − vix_spot) / vix_spot
roll_yield_daily    float64          −(vx1 − vix_spot) / vx1_dte
vix3m               float64          CBOE VIX3M close (if available)
vvix                float64          CBOE VVIX close (if available)

Usage::

    python RenTech/data_pipeline/download_cboe_vix_futures.py
    python RenTech/data_pipeline/download_cboe_vix_futures.py --force-reload --start 2010-01-01
"""

from __future__ import annotations

import argparse
from datetime import date
from pathlib import Path

import numpy as np
import pandas as pd

DATA_DIR = Path(__file__).resolve().parents[1] / "data"
OUT_PATH = DATA_DIR / "vix_futures_cboe.parquet"
CONTRACTS_PATH = DATA_DIR / "vix_futures_contracts_long.parquet"


def _load_futures(force: bool) -> pd.DataFrame:
    """Load raw VIX futures (skip ``select_monthly_futures`` which drops post-2013 data)."""
    import vix_utils

    raw = vix_utils.load_vix_term_structure(forceReload=force)
    raw["Trade Date"] = pd.to_datetime(raw["Trade Date"]).dt.tz_localize(None)
    return raw


def _build_vx_series(raw: pd.DataFrame, tenor: int, col: str) -> pd.Series:
    """Extract a single price series for a given monthly tenor from raw data."""
    mask = raw["Tenor_Monthly"] == tenor
    sub = raw.loc[mask, ["Trade Date", col]].copy()
    sub = sub.dropna(subset=[col])
    sub = sub.sort_values("Trade Date")
    sub = sub.drop_duplicates(subset=["Trade Date"], keep="last")
    return sub.set_index("Trade Date")[col]


def _build_vx_field(raw: pd.DataFrame, tenor: int, col: str) -> pd.Series:
    mask = raw["Tenor_Monthly"] == tenor
    sub = raw.loc[mask, ["Trade Date", col]].copy()
    sub = sub.sort_values("Trade Date")
    sub = sub.drop_duplicates(subset=["Trade Date"], keep="last")
    return sub.set_index("Trade Date")[col]


def _load_cash(force: bool) -> pd.DataFrame:
    import vix_utils

    cash = vix_utils.get_vix_index_histories()
    cash["Trade Date"] = pd.to_datetime(cash["Trade Date"]).dt.tz_localize(None)
    return cash


def build(force_reload: bool = False, start: str | None = None) -> pd.DataFrame:
    print("Downloading VIX futures from CBOE …", flush=True)
    monthly = _load_futures(force_reload)
    print(f"  futures rows (raw): {len(monthly):,}", flush=True)

    print("Downloading VIX cash indices from CBOE …", flush=True)
    cash = _load_cash(force_reload)

    vx1_settle = _build_vx_series(monthly, 1, "Settle").rename("vx1_settle")
    vx2_settle = _build_vx_series(monthly, 2, "Settle").rename("vx2_settle")
    vx1_close = _build_vx_series(monthly, 1, "Close").rename("vx1_close")
    vx2_close = _build_vx_series(monthly, 2, "Close").rename("vx2_close")
    vx1_expiry = _build_vx_field(monthly, 1, "Expiry").rename("vx1_expiry")
    vx2_expiry = _build_vx_field(monthly, 2, "Expiry").rename("vx2_expiry")
    vx1_dte = _build_vx_field(monthly, 1, "Tenor_Days").rename("vx1_dte")
    vx2_dte = _build_vx_field(monthly, 2, "Tenor_Days").rename("vx2_dte")
    vx1_vol = _build_vx_field(monthly, 1, "Total Volume").rename("vx1_volume")
    vx2_vol = _build_vx_field(monthly, 2, "Total Volume").rename("vx2_volume")
    vx1_oi = _build_vx_field(monthly, 1, "Open Interest").rename("vx1_oi")
    vx2_oi = _build_vx_field(monthly, 2, "Open Interest").rename("vx2_oi")
    vx3_settle = _build_vx_series(monthly, 3, "Settle").rename("vx3_settle")
    vx3_close = _build_vx_series(monthly, 3, "Close").rename("vx3_close")
    vx3_dte = _build_vx_field(monthly, 3, "Tenor_Days").rename("vx3_dte")
    vx3_expiry = _build_vx_field(monthly, 3, "Expiry").rename("vx3_expiry")
    vx4_expiry = _build_vx_field(monthly, 4, "Expiry").rename("vx4_expiry")
    vx5_expiry = _build_vx_field(monthly, 5, "Expiry").rename("vx5_expiry")
    vx4_settle = _build_vx_series(monthly, 4, "Settle").rename("vx4_settle")
    vx4_close = _build_vx_series(monthly, 4, "Close").rename("vx4_close")
    vx5_settle = _build_vx_series(monthly, 5, "Settle").rename("vx5_settle")
    vx5_close = _build_vx_series(monthly, 5, "Close").rename("vx5_close")

    # Use VIX spot dates as the base (full daily coverage) and left-join futures
    vix_spot = cash.loc[cash["Symbol"] == "VIX", ["Trade Date", "Close"]].copy()
    vix_spot = vix_spot.drop_duplicates(subset=["Trade Date"], keep="last")
    vix_spot = vix_spot.set_index("Trade Date")["Close"].rename("vix_spot")

    panel = vix_spot.to_frame()
    panel.index.name = "trade_date"
    for s in [vx1_settle, vx2_settle, vx3_settle, vx4_settle, vx5_settle,
              vx1_close, vx2_close, vx3_close, vx4_close, vx5_close,
              vx1_expiry, vx2_expiry, vx3_expiry, vx4_expiry, vx5_expiry,
              vx1_dte, vx2_dte, vx3_dte,
              vx1_vol, vx2_vol, vx1_oi, vx2_oi]:
        panel = panel.join(s, how="left")

    vix3m = cash.loc[cash["Symbol"] == "VIX3M", ["Trade Date", "Close"]].copy()
    vix3m = vix3m.drop_duplicates(subset=["Trade Date"], keep="last")
    vix3m = vix3m.set_index("Trade Date")["Close"].rename("vix3m")
    panel = panel.join(vix3m, how="left")

    vvix = cash.loc[cash["Symbol"] == "VVIX", ["Trade Date", "Close"]].copy()
    vvix = vvix.drop_duplicates(subset=["Trade Date"], keep="last")
    vvix = vvix.set_index("Trade Date")["Close"].rename("vvix")
    panel = panel.join(vvix, how="left")

    # ---- derived signals ----
    s1 = panel["vx1_settle"].fillna(panel["vx1_close"])
    s2 = panel["vx2_settle"].fillna(panel["vx2_close"])
    s3 = panel["vx3_settle"].fillna(panel["vx3_close"])
    s4 = panel["vx4_settle"].fillna(panel["vx4_close"])
    s5 = panel["vx5_settle"].fillna(panel["vx5_close"])

    panel["vx1"] = s1
    panel["vx2"] = s2
    panel["vx3"] = s3
    panel["vx4"] = s4
    panel["vx5"] = s5

    # Calendar spreads in VIX points (near − far): e.g. vx1=18, vx2=20 → spread_m1_m2 = −2
    panel["spread_m1_m2"] = s1 - s2
    panel["spread_m2_m3"] = s2 - s3
    panel["spread_m3_m4"] = s3 - s4
    panel["spread_m4_m5"] = s4 - s5

    panel["contango_ratio"] = (s2 / s1) - 1.0
    # Front-two-month roll cost (e.g. 0.09 ≈ 9% between M1 and M2).
    panel["roll_cost_m1_m2"] = panel["contango_ratio"]
    # VX1 vs 60-day trend (futures level, not VIX cash).
    panel["vx1_sma60"] = panel["vx1"].rolling(60, min_periods=30).mean()
    panel["vx1_elev_vs_60d"] = (panel["vx1"] / panel["vx1_sma60"]) - 1.0
    # VX3/VX1 − 1: ~3-month vs ~1-month **futures** (not VIX3M cash index).
    panel["vx3_vx1_ratio"] = (s3 / s1) - 1.0
    panel["vx3_vx1_ratio_ffill"] = panel["vx3_vx1_ratio"].ffill()
    panel["contango_flag"] = np.where(
        panel["contango_ratio"] > 0.005, 1,
        np.where(panel["contango_ratio"] < -0.005, -1, 0),
    ).astype(np.int8)

    vx1_price = s1
    panel["basis_pct"] = (vx1_price - panel["vix_spot"]) / panel["vix_spot"]

    dte = panel["vx1_dte"].replace(0, np.nan)
    panel["roll_yield_daily"] = -(vx1_price - panel["vix_spot"]) / dte

    # VIX3M / VIX ratio — gap-free contango proxy (available whenever both cash
    # indices are published, independent of futures settlement coverage).
    panel["vix3m_vix_ratio"] = panel["vix3m"] / panel["vix_spot"]

    # Forward-fill contango_ratio / contango_flag across gaps in VX1 data so
    # strategies relying on the regime label don't lose signal on missing days.
    panel["contango_ratio_ffill"] = panel["contango_ratio"].ffill()
    panel["contango_flag_ffill"] = panel["contango_flag"].replace(0, np.nan).ffill().fillna(0).astype(np.int8)

    panel = panel.sort_index()
    if start:
        panel = panel.loc[start:]

    return panel, monthly


def build_contracts_long(raw: pd.DataFrame) -> pd.DataFrame:
    """One row per (trade_date, expiry) with settle/close price."""
    df = raw.copy()
    df["trade_date"] = pd.to_datetime(df["Trade Date"]).dt.tz_localize(None)
    df["expiry"] = pd.to_datetime(df["Expiry"], errors="coerce").dt.tz_localize(None)
    df["price"] = df["Settle"].fillna(df["Close"]).astype(np.float64)
    out = df[
        ["trade_date", "expiry", "Tenor_Monthly", "Tenor_Days", "price"]
    ].dropna(subset=["price", "expiry"])
    out = out.sort_values(["trade_date", "expiry"])
    out = out.drop_duplicates(subset=["trade_date", "expiry"], keep="last")
    return out.reset_index(drop=True)


def main() -> None:
    ap = argparse.ArgumentParser(description="Download VIX futures from CBOE and build contango panel.")
    ap.add_argument("--force-reload", action="store_true", help="Re-download from CBOE even if cached.")
    ap.add_argument("--start", type=str, default=None, help="Trim output to YYYY-MM-DD start date.")
    ap.add_argument("--out", type=str, default=str(OUT_PATH), help=f"Output path (default {OUT_PATH}).")
    args = ap.parse_args()

    panel, monthly = build(force_reload=args.force_reload, start=args.start)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    panel.to_parquet(out)

    contracts = build_contracts_long(monthly)
    contracts.to_parquet(CONTRACTS_PATH, index=False)

    print(f"\nSaved {len(panel):,} rows  ->  {out}")
    print(f"Saved {len(contracts):,} contract rows -> {CONTRACTS_PATH}")
    print(f"Date range: {panel.index.min().date()} .. {panel.index.max().date()}")
    print(f"Contango days:       {(panel['contango_flag'] == 1).sum():,}")
    print(f"Backwardation days:  {(panel['contango_flag'] == -1).sum():,}")
    print(f"Flat/ambiguous:      {(panel['contango_flag'] == 0).sum():,}")
    print()
    print(panel[["vix_spot", "vx1_settle", "vx2_settle", "contango_ratio", "basis_pct"]].describe().round(4))


if __name__ == "__main__":
    main()
