"""Shared loaders for VRP research scripts (Theta chunks + yfinance SPY/VIX)."""

from __future__ import annotations

from pathlib import Path

import pandas as pd

from RenTech.core.theta_chunks_loader import ThetaChunksLoader, theta_chunks_date_bounds
from RenTech.strategy_stack.vrp_backtester import (
    load_spy_vix_from_yfinance,
    normalize_spy_df,
    trading_days_intersecting_spy,
)

_REPO = Path(__file__).resolve().parents[2]
DEFAULT_THETA_DIR = _REPO / "RenTech" / "data" / "theta_chunks"


def load_theta_run(
    theta_dir: Path | None = None,
    *,
    start: str | None = None,
    end: str | None = None,
) -> tuple[ThetaChunksLoader, pd.DataFrame, list[pd.Timestamp], pd.Timestamp, pd.Timestamp]:
    """
    Theta monthly chunks + aligned SPY/VIX panel and overlapping trading days.

    Optional ``start`` / ``end`` (YYYY-MM-DD) clip the day list (inclusive).
    """
    d = Path(theta_dir or DEFAULT_THETA_DIR).expanduser()
    d0, d1 = theta_chunks_date_bounds(d)
    if start:
        d0 = max(d0, pd.Timestamp(start))
    if end:
        d1 = min(d1, pd.Timestamp(end))
    yf_start = (d0 - pd.Timedelta(days=400)).strftime("%Y-%m-%d")
    yf_end = (d1 + pd.Timedelta(days=14)).strftime("%Y-%m-%d")
    spy = normalize_spy_df(load_spy_vix_from_yfinance(yf_start, yf_end))
    ld = ThetaChunksLoader(d, spy_df=spy)
    days = trading_days_intersecting_spy(ld, spy.index, d0, d1)
    if start:
        t0 = pd.Timestamp(start)
        days = [x for x in days if x >= t0]
    if end:
        t1 = pd.Timestamp(end)
        days = [x for x in days if x <= t1]
    return ld, spy, days, d0, d1


def attach_vvix(spy_df: pd.DataFrame, start: str, end: str) -> pd.DataFrame:
    """Add ``vvix_close`` column aligned to ``spy_df`` (NaN if download fails)."""
    from RenTech.strategy_stack.vrp_backtester import _extract_close_series

    out = spy_df.copy()
    try:
        import yfinance as yf

        raw = yf.download("^VVIX", start=start, end=end, progress=False, auto_adjust=False, threads=False)
        if raw is not None and not raw.empty:
            vv = _extract_close_series(raw, "^VVIX")
            out["vvix_close"] = vv.reindex(out.index).ffill().bfill()
        else:
            out["vvix_close"] = float("nan")
    except Exception:
        out["vvix_close"] = float("nan")
    return out


def year_buckets(days: list[pd.Timestamp]) -> dict[int, list[pd.Timestamp]]:
    by: dict[int, list[pd.Timestamp]] = {}
    for d in days:
        y = int(pd.Timestamp(d).year)
        by.setdefault(y, []).append(d)
    return {y: sorted(v) for y, v in sorted(by.items())}