"""
XGBoost implied-volatility surface model for SPY options.

Trains on historical EOD Parquet (iVolatility-style) merged with **SPY close** as the
underlying and **VIX** as a regime feature. Use :func:`generate_synthetic_chain` to
build an :class:`~RenTech.core.options_data_loader.OptionChain` compatible with
``IVolatilityLoader`` / ``vrp_backtester`` when only spot + VIX are known.

Dependencies
------------
``pandas``, ``numpy``, ``pyarrow``, ``scipy``, ``scikit-learn``, ``xgboost``, ``joblib``,
and ``yfinance`` (for default SPY/VIX panel in training).

Example
-------
>>> from pathlib import Path
>>> from RenTech.core.iv_surface_model import train_surface_model, generate_synthetic_chain
>>> metrics = train_surface_model(Path("SPY-Option-Data/spy_options_eod_combined.parquet"))
>>> chain = generate_synthetic_chain(
...     "RenTech/core/spy_iv_surface.joblib",
...     spy_price=500.0,
...     vix_level=25.0,
...     target_dte=45,
...     strike_range=0.20,
...     option_types=("P",),
... )
"""

from __future__ import annotations

import math
import sys
import warnings
from pathlib import Path
from typing import Any, Literal, Sequence

import numpy as np
import pandas as pd

# Repo root on sys.path when running as ``python3 RenTech/core/iv_surface_model.py``.
_REPO_ROOT = Path(__file__).resolve().parents[2]
if str(_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))

from RenTech.core.options_data_loader import OptionChain, OptionContract

# ---------------------------------------------------------------------------
# Paths & constants
# ---------------------------------------------------------------------------

# Default artifact next to this module (override in train_surface_model).
DEFAULT_MODEL_FILENAME = "spy_iv_surface.joblib"

FEATURE_COLUMNS: list[str] = ["vix_level", "dte", "moneyness", "is_call"]
TARGET_COLUMN = "implied_volatility"  # training label (sourced from Parquet ``iv``)


# ---------------------------------------------------------------------------
# 1. Market data helpers (keep ``core`` independent of ``strategy_stack``)
# ---------------------------------------------------------------------------


def _extract_close_series(raw: pd.DataFrame, label: str) -> pd.Series:
    """Single-ticker yfinance table → numeric close series indexed by date."""
    if raw is None or raw.empty:
        raise ValueError(f"yfinance returned no rows for {label}")
    if isinstance(raw.columns, pd.MultiIndex):
        lv0 = list(raw.columns.get_level_values(0))
        if "Adj Close" in lv0:
            col = raw["Adj Close"]
        elif "Close" in lv0:
            col = raw["Close"]
        else:
            col = raw.iloc[:, 0]
        if isinstance(col, pd.DataFrame):
            col = col.iloc[:, 0]
    else:
        col = raw["Adj Close"] if "Adj Close" in raw.columns else raw["Close"]
    ser = col.squeeze()
    if isinstance(ser, pd.DataFrame):
        ser = ser.iloc[:, 0]
    ser = pd.to_numeric(ser, errors="coerce").astype(float)
    ser.index = pd.to_datetime(ser.index).normalize()
    return ser[~ser.index.duplicated(keep="last")].sort_index()


def load_spy_vix_panel(start: str, end: str) -> pd.DataFrame:
    """
    Download **SPY** and **^VIX** closes for ``[start, end)`` and return a daily frame
    indexed by date with columns ``close`` (SPY) and ``vix_close``.

    VIX is forward-filled then back-filled once so every SPY row has a level.
    """
    try:
        import yfinance as yf
    except ImportError as e:
        raise ImportError("install yfinance: pip install yfinance") from e

    spy_raw = yf.download("SPY", start=start, end=end, progress=False, auto_adjust=False, threads=False)
    vix_raw = yf.download("^VIX", start=start, end=end, progress=False, auto_adjust=False, threads=False)
    spy_close = _extract_close_series(spy_raw, "SPY").rename("close")
    vix_close = _extract_close_series(vix_raw, "^VIX").rename("vix_close")
    df = pd.concat([spy_close, vix_close], axis=1, join="outer").sort_index()
    df["vix_close"] = df["vix_close"].ffill().bfill()
    df = df.dropna(subset=["close"])
    return df


def _parquet_date_bounds(parquet_path: str | Path) -> tuple[pd.Timestamp, pd.Timestamp]:
    """Min/max ``date`` in the Parquet file (single-column scan, cheap)."""
    try:
        import pyarrow.parquet as pq
    except ImportError as e:
        raise ImportError("pyarrow required for Parquet: pip install pyarrow") from e

    path = Path(parquet_path).expanduser()
    pf = pq.ParquetFile(path)
    dmin: pd.Timestamp | None = None
    dmax: pd.Timestamp | None = None
    for i in range(pf.num_row_groups):
        col = pf.read_row_group(i, columns=["date"]).column(0)
        s = pd.to_datetime(col.to_pandas(), errors="coerce").dt.normalize()
        s = s.dropna()
        if s.empty:
            continue
        lo, hi = s.min(), s.max()
        dmin = lo if dmin is None else min(dmin, lo)
        dmax = hi if dmax is None else max(dmax, hi)
    if dmin is None or dmax is None:
        raise ValueError(f"No dates found in {path}")
    return pd.Timestamp(dmin), pd.Timestamp(dmax)


# ---------------------------------------------------------------------------
# 2. Data preparation
# ---------------------------------------------------------------------------


def prepare_training_data(
    parquet_path: str | Path,
    spy_vix: pd.DataFrame | None = None,
    *,
    yfinance_end_pad_days: int = 5,
    max_rows: int | None = None,
) -> tuple[pd.DataFrame, pd.Series]:
    """
    Load options rows from Parquet, merge SPY (underlying) and VIX, engineer features.

    Filters (liquidity / sanity)
    ----------------------------
    - ``iv`` > 0  (implied volatility, stored as **decimal**, e.g. 0.25 = 25% vol)
    - ``bid`` > 0
    - calendar ``dte`` >= 1
    - ``volume`` > 0 **if** a ``volume`` column exists; otherwise that filter is **skipped**
      (the bundled ``spy_options_eod_combined.parquet`` has no volume column).

    Features
    --------
    - ``moneyness`` = strike / underlying (SPY close on ``date``)
    - ``is_call`` = 1 if call else 0
    - ``vix_level`` = VIX close on ``date``
    - ``dte`` = (expiration - date).days

    Target
    ------
    - ``implied_volatility`` ← Parquet column ``iv``

    Parameters
    ----------
    parquet_path
        Path to combined EOD Parquet.
    spy_vix
        DataFrame indexed by normalized date with ``close`` and ``vix_close``. If
        ``None``, dates are inferred from the Parquet file and yfinance is used.
    yfinance_end_pad_days
        When downloading from yfinance, extend the end date by this many days so the
        last chain dates still get a VIX print.
    max_rows
        If set, stop after accumulating this many **post-filter** rows (subsampling for
        quick experiments).

    Returns
    -------
    X, y
        ``X`` has columns :data:`FEATURE_COLUMNS`; ``y`` is the IV series aligned to ``X``.
    """
    try:
        import pyarrow.parquet as pq
    except ImportError as e:
        raise ImportError("pyarrow required: pip install pyarrow") from e

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

    if spy_vix is None:
        d0, d1 = _parquet_date_bounds(path)
        start_s = d0.strftime("%Y-%m-%d")
        end_s = (d1 + pd.Timedelta(days=yfinance_end_pad_days)).strftime("%Y-%m-%d")
        spy_vix = load_spy_vix_panel(start_s, end_s)

    # Join keys: normalized pandas timestamps
    mkt = spy_vix.copy()
    mkt.index = pd.to_datetime(mkt.index).normalize()
    mkt = mkt[~mkt.index.duplicated(keep="last")]
    if "close" not in mkt.columns or "vix_close" not in mkt.columns:
        raise ValueError("spy_vix must contain columns 'close' and 'vix_close'.")

    need_cols = ["date", "expiration", "strike", "option_type", "bid", "iv"]
    pf = pq.ParquetFile(path)
    has_volume = "volume" in pf.schema_arrow.names

    chunks: list[pd.DataFrame] = []
    total = 0

    for rg in range(pf.num_row_groups):
        cols = need_cols + (["volume"] if has_volume else [])
        table = pf.read_row_group(rg, columns=cols)
        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")

        dte = (df["expiration"] - df["date"]).dt.days.astype(np.int64)
        df["dte"] = dte

        df = df.merge(
            mkt[["close", "vix_close"]].rename(columns={"close": "underlying_price"}),
            left_on="date",
            right_index=True,
            how="inner",
        )

        ok = (
            (df["iv"] > 0.0)
            & np.isfinite(df["iv"])
            & (df["bid"] > 0.0)
            & (df["dte"] >= 1)
            & (df["underlying_price"] > 0.0)
        )
        if has_volume:
            df["volume"] = pd.to_numeric(df["volume"], errors="coerce").fillna(0.0)
            ok &= df["volume"] > 0

        df = df.loc[ok].copy()
        if df.empty:
            continue

        df["moneyness"] = df["strike"] / df["underlying_price"]
        df["is_call"] = (df["option_type"].astype(str).str.upper().str.startswith("C")).astype(np.int8)
        df["vix_level"] = pd.to_numeric(df["vix_close"], errors="coerce")
        df = df[np.isfinite(df["vix_level"]) & (df["moneyness"] > 0.0)]

        out = pd.DataFrame(
            {
                "vix_level": df["vix_level"].astype(np.float64),
                "dte": df["dte"].astype(np.float64),
                "moneyness": df["moneyness"].astype(np.float64),
                "is_call": df["is_call"].astype(np.float64),
                TARGET_COLUMN: df["iv"].astype(np.float64),
            }
        )
        chunks.append(out)
        total += len(out)
        if max_rows is not None and total >= max_rows:
            break

    if not chunks:
        raise ValueError("No training rows after filters; check Parquet and spy_vix alignment.")

    full = pd.concat(chunks, ignore_index=True)
    if max_rows is not None and len(full) > max_rows:
        full = full.iloc[:max_rows].copy()

    y = full.pop(TARGET_COLUMN)
    X = full
    return X, y


# ---------------------------------------------------------------------------
# 3. Training
# ---------------------------------------------------------------------------


def train_surface_model(
    parquet_path: str | Path,
    spy_vix: pd.DataFrame | None = None,
    *,
    model_path: str | Path | None = None,
    random_state: int = 42,
    max_rows: int | None = None,
) -> dict[str, Any]:
    """
    Prepare data, 80/20 train/test split, fit ``XGBRegressor``, evaluate, persist model.

    Model hyperparameters (per spec): ``n_estimators=200``, ``max_depth=5``,
    ``learning_rate=0.05``, ``n_jobs=-1``.

    Saves with **joblib** to ``model_path`` (default: :file:`spy_iv_surface.joblib` next
    to this module).

    Returns
    -------
    dict
        Keys include ``mae_train``, ``mae_test``, ``r2_train``, ``r2_test``,
        ``n_train``, ``n_test``, ``model_path``.
    """
    try:
        import joblib
        from sklearn.metrics import mean_absolute_error, r2_score
        from sklearn.model_selection import train_test_split
        from xgboost import XGBRegressor
    except ImportError as e:
        raise ImportError(
            "train_surface_model needs scikit-learn, xgboost, joblib. "
            "pip install scikit-learn xgboost joblib"
        ) from e

    X, y = prepare_training_data(parquet_path, spy_vix, max_rows=max_rows)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=random_state, shuffle=True
    )

    model = XGBRegressor(
        n_estimators=200,
        max_depth=5,
        learning_rate=0.05,
        n_jobs=-1,
        random_state=random_state,
        verbosity=0,
    )
    model.fit(X_train, y_train)

    pred_tr = model.predict(X_train)
    pred_te = model.predict(X_test)

    mae_tr = float(mean_absolute_error(y_train, pred_tr))
    mae_te = float(mean_absolute_error(y_test, pred_te))
    r2_tr = float(r2_score(y_train, pred_tr))
    r2_te = float(r2_score(y_test, pred_te))

    out_path = Path(model_path) if model_path is not None else Path(__file__).resolve().parent / DEFAULT_MODEL_FILENAME
    out_path.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump(model, out_path)

    return {
        "mae_train": mae_tr,
        "mae_test": mae_te,
        "r2_train": r2_tr,
        "r2_test": r2_te,
        "n_train": int(len(X_train)),
        "n_test": int(len(X_test)),
        "model_path": str(out_path.resolve()),
    }


# ---------------------------------------------------------------------------
# 4. Black–Scholes (vectorized)
# ---------------------------------------------------------------------------


def _d1_d2(
    S: np.ndarray,
    K: np.ndarray,
    T: np.ndarray,
    r: float,
    sigma: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Log-normal d1, d2; safe for array inputs."""
    eps = 1e-12
    T = np.maximum(np.asarray(T, dtype=np.float64), eps)
    sigma = np.maximum(np.asarray(sigma, dtype=np.float64), eps)
    S = np.asarray(S, dtype=np.float64)
    K = np.asarray(K, dtype=np.float64)
    sqrt_t = np.sqrt(T)
    d1 = (np.log(S / K) + (r + 0.5 * sigma**2) * T) / (sigma * sqrt_t)
    d2 = d1 - sigma * sqrt_t
    return d1, d2


def synthetic_black_scholes(
    S: float | np.ndarray,
    K: float | np.ndarray,
    T: float | np.ndarray,
    r: float,
    sigma: float | np.ndarray,
    option_type: str | np.ndarray,
) -> np.ndarray:
    """
    Vectorized Black–Scholes **option price per share** (multiply by 100 for contract $).

    Parameters
    ----------
    S
        Spot (e.g. SPY price).
    K
        Strike(s).
    T
        Time to expiry in **years** (e.g. ``dte / 365.0``).
    r
        Continuously compounded risk-free rate (default in callers often ``0.04``).
    sigma
        Annualized volatility as a decimal (matches XGBoost target IV scale).
    option_type
        ``'C'`` / ``'P'`` or an array of ``'C'``/``'P'`` per row.

    Returns
    -------
    np.ndarray
        Theoretical premium **per share** (same units as typical EOD ``mid`` in this repo).
    """
    from scipy.stats import norm

    S = np.asarray(S, dtype=np.float64)
    K = np.asarray(K, dtype=np.float64)
    T = np.asarray(T, dtype=np.float64)
    sigma = np.asarray(sigma, dtype=np.float64)

    d1, d2 = _d1_d2(S, K, T, r, sigma)
    disc = np.exp(-r * T)

    if isinstance(option_type, str):
        is_call = option_type.upper().startswith("C")
        if is_call:
            price = S * norm.cdf(d1) - K * disc * norm.cdf(d2)
        else:
            price = K * disc * norm.cdf(-d2) - S * norm.cdf(-d1)
        return np.asarray(price, dtype=np.float64)

    # Per-row C/P (object or string array)
    ot = np.asarray(option_type, dtype=object)
    is_call = np.vectorize(lambda x: str(x).upper().startswith("C"), otypes=[bool])(ot)
    pc = S * norm.cdf(d1) - K * disc * norm.cdf(d2)
    pp = K * disc * norm.cdf(-d2) - S * norm.cdf(-d1)
    return np.where(is_call, pc, pp).astype(np.float64)


def _bs_greeks_row(
    S: np.ndarray,
    K: np.ndarray,
    T: np.ndarray,
    r: float,
    sigma: np.ndarray,
    is_call: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Delta, gamma, vega (per +1 vol point ≈ 0.01), theta (per year / 365 ≈ per day)."""
    from scipy.stats import norm

    eps = 1e-12
    T = np.maximum(T, eps)
    sigma = np.maximum(sigma, eps)
    d1, d2 = _d1_d2(S, K, T, r, sigma)
    disc = np.exp(-r * T)
    sqrt_t = np.sqrt(T)
    pdf1 = norm.pdf(d1)

    gamma = pdf1 / (S * sigma * sqrt_t)
    vega_share = S * pdf1 * sqrt_t  # per 100% vol unit; /100 for per 1 vol point
    vega = vega_share / 100.0

    delta = np.where(is_call, norm.cdf(d1), norm.cdf(d1) - 1.0)

    # Theta per year (common convention); divide by 365 for rough per-day theta
    theta_call = -S * pdf1 * sigma / (2.0 * sqrt_t) - r * K * disc * norm.cdf(d2)
    theta_put = -S * pdf1 * sigma / (2.0 * sqrt_t) + r * K * disc * norm.cdf(-d2)
    theta = np.where(is_call, theta_call, theta_put) / 365.0

    return delta, gamma, theta, vega


# ---------------------------------------------------------------------------
# 5. Synthetic OptionChain
# ---------------------------------------------------------------------------


def generate_synthetic_chain(
    model_or_path: Any,
    spy_price: float,
    vix_level: float,
    target_dte: int,
    *,
    strike_range: float = 0.20,
    strike_step: float = 1.0,
    r: float = 0.04,
    as_of: pd.Timestamp | None = None,
    option_types: Sequence[Literal["C", "P"]] = ("C", "P"),
    spread_fraction: float = 0.03,
) -> OptionChain:
    """
    Build an :class:`OptionChain` using XGBoost-predicted IVs and Black–Scholes quotes.

    Parameters
    ----------
    model_or_path
        Fitted ``XGBRegressor`` or path to a ``joblib`` file saved by
        :func:`train_surface_model`.
    spy_price
        Spot for moneyness and pricing.
    vix_level
        VIX close (or scenario level) fed into the model.
    target_dte
        Target **calendar** days to expiration; ``expiration = as_of + target_dte`` days.
    strike_range
        Fractional width around spot: strikes run from ``S*(1-range)`` to ``S*(1+range)``.
    strike_step
        Strike grid step in **dollars** (e.g. ``1.0`` for fine SPY strikes).
    r
        Risk-free rate for BS (default 4%).
    as_of
        Chain date; default **today** (normalized).
    option_types
        Which sides to synthesize (default both calls and puts).
    spread_fraction
        Half bid–ask width as a fraction of mid: ``bid = mid*(1-f/2)``, ``ask = mid*(1+f/2)``.

    Returns
    -------
    OptionChain
        Populated with synthetic ``OptionContract`` rows (``mid`` from BS, ``iv`` from model,
        greeks from analytical BS).
    """
    try:
        import joblib
    except ImportError as e:
        raise ImportError("joblib required: pip install joblib") from e

    if isinstance(model_or_path, (str, Path)):
        model = joblib.load(model_or_path)
    else:
        model = model_or_path

    if as_of is None:
        as_of = pd.Timestamp.now(tz=None).normalize()
    else:
        as_of = pd.Timestamp(as_of).normalize()

    if target_dte < 1:
        raise ValueError("target_dte must be >= 1")
    expiration = as_of + pd.Timedelta(days=int(target_dte))
    T = float(target_dte) / 365.0

    lo = spy_price * (1.0 - strike_range)
    hi = spy_price * (1.0 + strike_range)
    strikes = np.arange(math.floor(lo), math.ceil(hi) + strike_step, strike_step, dtype=np.float64)
    strikes = strikes[strikes > 0]

    rows: list[OptionContract] = []
    for opt in option_types:
        want_call = str(opt).upper().startswith("C")
        is_call_int = np.ones(len(strikes), dtype=np.float64) if want_call else np.zeros(len(strikes), dtype=np.float64)
        moneyness = strikes / float(spy_price)
        X_pred = pd.DataFrame(
            {
                "vix_level": np.full(len(strikes), float(vix_level), dtype=np.float64),
                "dte": np.full(len(strikes), float(target_dte), dtype=np.float64),
                "moneyness": moneyness,
                "is_call": is_call_int,
            }
        )
        iv_hat = model.predict(X_pred)
        # Clip to sane decimals for numerics
        iv_hat = np.clip(iv_hat.astype(np.float64), 1e-4, 5.0)

        S = np.full_like(strikes, float(spy_price), dtype=np.float64)
        Tarr = np.full_like(strikes, T, dtype=np.float64)
        ot = np.full(len(strikes), "C" if want_call else "P", dtype=object)

        mid = synthetic_black_scholes(S, strikes, Tarr, r, iv_hat, ot)
        mid = np.maximum(mid, 0.0)

        is_call_arr = np.array([want_call] * len(strikes), dtype=bool)
        delta, gamma, theta, vega = _bs_greeks_row(S, strikes, Tarr, r, iv_hat, is_call_arr)

        half = spread_fraction / 2.0
        bid = np.clip(mid * (1.0 - half), 0.0, None)
        ask = mid * (1.0 + half)

        letter = "C" if want_call else "P"
        for i in range(len(strikes)):
            rows.append(
                OptionContract(
                    date=as_of,
                    expiration=expiration,
                    strike=float(strikes[i]),
                    option_type=letter,
                    bid=float(bid[i]),
                    ask=float(ask[i]),
                    mid=float(mid[i]),
                    iv=float(iv_hat[i]),
                    delta=float(delta[i]),
                    gamma=float(gamma[i]),
                    theta=float(theta[i]),
                    vega=float(vega[i]),
                )
            )

    return OptionChain(as_of=as_of, contracts=rows)


# ---------------------------------------------------------------------------
# CLI / smoke test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import os

    _repo = Path(__file__).resolve().parents[2]
    _pq = _repo / "SPY-Option-Data" / "spy_options_eod_combined.parquet"

    if not _pq.is_file():
        warnings.warn(f"Parquet not found at {_pq}; train step skipped.", stacklevel=1)
    else:
        _max_env = os.environ.get("IV_SURFACE_MAX_ROWS", "").strip()
        _max_rows = int(_max_env) if _max_env.isdigit() else None
        print("Training XGBoost IV surface model (this may take several minutes)...")
        if _max_rows:
            print(f"  (subsample: IV_SURFACE_MAX_ROWS={_max_rows})")
        _metrics = train_surface_model(_pq, max_rows=_max_rows)
        print("Training metrics:")
        for k, v in _metrics.items():
            print(f"  {k}: {v}")

    _model_path = Path(__file__).resolve().parent / DEFAULT_MODEL_FILENAME
    if not _model_path.is_file():
        raise SystemExit(f"No model at {_model_path}; training must succeed first.")

    print()
    print("Synthetic 45-DTE put chain — SPY = $500, VIX = 25 (scenario)")
    _chain = generate_synthetic_chain(
        _model_path,
        spy_price=500.0,
        vix_level=25.0,
        target_dte=45,
        strike_range=0.20,
        strike_step=1.0,
        option_types=("P",),
        as_of=pd.Timestamp("2012-06-15"),  # explicit as_of for reproducible expiration math
    )
    _puts = [c for c in _chain.contracts if c.option_type == "P"]
    _puts.sort(key=lambda c: c.strike)
    _df = pd.DataFrame(
        {
            "strike": [c.strike for c in _puts],
            "bid": [c.bid for c in _puts],
            "mid": [c.mid for c in _puts],
            "ask": [c.ask for c in _puts],
            "iv": [c.iv for c in _puts],
            "delta": [c.delta for c in _puts],
        }
    )
    # Compact console view: every 5th row + head/tail
    _idx = sorted(set(range(0, len(_df), 5)) | {0, len(_df) - 1})
    print(_df.iloc[sorted(_idx)].to_string(index=False, float_format=lambda x: f"{x:.4f}"))
    print(f"\nTotal synthetic puts: {len(_puts)}  |  as_of={_chain.as_of.date()}  exp={_puts[0].expiration.date() if _puts else 'n/a'}")