#!/usr/bin/env python3
"""
Attribute max drawdown to closed trades (equity updates on exit only, same as VRPBacktester.metrics).

Usage (mirrors vrp_backtest_theta.py window)::

    python RenTech/strategy_stack/vrp_drawdown_report.py
    python RenTech/strategy_stack/vrp_drawdown_report.py --max-days 500
"""

from __future__ import annotations

import argparse
import math
import sys
from collections import defaultdict
from pathlib import Path

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

import pandas as pd

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

_DEFAULT_THETA_DIR = _REPO_ROOT / "RenTech" / "data" / "theta_chunks"

_REGIME_LABEL = {
    "pmcc": "R1 strangle",
    "diagonal": "R2 diagonal",
    "naked": "R3 put spread",
    "credit_spread": "R4 credit spread",
}


def max_drawdown_episode(
    equity: list[tuple[pd.Timestamp, float]],
) -> tuple[float, int, int, float, float]:
    """
    Running-max drawdown on discrete equity points.
    Returns (max_dd_pct, peak_idx, trough_idx, cap_at_peak, cap_at_trough).
    """
    if len(equity) < 2:
        return 0.0, 0, 0, float(equity[0][1]), float(equity[0][1])

    peak_val = float(equity[0][1])
    peak_idx = 0
    max_dd = 0.0
    best_peak_idx = 0
    best_trough_idx = 0

    for i in range(1, len(equity)):
        _, cap = equity[i]
        cap = float(cap)
        if cap > peak_val:
            peak_val = cap
            peak_idx = i
        dd = (peak_val - cap) / peak_val if peak_val > 0 else 0.0
        if dd > max_dd:
            max_dd = dd
            best_peak_idx = peak_idx
            best_trough_idx = i

    cap_peak = float(equity[best_peak_idx][1])
    cap_trough = float(equity[best_trough_idx][1])
    return max_dd, best_peak_idx, best_trough_idx, cap_peak, cap_trough


def trades_in_drawdown_slice(
    trades: list[ClosedTrade],
    peak_idx: int,
    trough_idx: int,
) -> list[ClosedTrade]:
    """
    equity[k] = capital after k-th trade closes (k>=1); equity[0]=initial.
    Trades from index peak_idx .. trough_idx-1 inclusive move equity[peak_idx] -> equity[trough_idx].
    """
    if trough_idx <= peak_idx:
        return []
    return trades[peak_idx:trough_idx]


def main() -> None:
    ap = argparse.ArgumentParser(description="VRP max drawdown attribution (Theta pipeline)")
    ap.add_argument("--theta-dir", type=Path, default=_DEFAULT_THETA_DIR)
    ap.add_argument("--capital", type=float, default=DEFAULT_STARTING_CAPITAL)
    ap.add_argument("--start", type=str, default="2016-01-04")
    ap.add_argument("--end", type=str, default="2026-04-02")
    ap.add_argument("--max-days", type=int, default=0)
    ap.add_argument("--no-progress", action="store_true")
    ap.add_argument(
        "--no-vol-scaling",
        action="store_true",
        help="Disable VVIX / VIX-momentum risk scaling (default: ON).",
    )
    ap.add_argument(
        "--no-r2-crossover-filters",
        action="store_true",
        help="Disable R2 entry filters (SPY>SMA50, VIX MA20/max(VIX,50); default: ON).",
    )
    args = ap.parse_args()

    theta_dir = args.theta_dir.expanduser()
    d0, d1 = theta_chunks_date_bounds(theta_dir)
    yf_start = (d0 - pd.Timedelta(days=400)).strftime("%Y-%m-%d")
    yf_end = (d1 + pd.Timedelta(days=14)).strftime("%Y-%m-%d")
    spy_wide = normalize_spy_df(load_spy_vix_from_yfinance(yf_start, yf_end))
    ld = ThetaChunksLoader(theta_dir, spy_df=spy_wide)
    days = trading_days_intersecting_spy(ld, spy_wide.index, d0, d1)
    if args.start.strip():
        t0 = pd.Timestamp(args.start.strip())
        days = [d for d in days if d >= t0]
    if args.end.strip():
        t1 = pd.Timestamp(args.end.strip())
        days = [d for d in days if d <= t1]
    if int(args.max_days) > 0:
        days = days[: int(args.max_days)]

    bt = VRPBacktester(
        ld,
        initial_capital=float(args.capital),
        spy_df=spy_wide,
        vol_risk_scaling=not bool(args.no_vol_scaling),
        r2_crossover_filters=not bool(args.no_r2_crossover_filters),
    )
    bt.run_backtest(trading_days=days, show_progress=not bool(args.no_progress))

    eq = bt._equity_curve
    trades = bt.trade_log
    m = bt.metrics()

    max_dd, pidx, tidx, cap_p, cap_t = max_drawdown_episode(eq)
    slice_trades = trades_in_drawdown_slice(trades, pidx, tidx)

    by_reg: dict[Regime, float] = defaultdict(float)
    for tr in slice_trades:
        by_reg[tr.regime] += float(tr.pnl_usd)

    print("=" * 72)
    print("DRAWDOWN (discrete: capital updates on trade exit only)")
    print("=" * 72)
    print(f"  Reported max DD (metrics): {m['max_drawdown']:.2%}")
    print(f"  Recomputed max DD:         {max_dd:.2%}")
    print(f"  Peak equity point index:   {pidx}  date≈{eq[pidx][0].date()}  ${cap_p:,.2f}")
    print(f"  Trough equity point index:   {tidx}  date≈{eq[tidx][0].date()}  ${cap_t:,.2f}")
    print(f"  $ change peak→trough:       ${cap_t - cap_p:,.2f}")
    print(f"  Trades in episode:         {len(slice_trades)}")
    print()
    print("  PnL by regime (within episode):")
    for r in ("pmcc", "diagonal", "naked", "credit_spread"):
        v = by_reg.get(r, 0.0)
        if abs(v) < 1e-9 and r not in by_reg:
            continue
        print(f"    {_REGIME_LABEL[r]:20} ${v:>12,.2f}")
    print()
    worst = sorted(slice_trades, key=lambda x: x.pnl_usd)[:12]
    print("  Worst trades in episode (PnL $):")
    for tr in worst:
        print(
            f"    {tr.exit_date.date()}  {_REGIME_LABEL.get(tr.regime, tr.regime):14} "
            f"{tr.exit_reason:14}  ${tr.pnl_usd:>10,.2f}  days={tr.days_in_trade}"
        )
    print()
    print("=" * 72)
    print("NOTES")
    print("  • DD is measured on realized equity at exits; no intra-trade MTM in this engine.")
    print("  • Episode = from prior equity peak to worst subsequent trough (running max).")
    print("=" * 72)


if __name__ == "__main__":
    main()
