"""
Load SPY option chains from ThetaData 15:45 ET monthly Parquet chunks (``build_theta_dataset``).

Schema: ``quote_datetime``, ``expiration``, ``strike``, ``right``, ``bid``, ``ask``,
``implied_vol``, ``delta``.

If ``implied_vol`` / ``delta`` are missing (e.g. ``--skip-greeks`` pulls), this loader
uses **Black–Scholes** with **SPY close** from the provided ``spy_df`` to infer IV from
mid and populate delta/gamma/theta/vega so :class:`~RenTech.core.options_data_loader.OptionChain`
targeting (e.g. VRP backtester) works.
"""

from __future__ import annotations

import math
from datetime import date
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Iterator

import numpy as np
import pandas as pd
from scipy.optimize import brentq

from RenTech.core.iv_surface_model import _bs_greeks_row, synthetic_black_scholes
from RenTech.core.options_data_loader import OptionChain, OptionContract

_DEFAULT_R = 0.04


def theta_chunks_date_bounds(theta_dir: str | Path) -> tuple[pd.Timestamp, pd.Timestamp]:
    """Min/max **session date** (America/New_York calendar day) across ``spy_1545_*.parquet``."""
    d = Path(theta_dir).expanduser()
    if not d.is_dir():
        raise NotADirectoryError(f"Not a directory: {d}")
    paths = sorted(d.glob("spy_1545_*.parquet"))
    if not paths:
        raise FileNotFoundError(f"No spy_1545_*.parquet under {d}")
    lo: pd.Timestamp | None = None
    hi: pd.Timestamp | 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"])
        q = q.dropna()
        if q.empty:
            continue
        a, b = q.min(), q.max()
        lo = a if lo is None else min(lo, a)
        hi = b if hi is None else max(hi, b)
    if lo is None or hi is None:
        raise ValueError(f"No quote_datetime rows in chunks under {d}")
    return pd.Timestamp(lo), pd.Timestamp(hi)


def _session_dates_series(quote_datetime: pd.Series) -> pd.Series:
    """
    NY session **calendar date** as timezone-naive midnight (matches yfinance SPY index).
    """
    qd = pd.to_datetime(quote_datetime, utc=False)
    if qd.dt.tz is None:
        qd = qd.dt.tz_localize(
            "America/New_York",
            ambiguous="NaT",
            nonexistent="shift_forward",
        )
    else:
        qd = qd.dt.tz_convert("America/New_York")
    return pd.to_datetime(qd.dt.date)


def _solve_iv_from_mid(
    mid: float,
    S: float,
    K: float,
    T: float,
    r: float,
    is_call: bool,
) -> float:
    """Brentq IV in (1%, 300%) so BS price matches mid."""
    if not (math.isfinite(mid) and mid > 0 and S > 0 and K > 0 and T > 0):
        return 0.25

    def err(sig: float) -> float:
        p = float(
            synthetic_black_scholes(
                S, K, T, r, float(np.clip(sig, 1e-4, 5.0)), "C" if is_call else "P"
            )
        )
        return p - mid

    try:
        return float(brentq(err, 0.01, 3.0, maxiter=80))
    except ValueError:
        return 0.25


def _row_to_contract(
    r: Mapping[str, Any] | pd.Series,
    as_of: pd.Timestamp,
    spy_close: float,
    r_rate: float,
) -> OptionContract | None:
    bid = float(r["bid"]) if pd.notna(r.get("bid")) else float("nan")
    ask = float(r["ask"]) if pd.notna(r.get("ask")) else float("nan")
    if not (math.isfinite(bid) and math.isfinite(ask) and bid > 0 and ask > 0):
        return None
    mid = 0.5 * (bid + ask)

    opt = str(r.get("right", r.get("option_type", ""))).strip().upper()
    if opt.startswith("C"):
        ot = "C"
    elif opt.startswith("P"):
        ot = "P"
    else:
        ot = opt[:1] if opt else ""
    if ot not in {"C", "P"}:
        return None
    is_call = ot == "C"

    exp = pd.Timestamp(r["expiration"]).normalize()
    as_n = pd.Timestamp(as_of).normalize()
    dte = int((exp - as_n).days)
    T = max(float(dte) / 365.0, 1.0 / 3650.0)
    K = float(r["strike"])
    S = float(spy_close)
    if not (math.isfinite(S) and S > 0 and math.isfinite(K) and K > 0):
        return None

    iv_raw = r.get("implied_vol")
    d_raw = r.get("delta")
    has_iv = pd.notna(iv_raw) and math.isfinite(float(iv_raw)) and float(iv_raw) > 0
    has_d = pd.notna(d_raw) and math.isfinite(float(d_raw))

    if has_iv:
        iv = float(iv_raw)
    else:
        iv = _solve_iv_from_mid(mid, S, K, T, r_rate, is_call)

    S_arr = np.array([S], dtype=np.float64)
    K_arr = np.array([K], dtype=np.float64)
    T_arr = np.array([T], dtype=np.float64)
    sig_arr = np.array([iv], dtype=np.float64)
    is_call_arr = np.array([is_call], dtype=bool)
    delta_b, gamma_b, theta_b, vega_b = _bs_greeks_row(S_arr, K_arr, T_arr, r_rate, sig_arr, is_call_arr)

    delta = float(d_raw) if has_d else float(delta_b[0])
    gamma = float(gamma_b[0])
    theta = float(theta_b[0])
    vega = float(vega_b[0])

    return OptionContract(
        date=as_n,
        expiration=exp,
        strike=K,
        option_type=ot,
        bid=bid,
        ask=ask,
        mid=float(mid),
        iv=float(iv),
        delta=delta,
        gamma=gamma,
        theta=theta,
        vega=vega,
    )


class ThetaChunksLoader:
    """
    Point-in-time chains from ``RenTech/data/theta_chunks/spy_1545_YYYY_MM.parquet``.

    Parameters
    ----------
    theta_dir
        Directory containing monthly Parquet files.
    spy_df
        Panel with ``close`` (and same index convention as backtester) — used to infer
        IV/greeks when Theta rows omit them.
    r
        Risk-free rate for BS (default 4% continuous).
    """

    def __init__(
        self,
        theta_dir: str | Path,
        spy_df: pd.DataFrame,
        *,
        r: float = _DEFAULT_R,
    ) -> None:
        self.theta_dir = Path(theta_dir).expanduser()
        if not self.theta_dir.is_dir():
            raise NotADirectoryError(f"Not a directory: {self.theta_dir}")
        self.spy_df = spy_df
        self.r = float(r)
        self._dates_cache: list[pd.Timestamp] | None = None
        self._spy_id: int | None = None
        self._close_by_date: dict[date, float] | None = None

    def _spy_close(self, as_of: pd.Timestamp) -> float:
        """Resolve SPY close by **calendar date** (robust to tz-naive vs tz-aware index)."""
        sdf = self.spy_df
        if self._spy_id != id(sdf):
            self._spy_id = id(sdf)
            self._close_by_date = {
                pd.Timestamp(ix).date(): float(row["close"])
                for ix, row in sdf.iterrows()
            }
        assert self._close_by_date is not None
        key = pd.Timestamp(as_of).normalize().date()
        if key not in self._close_by_date:
            raise KeyError(f"No SPY row for {key} in spy_df")
        return self._close_by_date[key]

    def iter_chain_dates(self) -> Iterator[pd.Timestamp]:
        if self._dates_cache is None:
            dates: set[pd.Timestamp] = set()
            for p in sorted(self.theta_dir.glob("spy_1545_*.parquet")):
                df = pd.read_parquet(p, columns=["quote_datetime"])
                if df.empty:
                    continue
                q = _session_dates_series(df["quote_datetime"])
                for u in q.dropna().unique():
                    dates.add(pd.Timestamp(u))
            self._dates_cache = sorted(dates)
        yield from self._dates_cache

    def get_chain_for_date(
        self,
        target_date: pd.Timestamp,
        *,
        strike_pct_lo: float | None = None,
        strike_pct_hi: float | None = None,
        min_dte: int | None = None,
        max_dte: int | None = None,
    ) -> OptionChain:
        """
        Parameters
        ----------
        strike_pct_lo, strike_pct_hi
            If set, keep strikes in ``[strike_pct_lo * S, strike_pct_hi * S]`` (else legacy 0.50–1.10).
        min_dte, max_dte
            If set, drop rows whose calendar DTE vs ``ts`` is outside ``[min_dte, max_dte]``
            (inclusive). Evaluated after strike scaling, before per-row BS work.
        """
        ts = pd.Timestamp(pd.Timestamp(target_date).date())
        y, m = ts.year, ts.month
        path = self.theta_dir / f"spy_1545_{y:04d}_{m:02d}.parquet"
        if not path.is_file():
            return OptionChain(as_of=ts, contracts=[])

        df = pd.read_parquet(path)
        if df.empty:
            return OptionChain(as_of=ts, contracts=[])

        qd = _session_dates_series(df["quote_datetime"])
        # Compare calendar dates (avoid ns-resolution / dtype mismatches on ==).
        mask = pd.to_datetime(qd, errors="coerce").dt.date == ts.date()
        df = df.loc[mask]
        if df.empty:
            return OptionChain(as_of=ts, contracts=[])

        # Include **calls and puts** so R1 (weekly strangle) and put regimes share one chain.
        rcol = df["right"] if "right" in df.columns else df.get("option_type", "")
        rup = rcol.astype(str).str.upper().str.strip().str[0]
        df = df.loc[rup.isin(["C", "P"])]
        if df.empty:
            return OptionChain(as_of=ts, contracts=[])

        try:
            spy_close = self._spy_close(ts)
        except KeyError:
            return OptionChain(as_of=ts, contracts=[])

        # ``build_theta_dataset`` stores strike as raw/10_000 (e.g. 45 => $450). Values stay
        # well below ~150; scale to dollars per share for filters and BS.
        strike = pd.to_numeric(df["strike"], errors="coerce")
        if bool(strike.max(skipna=True) < 150):
            strike = strike * 10.0
        df = df.copy()
        df["strike"] = strike

        pct_lo = 0.50 if strike_pct_lo is None else float(strike_pct_lo)
        pct_hi = 1.10 if strike_pct_hi is None else float(strike_pct_hi)
        lo_k = pct_lo * float(spy_close)
        hi_k = pct_hi * float(spy_close)
        df = df.loc[(df["strike"] >= lo_k) & (df["strike"] <= hi_k)]

        if min_dte is not None or max_dte is not None:
            exp = pd.to_datetime(df["expiration"], errors="coerce").dt.normalize()
            dte = (exp - ts).dt.days
            lo_d = int(min_dte) if min_dte is not None else -10**9
            hi_d = int(max_dte) if max_dte is not None else 10**9
            df = df.loc[(dte >= lo_d) & (dte <= hi_d)]

        cols = list(df.columns)
        contracts: list[OptionContract] = []
        for tup in df.itertuples(index=False, name=None):
            row = {cols[i]: tup[i] for i in range(len(cols))}
            c = _row_to_contract(row, ts, spy_close, self.r)
            if c is not None:
                contracts.append(c)

        return OptionChain(as_of=ts, contracts=contracts)