#!/usr/bin/env python3
"""
Fill **missing** ``implied_vol`` in Theta monthly Parquet chunks using the same Black-Scholes
inverse as :class:`~RenTech.core.theta_chunks_loader.ThetaChunksLoader` (Brentq on mid).

Also sets ``iv_inferred`` (bool) where IV was back-filled.

When ``implied_vol`` is available (vendor or inferred), **Black–Scholes delta** is written for
rows that still lack a vendor ``delta``, using SPY close as spot (same convention as the
loader). Set ``greeks_inferred`` where that BS delta was applied. With ``--overwrite-delta``,
vendor delta is replaced whenever BS inputs are valid (useful for quote-only chains).

Requires underlying close by session date (yfinance) for discounting / spot.

Usage::

    python RenTech/data_pipeline/enrich_theta_parquet_iv.py --theta-dir RenTech/data/theta_chunks
    python RenTech/data_pipeline/enrich_theta_parquet_iv.py --theta-dir RenTech/data/theta_chunks --dry-run
    python RenTech/data_pipeline/enrich_theta_parquet_iv.py --inplace --yes
    python RenTech/data_pipeline/enrich_theta_parquet_iv.py --snapshots both --inplace --yes

``--inplace`` overwrites each matching Parquet file (make a backup first unless ``--yes``).
"""

from __future__ import annotations

import argparse
import math
import sys
from pathlib import Path

_REPO = Path(__file__).resolve().parents[2]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

import numpy as np
import pandas as pd
import yfinance as yf

from RenTech.core.iv_surface_model import _bs_greeks_row
from RenTech.core.theta_chunks_loader import _session_dates_series, _solve_iv_from_mid
_DEFAULT_R = 0.04


def _scale_strike_series(strike: pd.Series, spy_close: float) -> pd.Series:
    s = pd.to_numeric(strike, errors="coerce")
    if bool(s.max(skipna=True) < 150):
        s = s * 10.0
    return s


def _row_iv(
    mid: float,
    S: float,
    K: float,
    T: float,
    r: float,
    is_call: bool,
) -> float:
    return float(_solve_iv_from_mid(mid, S, K, T, r, is_call))


def _apply_bs_delta(
    out: pd.DataFrame,
    close_by_date: dict,
    *,
    r_rate: float,
    overwrite_delta: bool,
) -> tuple[pd.DataFrame, int]:
    """
    Set ``delta`` from BS using final ``implied_vol`` + SPY close; ``greeks_inferred`` marks BS rows.
    """
    n = len(out)
    delta_orig = (
        pd.to_numeric(out["delta"], errors="coerce").to_numpy(dtype=np.float64)
        if "delta" in out.columns
        else np.full(n, np.nan, dtype=np.float64)
    )

    qd = _session_dates_series(out["quote_datetime"])
    sess = pd.to_datetime(qd).dt.normalize()
    dates = pd.Series(sess.dt.date, index=out.index)
    S_vals = pd.to_numeric(dates.map(close_by_date), errors="coerce").to_numpy(dtype=np.float64)

    K_raw = pd.to_numeric(out["strike"], errors="coerce").to_numpy(dtype=np.float64)
    mx = np.nanmax(K_raw) if np.any(np.isfinite(K_raw)) else float("nan")
    K = np.where(np.isfinite(K_raw), K_raw * (10.0 if math.isfinite(mx) and mx < 150 else 1.0), np.nan)

    exp = pd.to_datetime(out["expiration"], errors="coerce").dt.normalize()
    dte = (exp - sess).dt.days.to_numpy(dtype=np.float64)
    T = np.maximum(dte / 365.0, 1.0 / 3650.0)

    opt = out["right"].astype(str).str.upper().str.strip()
    is_call = opt.str.startswith("C").to_numpy()

    iv = pd.to_numeric(out["implied_vol"], errors="coerce").to_numpy(dtype=np.float64)

    mask = (
        np.isfinite(S_vals)
        & (S_vals > 0)
        & np.isfinite(K)
        & (K > 0)
        & np.isfinite(iv)
        & (iv > 0)
        & np.isfinite(T)
        & (T > 0)
        & np.isfinite(dte)
    )

    delta_bs_full = np.full(n, np.nan, dtype=np.float64)
    if bool(mask.any()):
        d_vec, _, _, _ = _bs_greeks_row(
            S_vals[mask],
            K[mask],
            T[mask],
            float(r_rate),
            iv[mask],
            is_call[mask],
        )
        delta_bs_full[mask] = d_vec

    vendor_ok = np.isfinite(delta_orig)
    if overwrite_delta:
        write_bs = mask
        greeks_inferred = mask.copy()
    else:
        write_bs = mask & ~vendor_ok
        greeks_inferred = write_bs

    new_delta = np.where(
        mask,
        np.where(
            overwrite_delta | ~vendor_ok,
            delta_bs_full,
            delta_orig,
        ),
        delta_orig,
    )
    out = out.copy()
    out["delta"] = new_delta.astype("float32")
    out["greeks_inferred"] = greeks_inferred.astype(bool)
    return out, int(greeks_inferred.sum())


def enrich_dataframe(
    df: pd.DataFrame,
    close_by_date: dict,
    *,
    r_rate: float = _DEFAULT_R,
    overwrite_delta: bool = False,
    skip_bs_delta: bool = False,
) -> tuple[pd.DataFrame, int, int]:
    """Return copy with implied_vol filled where missing; counts (filled_iv, filled_delta_bs)."""
    if df.empty:
        return df, 0, 0

    out = df.copy()
    qd = _session_dates_series(out["quote_datetime"])
    out["_session"] = pd.to_datetime(qd).dt.normalize()

    filled = 0
    ivs = []
    inferred_flags = []

    for idx, row in out.iterrows():
        raw_iv = row.get("implied_vol")
        has_iv = pd.notna(raw_iv) and math.isfinite(float(raw_iv)) and float(raw_iv) > 0
        bid = float(row["bid"]) if pd.notna(row.get("bid")) else float("nan")
        ask = float(row["ask"]) if pd.notna(row.get("ask")) else float("nan")
        if not (math.isfinite(bid) and math.isfinite(ask) and bid > 0 and ask > 0):
            ivs.append(float(raw_iv) if has_iv else float("nan"))
            inferred_flags.append(False)
            continue

        mid = 0.5 * (bid + ask)
        sess = row["_session"]
        key = pd.Timestamp(sess).date()
        if key not in close_by_date:
            ivs.append(float(raw_iv) if has_iv else float("nan"))
            inferred_flags.append(False)
            continue
        S = float(close_by_date[key])

        opt = str(row.get("right", row.get("option_type", ""))).strip().upper()
        ot = "C" if opt.startswith("C") else ("P" if opt.startswith("P") else "")
        if ot not in {"C", "P"}:
            ivs.append(float(raw_iv) if has_iv else float("nan"))
            inferred_flags.append(False)
            continue
        is_call = ot == "C"

        exp = pd.Timestamp(row["expiration"]).normalize()
        dte = int((exp - pd.Timestamp(sess).normalize()).days)
        T = max(float(dte) / 365.0, 1.0 / 3650.0)

        K = float(_scale_strike_series(pd.Series([row["strike"]]), S).iloc[0])
        if not (math.isfinite(S) and S > 0 and math.isfinite(K) and K > 0):
            ivs.append(float(raw_iv) if has_iv else float("nan"))
            inferred_flags.append(False)
            continue

        if has_iv:
            ivs.append(float(raw_iv))
            inferred_flags.append(False)
        else:
            ivs.append(_row_iv(mid, S, K, T, r_rate, is_call))
            inferred_flags.append(True)
            filled += 1

    out["implied_vol"] = ivs
    out["iv_inferred"] = inferred_flags
    out = out.drop(columns=["_session"], errors="ignore")

    if skip_bs_delta:
        if "greeks_inferred" not in out.columns:
            out["greeks_inferred"] = False
        return out, filled, 0

    out, n_g = _apply_bs_delta(out, close_by_date, r_rate=r_rate, overwrite_delta=overwrite_delta)
    return out, filled, n_g


def _collect_chunk_paths(theta_dir: Path, snapshots: str, roots: tuple[str, ...]) -> list[Path]:
    roots_l = tuple(r.lower() for r in roots)
    p1545: list[Path] = []
    p1000: list[Path] = []
    for r in roots_l:
        # Avoid recursively re-processing already enriched outputs like *_ivfilled.parquet.
        # We only want the raw monthly chunks:
        #   <root>_1545_YYYY_MM.parquet and <root>_1000_YYYY_MM.parquet
        p1545.extend(sorted(theta_dir.glob(f"{r}_1545_[0-9][0-9][0-9][0-9]_[0-9][0-9].parquet")))
        p1000.extend(sorted(theta_dir.glob(f"{r}_1000_[0-9][0-9][0-9][0-9]_[0-9][0-9].parquet")))
    if snapshots == "1545":
        return p1545
    if snapshots == "1000":
        return p1000
    return sorted(set(p1545) | set(p1000), key=lambda x: x.name)


def _root_from_chunk(path: Path) -> str:
    return path.stem.split("_")[0].upper()


def _ticker_for_root(root: str) -> str:
    # Most ETF roots map 1:1. Keep explicit mappings for index-style roots.
    m = {
        "SPX": "^SPX",
        "VIX": "^VIX",
        "NDX": "^NDX",
        "RUT": "^RUT",
    }
    ru = root.upper()
    return m.get(ru, ru)


def _load_close_by_date_for_root(root: str, start: pd.Timestamp, end: pd.Timestamp) -> dict:
    ticker = _ticker_for_root(root)
    # end is exclusive in yf.download
    start_s = (pd.Timestamp(start) - pd.Timedelta(days=5)).strftime("%Y-%m-%d")
    end_s = (pd.Timestamp(end) + pd.Timedelta(days=6)).strftime("%Y-%m-%d")
    hist = yf.download(
        ticker,
        start=start_s,
        end=end_s,
        auto_adjust=True,
        progress=False,
        interval="1d",
    )
    if hist is None or hist.empty:
        raise RuntimeError(f"No yfinance history for root={root} ticker={ticker} in [{start_s}, {end_s})")
    if "Close" not in hist.columns:
        raise RuntimeError(f"Missing Close column for root={root} ticker={ticker}")
    close_raw = hist["Close"]
    if isinstance(close_raw, pd.DataFrame):
        # yfinance can return a 2D Close frame (e.g. ticker-level columns).
        if close_raw.shape[1] == 0:
            raise RuntimeError(f"Empty Close frame for root={root} ticker={ticker}")
        close_raw = close_raw.iloc[:, 0]
    close = pd.to_numeric(close_raw, errors="coerce").dropna()
    return {pd.Timestamp(ix).normalize().date(): float(v) for ix, v in close.items()}


def main() -> None:
    ap = argparse.ArgumentParser(description="Back-fill missing IV and BS delta in Theta Parquet chunks")
    ap.add_argument("--theta-dir", type=Path, default=_REPO / "RenTech" / "data" / "theta_chunks")
    ap.add_argument(
        "--roots",
        type=str,
        default="SPY",
        help="Comma list of roots to process, e.g. SPY or TLT,GLD,IWM,QQQ,USO",
    )
    ap.add_argument(
        "--snapshots",
        choices=("1545", "1000", "both"),
        default="both",
        help="Which monthly SPY chunk files to process (default: both 15:45 and 10:00 ET).",
    )
    ap.add_argument(
        "--overwrite-delta",
        action="store_true",
        help="Replace vendor delta with BS delta whenever IV+spot are valid.",
    )
    ap.add_argument(
        "--skip-bs-delta",
        action="store_true",
        help="Only back-fill implied_vol; leave delta logic unchanged (legacy).",
    )
    ap.add_argument("--dry-run", action="store_true", help="Count rows only; do not write")
    ap.add_argument("--inplace", action="store_true", help="Overwrite original Parquet files")
    ap.add_argument("--yes", action="store_true", help="With --inplace, skip confirmation")
    args = ap.parse_args()

    theta_dir = args.theta_dir.expanduser()
    if not theta_dir.is_dir():
        print(f"ERROR: not a directory: {theta_dir}", file=sys.stderr)
        sys.exit(1)

    roots = tuple(r.strip().upper() for r in str(args.roots).split(",") if r.strip())
    if not roots:
        print("ERROR: --roots produced no values", file=sys.stderr)
        sys.exit(1)

    paths = _collect_chunk_paths(theta_dir, str(args.snapshots), roots)
    if not paths:
        print(
            f"ERROR: no <root>_1545_*.parquet / <root>_1000_*.parquet under {theta_dir} "
            f"for roots={roots} --snapshots={args.snapshots}",
            file=sys.stderr,
        )
        sys.exit(1)

    # Wide date window across selected chunks, then per-root close map.
    dmin, dmax = None, None
    for p in paths:
        ts = pd.read_parquet(p, columns=["quote_datetime"])
        if ts.empty:
            continue
        q = _session_dates_series(ts["quote_datetime"]).dropna()
        if q.empty:
            continue
        lo, hi = q.min(), q.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:
        print("ERROR: no dates in chunks", file=sys.stderr)
        sys.exit(1)
    close_by_date_by_root: dict[str, dict] = {}
    for r in sorted({_root_from_chunk(p) for p in paths}):
        close_by_date_by_root[r] = _load_close_by_date_for_root(r, pd.Timestamp(dmin), pd.Timestamp(dmax))

    total_rows = 0
    total_missing_before = 0
    total_filled_iv = 0
    total_filled_delta = 0

    if args.inplace and not args.yes:
        print("ERROR: --inplace requires --yes (confirm overwrite)", file=sys.stderr)
        sys.exit(1)

    for p in paths:
        df = pd.read_parquet(p)
        root = _root_from_chunk(p)
        close_by_date = close_by_date_by_root.get(root)
        if not close_by_date:
            print(f"ERROR: no close map loaded for root={root} ({p.name})", file=sys.stderr)
            sys.exit(1)
        total_rows += len(df)
        if "implied_vol" in df.columns:
            iv0 = pd.to_numeric(df["implied_vol"], errors="coerce")
            total_missing_before += int((iv0.isna() | (iv0 <= 0)).sum())
        else:
            total_missing_before += len(df)

        new_df, n_iv, n_d = enrich_dataframe(
            df,
            close_by_date,
            overwrite_delta=bool(args.overwrite_delta),
            skip_bs_delta=bool(args.skip_bs_delta),
        )
        total_filled_iv += n_iv
        total_filled_delta += n_d

        if args.dry_run:
            print(f"{p.name}: root={root} rows={len(df)} filled_iv={n_iv} bs_delta_applied={n_d}")
            continue

        out_path = p if args.inplace else p.with_name(p.stem + "_ivfilled.parquet")
        new_df.to_parquet(out_path, index=False)
        print(f"Wrote {out_path.name} root={root} rows={len(new_df)} filled_iv={n_iv} bs_delta_applied={n_d}")

    print("---")
    print(
        f"files={len(paths)} total_rows={total_rows} missing_iv_before≈{total_missing_before} "
        f"filled_iv={total_filled_iv} bs_delta_applied={total_filled_delta}"
    )


if __name__ == "__main__":
    main()