#!/usr/bin/env python3
"""
Train a **LightGBM** implied-volatility surface model on historical SPY options.

Pipeline stages
---------------
1. **Ingest** — Pull daily option-chain snapshots from Massive (``api.massive.com``),
   paginate with ``next_url``, cache append-only **Parquet** (memory-safe day-by-day writes).
2. **Features** — Moneyness, annualized time-to-maturity, VIX level / 5d change, SPY 5d return,
   call/put flag (aligned to ``underlying_df`` index).
3. **Train** — Chronological split: first **18 months** train, next **6 months** validation
   (no shuffle).
4. **Evaluate** — Validation **RMSE**; **feature importance** plot (matplotlib) saved next to this file.
5. **Export** — ``joblib.dump`` → :file:`RenTech/ml_pipeline/iv_surface_model.pkl`.

Environment
-----------
``MASSIVE_API_KEY`` — required for live download (set in ``.env``; loaded via ``python-dotenv``).

Optional:

* ``MASSIVE_API_BASE`` — default ``https://api.massive.com``. Keys issued by **Polygon.io** normally
  use ``https://api.polygon.io`` (same ``/v3/snapshot/options/...`` paths). **403 Forbidden** from
  Massive’s edge usually means: wrong host for the key, or your **plan does not include** that options
  endpoint (check the dashboard product matrix).
* ``POLYGON_API_KEY`` — alias for ``MASSIVE_API_KEY`` if you store the key under Polygon’s name.
* ``MASSIVE_AS_OF_PARAM`` — query key for historical snapshot date if your plan documents one
  (e.g. ``as_of``); appended per trading day when set.
* ``MASSIVE_SEND_AUTH_HEADERS`` — default **off**: only ``apiKey`` is sent on the query string
  (Polygon’s documented pattern). Set to ``1`` / ``true`` if your provider requires
  ``Authorization`` / ``X-API-Key`` as well.

Flat files
----------
Massive publishes **options** flat files (S3 daily CSV) for very large history; this module focuses on
the REST snapshot path with pagination. For TB-scale backfills, prefer their flat-file catalog and
ingest into the same Parquet schema produced by :func:`normalize_massive_chain_response`.

---------------------------------------------------------------------------
SyntheticLoader integration (future / migration)
---------------------------------------------------------------------------
Today, :class:`~RenTech.core.synthetic_data_loader.SyntheticLoader` loads an **XGBoost** model from
:class:`~RenTech.core.iv_surface_model` with features
``[vix_level, dte, moneyness, is_call]`` (``dte`` in **calendar days**).

After training with this script, you can load the new artifact in-process::

    import joblib
    from pathlib import Path

    model = joblib.load(Path(__file__).resolve().parent / "iv_surface_model.pkl")

At inference time, build a DataFrame (or 2D array **in the same column order** as
:data:`FEATURE_COLUMNS`) with:

* ``moneyness`` = strike / spot
* ``time_to_maturity`` = dte_calendar_days / 365.25  (match training definition)
* ``vix_level``, ``vix_5d_change``, ``spy_5d_return`` from your live macro panel
* ``call_put_flag`` as categorical ``\"C\"`` / ``\"P\"`` (or the same codes used at train time)

Then ``model.predict(X)`` yields IV in **decimal** form (e.g. ``0.22`` = 22% vol), consistent with
:func:`~RenTech.core.iv_surface_model.synthetic_black_scholes`. Wire that into
``SyntheticOptionChain`` by replacing ``_predict_iv_matrix`` to use the wider feature set and this
pickle path once you deprecate the legacy XGBoost surface.

Dependencies
------------
``pandas``, ``numpy``, ``pyarrow``, ``requests``, ``python-dotenv``, ``scikit-learn``,
``lightgbm``, ``joblib``, ``matplotlib``, and (for macro panel) ``yfinance``.

CLI date window
---------------
With no ``--start`` / ``--end``, the script uses **today** as the end date and **two calendar years
earlier** as the start. Override either bound as needed.
"""

from __future__ import annotations

import json
import math
import os
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse

import numpy as np
import pandas as pd

# --- Repo root on sys.path -------------------------------------------------
_ML_DIR = Path(__file__).resolve().parent
_REPO_ROOT = _ML_DIR.parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

# ---------------------------------------------------------------------------
# Global paths & model I/O
# ---------------------------------------------------------------------------

DEFAULT_CACHE_PARQUET = _ML_DIR / "massive_spy_options_eod.parquet"
DEFAULT_MODEL_PATH = _ML_DIR / "iv_surface_model.pkl"
DEFAULT_FEATURE_IMPORTANCE_PNG = _ML_DIR / "iv_surface_lgbm_feature_importance.png"

# Massive REST base URL.
# Massive support notes often refer to `MASSIVE_API_BASE_URL`; keep `MASSIVE_API_BASE` as a fallback.
DEFAULT_MASSIVE_API_BASE = (
    os.environ.get("MASSIVE_API_BASE_URL")
    or os.environ.get("MASSIVE_API_BASE")
    or "https://api.massive.com"
).rstrip("/")

# Feature order must match inference code that loads iv_surface_model.pkl
FEATURE_COLUMNS: list[str] = [
    "moneyness",
    "time_to_maturity",
    "vix_level",
    "vix_5d_change",
    "spy_5d_return",
    "call_put_flag",
]
TARGET_COLUMN = "implied_volatility"


# ---------------------------------------------------------------------------
# .env → MASSIVE_API_KEY
# ---------------------------------------------------------------------------


def _load_dotenv() -> None:
    """Load ``.env`` from repo root and/or ``trading_bot_live/.env`` if ``python-dotenv`` is installed."""
    try:
        from dotenv import load_dotenv
    except ImportError:
        return
    for env_path in (_REPO_ROOT / ".env", _REPO_ROOT / "trading_bot_live" / ".env"):
        if env_path.is_file():
            load_dotenv(env_path, override=False)


def _get_api_key() -> str:
    _load_dotenv()
    key = (os.environ.get("MASSIVE_API_KEY") or os.environ.get("POLYGON_API_KEY") or "").strip()
    if not key:
        raise RuntimeError(
            "MASSIVE_API_KEY or POLYGON_API_KEY is not set. Add it to .env or export it before fetching."
        )
    return key


class MassiveHttpError(Exception):
    """Non-2xx from Massive/Polygon REST; carries status and provider message (no secrets)."""

    def __init__(self, status_code: int, redacted_url: str, detail: str) -> None:
        self.status_code = int(status_code)
        self.redacted_url = redacted_url
        self.detail = detail
        super().__init__(f"HTTP {self.status_code} {redacted_url} — {detail}")


def _redact_url_for_log(url: str) -> str:
    """Remove apiKey from URLs for logs and exceptions."""
    parsed = urlparse(url)
    q = [(k, v) for k, v in parse_qsl(parsed.query, keep_blank_values=True) if k.lower() != "apikey"]
    new_q = urlencode(q)
    return urlunparse((parsed.scheme, parsed.netloc, parsed.path, parsed.params, new_q, parsed.fragment))


def _http_error_detail(response: Any) -> str:
    """Best-effort JSON/text snippet from a failed requests response."""
    text = (getattr(response, "text", None) or "").strip()
    if not text:
        return "(empty body)"
    try:
        data = response.json()
    except Exception:
        return text[:500]
    if isinstance(data, dict):
        for k in ("message", "error", "detail", "status"):
            if k in data and data[k]:
                return str(data[k])[:500]
        return str(data)[:500]
    return text[:500]


def _massive_auth_mode() -> str:
    """
    Massive auth mode:

    - **header** (default): `Authorization: Bearer <key>` (Massive support guidance)
    - **query**: `?apiKey=<key>` (Polygon-style deployments)
    - **both**: send both header + query (rare, but supported as a fallback)
    """
    v = (os.environ.get("MASSIVE_AUTH_MODE") or "header").strip().lower()
    if v in ("header", "bearer"):
        return "header"
    if v in ("query", "apikey", "api_key"):
        return "query"
    if v in ("both", "all"):
        return "both"
    return "header"


# ---------------------------------------------------------------------------
# 1) Massive API — historical chains + Parquet cache
# ---------------------------------------------------------------------------


def _attach_api_key(url: str, api_key: str) -> str:
    """Append ``apiKey=`` to URL query string (Polygon-style pattern)."""
    parsed = urlparse(url)
    q = dict(parse_qsl(parsed.query, keep_blank_values=True))
    q["apiKey"] = api_key
    new_q = urlencode(q)
    return urlunparse((parsed.scheme, parsed.netloc, parsed.path, parsed.params, new_q, parsed.fragment))


def _session_headers() -> dict[str, str]:
    """
    Default: Massive support guidance is **Authorization: Bearer <key>**.

    Set ``MASSIVE_AUTH_MODE=query`` to send only ``apiKey`` on the query string.
    """
    h: dict[str, str] = {"Accept": "application/json", "User-Agent": "trading_bot-iv-train/1.0"}
    mode = _massive_auth_mode()
    if mode not in ("header", "both"):
        return h
    key = (
        os.environ.get("MASSIVE_API_KEY") or os.environ.get("POLYGON_API_KEY") or ""
    ).strip()
    if key:
        h["Authorization"] = f"Bearer {key}"
    return h


@dataclass
class MassiveClient:
    """Thin Massive REST client: one chain snapshot per request, paginated via ``next_url``."""

    api_base: str = DEFAULT_MASSIVE_API_BASE
    api_key: str = ""
    request_sleep_sec: float = 0.15

    def __post_init__(self) -> None:
        if not self.api_key:
            self.api_key = _get_api_key()

    def _full_url(self, path_or_url: str) -> str:
        mode = _massive_auth_mode()
        if path_or_url.startswith("http://") or path_or_url.startswith("https://"):
            return _attach_api_key(path_or_url, self.api_key) if mode in ("query", "both") else path_or_url
        path = path_or_url if path_or_url.startswith("/") else f"/{path_or_url}"
        url = f"{self.api_base}{path}"
        return _attach_api_key(url, self.api_key) if mode in ("query", "both") else url

    def get_json(self, path_or_url: str) -> dict[str, Any]:
        try:
            import requests
        except ImportError as e:
            raise ImportError("pip install requests") from e

        url = self._full_url(path_or_url)
        r = requests.get(url, headers=_session_headers(), timeout=120)
        if r.status_code >= 400:
            raise MassiveHttpError(
                r.status_code,
                _redact_url_for_log(url),
                _http_error_detail(r),
            )
        return r.json()


def iter_option_chain_pages(
    underlying: str,
    *,
    extra_params: dict[str, Any] | None = None,
    client: MassiveClient | None = None,
) -> Iterator[list[dict[str, Any]]]:
    """
    Yield each page of ``results`` from
    ``GET /v3/snapshot/options/{underlying}`` with optional query filters.

    Pagination: follow ``next_url`` until exhausted. Keeps only one page in memory at a time.
    """
    c = client or MassiveClient()
    params = dict(extra_params or {})
    q = urlencode({k: v for k, v in params.items() if v is not None})
    path = f"/v3/snapshot/options/{underlying.upper()}"
    if q:
        path = f"{path}?{q}"
    data = c.get_json(path)
    while True:
        results = data.get("results") or []
        if results:
            yield results
        nxt = data.get("next_url")
        if not nxt:
            break
        time.sleep(c.request_sleep_sec)
        data = c.get_json(nxt)
        if data.get("status") not in (None, "OK", "DELAYED"):
            break


def normalize_massive_chain_response(
    results: list[dict[str, Any]],
    as_of: pd.Timestamp,
    symbol: str,
) -> pd.DataFrame:
    """
    Flatten one API page of chain contracts into a typed, memory-tight DataFrame row set.

    IV normalization: Massive samples sometimes show ``implied_volatility`` as a percent (e.g. ``5``
    meaning 5%). We store **decimal** IV (0.05) when raw values look like percentages (typical if
    median > 1.5).
    """
    rows: list[dict[str, Any]] = []
    day = _norm_day(as_of)

    for item in results:
        det = item.get("details") or {}
        strike = det.get("strike_price")
        exp_raw = det.get("expiration_date")
        ctype = (det.get("contract_type") or "").lower()
        iv_raw = item.get("implied_volatility")
        ua = item.get("underlying_asset") or {}
        und_px = ua.get("price")
        day_bar = item.get("day") or {}
        oi = item.get("open_interest")

        lq = item.get("last_quote") or {}
        bid = lq.get("bid")
        ask = lq.get("ask")
        mid = lq.get("midpoint")

        rows.append(
            {
                "date": day,
                "underlying": str(symbol).upper(),
                "expiration": pd.to_datetime(exp_raw, errors="coerce"),
                "strike": float(strike) if strike is not None else np.nan,
                "option_type": "C" if ctype == "call" else ("P" if ctype == "put" else ""),
                "implied_volatility_raw": float(iv_raw) if iv_raw is not None else np.nan,
                "underlying_price": float(und_px) if und_px is not None else np.nan,
                "bid": float(bid) if bid is not None else np.nan,
                "ask": float(ask) if ask is not None else np.nan,
                "mid": float(mid) if mid is not None else np.nan,
                "volume": float(day_bar.get("volume")) if day_bar.get("volume") is not None else np.nan,
                "open_interest": float(oi) if oi is not None else np.nan,
                "ticker": str(det.get("ticker") or ""),
            }
        )

    if not rows:
        return pd.DataFrame()

    df = pd.DataFrame(rows)
    # IV → decimal
    iv = pd.to_numeric(df["implied_volatility_raw"], errors="coerce")
    med = float(iv.dropna().median()) if iv.notna().any() else 0.0
    if med > 1.5:
        iv = iv / 100.0
    df["implied_volatility"] = iv.astype(np.float32)
    df.drop(columns=["implied_volatility_raw"], inplace=True, errors="ignore")

    # dtypes
    df["strike"] = df["strike"].astype(np.float32)
    df["underlying_price"] = df["underlying_price"].astype(np.float32)
    for col in ("bid", "ask", "mid", "volume", "open_interest"):
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce").astype(np.float32)
    df["expiration"] = pd.to_datetime(df["expiration"], errors="coerce").dt.normalize()
    df["option_type"] = df["option_type"].astype("category")
    return df


def fetch_historical_chains(
    symbol: str,
    start_date: str | pd.Timestamp,
    end_date: str | pd.Timestamp,
    *,
    cache_path: str | Path = DEFAULT_CACHE_PARQUET,
    trading_days: pd.DatetimeIndex | None = None,
    as_of_param: str | None = None,
    client: MassiveClient | None = None,
) -> Path:
    """
    Download daily SPY (or ``symbol``) option snapshots from Massive and append to Parquet.

    * **Memory:** one trading day at a time; each day is normalized, filtered, written, then freed.
    * **Cache:** skips dates already present in ``cache_path`` (checks distinct ``date`` in Parquet).
    * **Historical date:** set env ``MASSIVE_AS_OF_PARAM`` (e.g. ``as_of``) *or* pass
      ``as_of_param``; value sent as YYYY-MM-DD per day. If unset, each call returns the provider's
      default (usually latest) — suitable for building today's file, not multi-year history.

    Parameters
    ----------
    symbol
        Underlying ticker, e.g. ``\"SPY\"``.
    start_date, end_date
        Inclusive calendar bounds; intersected with ``trading_days`` when provided.
    cache_path
        Append-only Parquet path (single file, row groups appended via rewrite concat for simplicity
        on small/medium caches; for huge files, partition by year externally).
    trading_days
        Optional index of session dates (e.g. SPY calendar). If ``None``, uses business days between
        start and end.
    as_of_param
        Query key for historical snapshot (provider-specific). Defaults to env
        ``MASSIVE_AS_OF_PARAM``.

    Returns
    -------
    Path
        Resolved ``cache_path``.
    """
    try:
        import pyarrow as pa
        import pyarrow.parquet as pq
    except ImportError as e:
        raise ImportError("pip install pyarrow") from e

    sym = symbol.upper()
    d0 = _norm_day(start_date)
    d1 = _norm_day(end_date)
    if d0 > d1:
        raise ValueError("start_date must be <= end_date")

    cache_path = Path(cache_path).expanduser()
    cache_path.parent.mkdir(parents=True, exist_ok=True)

    if trading_days is None:
        bdays = pd.bdate_range(d0, d1, freq="C")
    else:
        bdays = pd.DatetimeIndex(
            sorted({_norm_day(x) for x in trading_days if d0 <= _norm_day(x) <= d1})
        )

    existing: set[pd.Timestamp] = set()
    if cache_path.is_file():
        pf = pq.ParquetFile(cache_path)
        # Single column scan for dates already cached
        for rg in range(pf.num_row_groups):
            col = pf.read_row_group(rg, columns=["date"]).column(0)
            s = pd.to_datetime(col.to_pandas(), errors="coerce").dt.normalize()
            existing.update(s.dropna().unique())

    as_of_key = as_of_param or os.environ.get("MASSIVE_AS_OF_PARAM", "").strip() or None
    c = client or MassiveClient()

    new_chunks: list[pd.DataFrame] = []
    for day in bdays:
        if day in existing:
            continue
        extra: dict[str, Any] = {"limit": 250}
        if as_of_key:
            extra[as_of_key] = day.strftime("%Y-%m-%d")

        day_frames: list[pd.DataFrame] = []
        try:
            for page in iter_option_chain_pages(sym, extra_params=extra, client=c):
                part = normalize_massive_chain_response(page, day, sym)
                if not part.empty:
                    day_frames.append(part)
                time.sleep(c.request_sleep_sec)
        except MassiveHttpError as e:
            if e.status_code in (401, 403):
                raise RuntimeError(
                    "Massive/Polygon returned HTTP "
                    f"{e.status_code} (not a Python bug). Provider said: {e.detail!r}\n"
                    f"Request (apiKey redacted): {e.redacted_url}\n"
                    "Typical fixes:\n"
                    "  • Set MASSIVE_API_BASE=https://api.polygon.io if your key is from Polygon.io "
                    "(or use Massive’s host if the key was issued there).\n"
                    "  • Confirm your subscription includes **Options → Snapshot / option chain**.\n"
                    "  • Try MASSIVE_SEND_AUTH_HEADERS=1 only if their docs require header auth.\n"
                    "  • Do not paste your API key in logs or chat; rotate it if it leaked.\n"
                    "Use --from-parquet with local EOD data to train without this API."
                ) from e
            print(f"[fetch_historical_chains] skip {day.date()}: {e}", file=sys.stderr)
            continue
        except Exception as e:
            # Log and continue — partial history better than total failure
            print(f"[fetch_historical_chains] skip {day.date()}: {e}", file=sys.stderr)
            continue

        if not day_frames:
            continue
        day_df = pd.concat(day_frames, ignore_index=True)
        new_chunks.append(day_df)
        existing.add(day)

    if not new_chunks:
        if not cache_path.is_file():
            raise RuntimeError(
                "No data fetched and no cache file exists. "
                "Check MASSIVE_API_KEY, network, and whether historical as_of is configured."
            )
        return cache_path

    combined_new = pd.concat(new_chunks, ignore_index=True)

    if cache_path.is_file():
        old = pd.read_parquet(cache_path)
        out = pd.concat([old, combined_new], ignore_index=True)
    else:
        out = combined_new

    # Optimize dtypes before write
    out = _optimize_options_df_dtypes(out)
    out.to_parquet(cache_path, index=False, compression="zstd")
    return cache_path


# ---------------------------------------------------------------------------
# 2) Macro panel (SPY + VIX) for feature merge
# ---------------------------------------------------------------------------


def load_underlying_panel(start: str, end: str) -> pd.DataFrame:
    """
    SPY close + VIX close, indexed by normalized date.

    Columns: ``close`` (SPY), ``vix_close`` (VIX).
    """
    try:
        import yfinance as yf
    except ImportError as e:
        raise ImportError("pip install yfinance") from e

    spy = yf.download("SPY", start=start, end=end, progress=False, auto_adjust=False, threads=False)
    vix = yf.download("^VIX", start=start, end=end, progress=False, auto_adjust=False, threads=False)

    def _close(df: pd.DataFrame, label: str) -> pd.Series:
        if df is None or df.empty:
            raise ValueError(f"No rows for {label}")
        if isinstance(df.columns, pd.MultiIndex):
            lv0 = list(df.columns.get_level_values(0))
            col = df["Adj Close"] if "Adj Close" in lv0 else df["Close"]
        else:
            col = df["Adj Close"] if "Adj Close" in df.columns else df["Close"]
        ser = col.squeeze()
        if isinstance(ser, pd.DataFrame):
            ser = ser.iloc[:, 0]
        ser = pd.to_numeric(ser, errors="coerce").astype(np.float32)
        ser.index = pd.to_datetime(ser.index).normalize()
        return ser[~ser.index.duplicated(keep="last")].sort_index()

    out = pd.concat(
        [_close(spy, "SPY").rename("close"), _close(vix, "^VIX").rename("vix_close")],
        axis=1,
        join="outer",
    ).sort_index()
    out["vix_close"] = out["vix_close"].ffill().bfill()
    out = out.dropna(subset=["close"])
    out["close"] = out["close"].astype(np.float32)
    out["vix_close"] = pd.to_numeric(out["vix_close"], errors="coerce").astype(np.float32)
    return out


# ---------------------------------------------------------------------------
# 3) Feature engineering
# ---------------------------------------------------------------------------


def _norm_day(ts: pd.Timestamp | object) -> pd.Timestamp:
    return pd.Timestamp(ts).normalize()


def _optimize_options_df_dtypes(df: pd.DataFrame) -> pd.DataFrame:
    """Downcast wide options tables for RAM: float32 + categories."""
    out = df.copy()
    for col in ("strike", "underlying_price", "bid", "ask", "mid", "volume", "open_interest"):
        if col in out.columns:
            out[col] = pd.to_numeric(out[col], errors="coerce").astype(np.float32)
    if "implied_volatility" in out.columns:
        out["implied_volatility"] = pd.to_numeric(out["implied_volatility"], errors="coerce").astype(
            np.float32
        )
    if "option_type" in out.columns:
        out["option_type"] = out["option_type"].astype("category")
    return out


def engineer_features(
    raw_options_df: pd.DataFrame,
    underlying_df: pd.DataFrame,
) -> pd.DataFrame:
    """
    Join options with trailing macro features from ``underlying_df``.

    Adds / requires:

    * ``moneyness`` = strike / underlying_price (from chain row; falls back to merge-day SPY close).
    * ``time_to_maturity`` = (expiration - date) in **annualized years** (ACT/365.25 style).
    * ``vix_level`` — VIX close on ``date``.
    * ``vix_5d_change`` — VIX close minus VIX 5 sessions ago.
    * ``spy_5d_return`` — SPY close / lag5 - 1.
    * ``call_put_flag`` — ``\"C\"`` / ``\"P\"`` category.

    Target column ``implied_volatility`` is passed through when present.
    """
    opt = raw_options_df.copy()
    opt["date"] = pd.to_datetime(opt["date"], errors="coerce").dt.normalize()
    opt["expiration"] = pd.to_datetime(opt["expiration"], errors="coerce").dt.normalize()

    mkt = underlying_df.copy()
    mkt.index = pd.to_datetime(mkt.index).normalize()
    mkt = mkt[~mkt.index.duplicated(keep="last")].sort_index()

    if not {"close", "vix_close"}.issubset(mkt.columns):
        raise ValueError("underlying_df must have columns 'close' and 'vix_close'.")

    mkt = mkt.astype({"close": np.float32, "vix_close": np.float32})
    mkt["spy_5d_return"] = (mkt["close"] / mkt["close"].shift(5) - 1.0).astype(np.float32)
    mkt["vix_5d_change"] = (mkt["vix_close"] - mkt["vix_close"].shift(5)).astype(np.float32)
    mkt["vix_level"] = mkt["vix_close"]

    opt = opt.merge(
        mkt[["close", "vix_level", "vix_5d_change", "spy_5d_return"]],
        left_on="date",
        right_index=True,
        how="inner",
    )

    und = pd.to_numeric(opt["underlying_price"], errors="coerce")
    spot = und.where(und > 0, np.nan)
    spot = spot.fillna(pd.to_numeric(opt["close"], errors="coerce"))
    strike = pd.to_numeric(opt["strike"], errors="coerce")
    opt["moneyness"] = (strike / spot).astype(np.float32)

    dte_days = (opt["expiration"] - opt["date"]).dt.days.astype(np.float32)
    opt["time_to_maturity"] = (dte_days / np.float32(365.25)).astype(np.float32)

    ot = opt["option_type"].astype(str).str.upper().str[:1]
    opt["call_put_flag"] = np.where(ot == "C", "C", "P")
    opt["call_put_flag"] = opt["call_put_flag"].astype("category")

    keep = (
        np.isfinite(opt["moneyness"])
        & (opt["moneyness"] > 0)
        & np.isfinite(opt["time_to_maturity"])
        & (opt["time_to_maturity"] > 0)
        & np.isfinite(opt["implied_volatility"])
        & (opt["implied_volatility"] > 0)
        & np.isfinite(opt["vix_level"])
    )
    # Require 5d history for macro features
    keep &= np.isfinite(opt["spy_5d_return"]) & np.isfinite(opt["vix_5d_change"])
    opt = opt.loc[keep].copy()

    feat = opt[
        [
            "date",
            "moneyness",
            "time_to_maturity",
            "vix_level",
            "vix_5d_change",
            "spy_5d_return",
            "call_put_flag",
            TARGET_COLUMN,
        ]
    ].copy()

    for c in ("moneyness", "time_to_maturity", "vix_level", "vix_5d_change", "spy_5d_return"):
        feat[c] = feat[c].astype(np.float32)
    return feat.sort_values("date").reset_index(drop=True)


# ---------------------------------------------------------------------------
# 4) Chronological split + LightGBM train / eval / plot / save
# ---------------------------------------------------------------------------


def chronological_train_val_split(
    df: pd.DataFrame,
    *,
    train_months: int = 18,
    val_months: int = 6,
    date_col: str = "date",
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """
    First ``train_months`` calendar months → train; following ``val_months`` → validation.

    If the available data span is **shorter** than ``train_months + val_months``, fall back to a
    **proportional** time split (same 18:6 ratio on the observed ``[min_date, max_date]`` interval)
    so smoke tests on Parquet subsamples still run.
    """
    if df.empty:
        raise ValueError("Empty dataframe for split.")
    d = df.sort_values(date_col).reset_index(drop=True)
    t0 = d[date_col].min()
    t1 = d[date_col].max()
    t_train_end = t0 + pd.DateOffset(months=train_months)
    t_val_end = t_train_end + pd.DateOffset(months=val_months)

    if t_val_end <= t1:
        train = d[(d[date_col] >= t0) & (d[date_col] < t_train_end)]
        val = d[(d[date_col] >= t_train_end) & (d[date_col] < t_val_end)]
    else:
        total_m = max(train_months + val_months, 1)
        frac_tr = train_months / total_m
        span = t1 - t0
        if span.days < 2:
            raise ValueError("Need at least 2 distinct days for chronological split.")
        split_at = t0 + span * frac_tr
        train = d[d[date_col] < split_at]
        val = d[d[date_col] >= split_at]

    if train.empty or val.empty:
        raise ValueError(
            f"Split produced empty train or val (train={len(train)}, val={len(val)}). "
            "Provide a longer date range or more rows."
        )
    return train, val


def train_lgbm_iv_surface(
    feat_df: pd.DataFrame,
    *,
    model_path: str | Path = DEFAULT_MODEL_PATH,
    importance_png: str | Path = DEFAULT_FEATURE_IMPORTANCE_PNG,
    train_months: int = 18,
    val_months: int = 6,
    random_state: int = 42,
) -> dict[str, Any]:
    """
    Fit ``LGBMRegressor``, report validation RMSE, save model + feature-importance figure.
    """
    try:
        import joblib
        import lightgbm as lgb
        import matplotlib.pyplot as plt
        from sklearn.metrics import mean_squared_error
    except ImportError as e:
        raise ImportError("pip install lightgbm matplotlib scikit-learn joblib") from e

    train_df, val_df = chronological_train_val_split(
        feat_df, train_months=train_months, val_months=val_months
    )

    X_train = train_df[FEATURE_COLUMNS].copy()
    y_train = train_df[TARGET_COLUMN].astype(np.float32).values
    X_val = val_df[FEATURE_COLUMNS].copy()
    y_val = val_df[TARGET_COLUMN].astype(np.float32).values

    cat_cols = ["call_put_flag"]
    for X in (X_train, X_val):
        for c in cat_cols:
            if c in X.columns:
                X[c] = X[c].astype("category")

    model = lgb.LGBMRegressor(
        n_estimators=500,
        learning_rate=0.05,
        num_leaves=63,
        max_depth=-1,
        subsample=0.8,
        colsample_bytree=0.8,
        random_state=random_state,
        verbose=-1,
        force_col_wise=True,
    )
    model.fit(
        X_train,
        y_train,
        eval_set=[(X_val, y_val)],
        eval_metric="rmse",
        categorical_feature=["call_put_flag"],
    )

    pred_val = model.predict(X_val)
    rmse = float(math.sqrt(mean_squared_error(y_val, pred_val)))

    model_path = Path(model_path)
    model_path.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump(
        {
            "model": model,
            "feature_columns": list(FEATURE_COLUMNS),
            "target_column": TARGET_COLUMN,
            "train_months": train_months,
            "val_months": val_months,
        },
        model_path,
    )

    # Feature importance (gain)
    importance_png = Path(importance_png)
    imp = pd.Series(model.feature_importances_, index=FEATURE_COLUMNS).sort_values(ascending=True)
    fig, ax = plt.subplots(figsize=(8, 4.5))
    imp.plot(kind="barh", ax=ax, color="steelblue")
    ax.set_title("LightGBM IV surface — feature importance (gain)")
    ax.set_xlabel("Gain")
    fig.tight_layout()
    fig.savefig(importance_png, dpi=150, bbox_inches="tight")
    plt.close(fig)

    return {
        "rmse_val": rmse,
        "n_train": int(len(X_train)),
        "n_val": int(len(X_val)),
        "model_path": str(model_path.resolve()),
        "feature_importance_png": str(importance_png.resolve()),
    }


# ---------------------------------------------------------------------------
# Optional: build training frame from local Parquet (iVolatility-style) for offline runs
# ---------------------------------------------------------------------------


def options_parquet_to_raw_df(
    parquet_path: str | Path,
    *,
    max_row_groups: int | None = None,
) -> pd.DataFrame:
    """
    Load a bundled EOD Parquet (columns at least: date, expiration, strike, option_type, iv,
    bid) and map to the same schema as :func:`normalize_massive_chain_response` output.
    """
    try:
        import pyarrow.parquet as pq
    except ImportError as e:
        raise ImportError("pip install pyarrow") from e

    path = Path(parquet_path).expanduser()
    if not path.is_file():
        raise FileNotFoundError(path)

    need = ["date", "expiration", "strike", "option_type", "bid", "iv"]
    pf = pq.ParquetFile(path)
    chunks: list[pd.DataFrame] = []
    nrg = pf.num_row_groups if max_row_groups is None else min(max_row_groups, pf.num_row_groups)
    for rg in range(nrg):
        table = pf.read_row_group(rg, columns=[c for c in need if c in pf.schema_arrow.names])
        df = table.to_pandas()
        df["date"] = pd.to_datetime(df["date"], errors="coerce").dt.normalize()
        df["expiration"] = pd.to_datetime(df["expiration"], errors="coerce").dt.normalize()
        df["strike"] = pd.to_numeric(df["strike"], errors="coerce")
        df["bid"] = pd.to_numeric(df["bid"], errors="coerce")
        df["iv"] = pd.to_numeric(df["iv"], errors="coerce")
        df["underlying"] = "SPY"
        df.rename(columns={"iv": "implied_volatility"}, inplace=True)
        # No underlying in file — filled in engineer_features via merge ``close``
        df["underlying_price"] = np.nan
        df["ask"] = np.nan
        df["mid"] = np.nan
        df["volume"] = np.nan
        df["open_interest"] = np.nan
        df["ticker"] = ""
        df["option_type"] = df["option_type"].astype(str).str.upper().str[:1].astype("category")
        chunks.append(df)
    out = pd.concat(chunks, ignore_index=True)
    return _optimize_options_df_dtypes(out)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def resolve_training_date_window(
    start_arg: str | None,
    end_arg: str | None,
) -> tuple[str, str]:
    """
    Default: last **two calendar years** ending **today** (normalized local date).

    * If ``end`` is omitted → use today.
    * If ``start`` is omitted → ``end`` minus two years (``pd.DateOffset(years=2)``).
    """
    today = pd.Timestamp.now(tz=None).normalize()
    end = _norm_day(end_arg) if end_arg else today
    if start_arg:
        start = _norm_day(start_arg)
    else:
        start = end - pd.DateOffset(years=2)
    if start > end:
        raise ValueError(f"--start ({start.date()}) must be on or before --end ({end.date()}).")
    return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")


def main() -> None:
    import argparse

    ap = argparse.ArgumentParser(description="Train LightGBM IV surface (Massive ingest + features)")
    ap.add_argument(
        "--from-parquet",
        type=Path,
        default=None,
        help="Skip Massive fetch; use existing options Parquet (iVolatility-style columns).",
    )
    ap.add_argument(
        "--max-row-groups",
        type=int,
        default=None,
        help="When using --from-parquet, only read the first N row groups (smoke test / RAM cap).",
    )
    ap.add_argument("--massive-cache", type=Path, default=DEFAULT_CACHE_PARQUET)
    ap.add_argument("--symbol", type=str, default="SPY")
    ap.add_argument(
        "--start",
        type=str,
        default=None,
        help="First date YYYY-MM-DD (default: two calendar years before --end).",
    )
    ap.add_argument(
        "--end",
        type=str,
        default=None,
        help="Last date YYYY-MM-DD (default: today).",
    )
    ap.add_argument("--model-out", type=Path, default=DEFAULT_MODEL_PATH)
    ap.add_argument("--importance-out", type=Path, default=DEFAULT_FEATURE_IMPORTANCE_PNG)
    ap.add_argument("--train-months", type=int, default=18)
    ap.add_argument("--val-months", type=int, default=6)
    args = ap.parse_args()

    args.start, args.end = resolve_training_date_window(args.start, args.end)
    print(f"Date window: {args.start} → {args.end} (trailing 2y when dates omitted)")

    yf_pad_start = (pd.Timestamp(args.start) - pd.Timedelta(days=400)).strftime("%Y-%m-%d")
    yf_end = (pd.Timestamp(args.end) + pd.Timedelta(days=5)).strftime("%Y-%m-%d")
    underlying = load_underlying_panel(yf_pad_start, yf_end)

    if args.from_parquet:
        raw = options_parquet_to_raw_df(args.from_parquet, max_row_groups=args.max_row_groups)
    else:
        td = underlying.index[(underlying.index >= _norm_day(args.start)) & (underlying.index <= _norm_day(args.end))]
        cache = fetch_historical_chains(
            args.symbol,
            args.start,
            args.end,
            cache_path=args.massive_cache,
            trading_days=td,
        )
        raw = pd.read_parquet(cache)

    raw = raw[
        (raw["date"] >= _norm_day(args.start))
        & (raw["date"] <= _norm_day(args.end))
        & (pd.to_numeric(raw["bid"], errors="coerce") > 0)
    ].copy()

    feat = engineer_features(raw, underlying)
    print(f"Engineered rows: {len(feat):,}  |  columns: {FEATURE_COLUMNS + [TARGET_COLUMN]}")

    metrics = train_lgbm_iv_surface(
        feat,
        model_path=args.model_out,
        importance_png=args.importance_out,
        train_months=args.train_months,
        val_months=args.val_months,
    )
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()