#!/usr/bin/env python3
"""
Canonical **equity sell-the-rip** sleeve (S&P 500 or Russell 3000).

**Signal:** prior close-to-close return ≥ +3% (``pct_rise``), close < SMA200 (in downtrend).
**Direction:** SHORT — falling prices produce positive P&L.
**Rank:** highest normalized ATR (most volatile downtrend names, most likely to resume lower).
**Hold:** 10 trading days · **Top 10** names · equal weight.

This is the mirror image of ``run_sp500_dip_standard.py``:
  - Dip sleeve: buy after a −3% down day in an uptrend.
  - Rip sleeve: short after a +3% bounce in a downtrend.

Expected behavior:
  - Earns when individual stocks are in persistent downtrends with dead-cat bounces.
  - Should contribute positively in 2022-style bear markets.
  - Slightly negative correlation to buy-the-dip and to SPY.

Example::

    cd /Users/robzingale/trading_bot
    PYTHONUNBUFFERED=1 .venv/bin/python RenTech/strategy_stack/run_sp500_sell_rip_standard.py \\
        --start 2016-01-04 --yahoo-period 10y

    # Russell 3000 universe:
    PYTHONUNBUFFERED=1 .venv/bin/python RenTech/strategy_stack/run_sp500_sell_rip_standard.py \\
        --universe russell3000 --start 2016-01-04 --yahoo-period 10y

    # Bear-market-only gate (only short when SPY < SMA200):
    PYTHONUNBUFFERED=1 .venv/bin/python RenTech/strategy_stack/run_sp500_sell_rip_standard.py \\
        --start 2016-01-04 --spy-bear-only
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

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

from RenTech.strategy_stack.data_loader import DataLoader
from RenTech.strategy_stack.equity_universe_loaders import load_equity_panel_dict
from RenTech.strategy_stack.main import _compute_daily_backtest_features
from RenTech.strategy_stack.multi_strategy_manager import SellTheRipSleeve

LOGS = _REPO / "RenTech" / "data" / "logs"
DEFAULT_OUT_PREFIX = {
    "sp500": LOGS / "sp500_sell_rip_standard",
    "russell3000": LOGS / "russell3000_sell_rip_standard",
}


def _filter_equity_by_history(
    equity_dict: dict[str, pd.DataFrame],
    min_first_date: pd.Timestamp,
) -> dict[str, pd.DataFrame]:
    kept: dict[str, pd.DataFrame] = {}
    for t, df in equity_dict.items():
        if df is None or df.empty:
            continue
        idx = pd.to_datetime(df.index).tz_localize(None)
        if idx.min() <= min_first_date:
            kept[t] = df
    return kept


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument(
        "--universe",
        default="sp500",
        choices=("sp500", "russell3000"),
        help="Stock universe (default sp500)",
    )
    ap.add_argument("--start", default="2016-01-04")
    ap.add_argument("--end", default="")
    ap.add_argument(
        "--yahoo-period",
        default="10y",
        help="Yahoo lookback (default 10y; use 20y for full 2016 data)",
    )
    ap.add_argument("--top-n", type=int, default=10, help="Max new shorts per day (default 10)")
    ap.add_argument("--hold-days", type=int, default=10, help="Hold trading days (default 10)")
    ap.add_argument(
        "--pct-rise-min",
        type=float,
        default=0.03,
        help="Enter short when prior daily return >= +this fraction (default 3%%)",
    )
    ap.add_argument("--capital", type=float, default=100_000.0, help="Notional for daily_pnl_usd")
    ap.add_argument("--out-prefix", type=Path, default=None)
    ap.add_argument("--max-tickers", type=int, default=0, help="Cap universe size (0 = full)")
    ap.add_argument(
        "--min-history-from",
        default="",
        help="Drop names whose first bar is after this date",
    )
    ap.add_argument("--refresh-cache", action="store_true")
    ap.add_argument(
        "--spy-bear-only",
        action="store_true",
        help="Only open new shorts when SPY < SMA200 (off by default)",
    )
    args = ap.parse_args()

    uni = str(args.universe).lower()
    out_default = DEFAULT_OUT_PREFIX[uni]
    out_prefix = args.out_prefix if args.out_prefix is not None else out_default

    equity_dict = load_equity_panel_dict(
        uni,
        args.yahoo_period,
        max_tickers=int(args.max_tickers),
        refresh_cache=bool(args.refresh_cache),
    )
    n_loaded = len(equity_dict)
    if args.min_history_from.strip():
        min_first = pd.Timestamp(args.min_history_from).normalize()
    else:
        # Use the earliest bar date in the loaded data as the floor; drop names
        # that started significantly later than the rest of the universe.
        all_starts = [
            pd.to_datetime(df.index).tz_localize(None).min()
            for df in equity_dict.values()
            if df is not None and not df.empty
        ]
        if all_starts:
            data_start = min(all_starts)
            # Allow names that started within ~3 months of the universe's earliest bar
            min_first = data_start + pd.Timedelta(days=90)
        else:
            min_first = pd.Timestamp(args.start).normalize() - pd.Timedelta(days=400)
    equity_dict = _filter_equity_by_history(equity_dict, min_first)
    min_names = 50 if uni == "sp500" else 200
    if len(equity_dict) < min_names:
        raise SystemExit(
            f"Too few {uni} names after history filter ({len(equity_dict)} of {n_loaded} loaded; "
            f"need first bar <= {min_first.date()}). Try --refresh-cache or --yahoo-period max."
        )

    spy_df = _compute_daily_backtest_features(
        DataLoader().fetch_daily("SPY", period=args.yahoo_period)
    )
    eng = SellTheRipSleeve(
        pct_rise_min=float(args.pct_rise_min),
        hold_trading_days=int(args.hold_days),
        only_when_spy_bear=bool(args.spy_bear_only),
    )
    daily_ret = eng.generate_returns(
        equity_dict,
        top_n=int(args.top_n),
        spy_df=spy_df,
        verbose=True,
    )
    daily_ret = daily_ret.sort_index()
    daily_ret.index = pd.to_datetime(daily_ret.index).tz_localize(None)
    mask = daily_ret.index >= pd.Timestamp(args.start)
    if args.end.strip():
        mask &= daily_ret.index <= pd.Timestamp(args.end)
    r = daily_ret.loc[mask].fillna(0.0).astype(np.float64)

    cap = float(args.capital)
    pnl = r * cap
    eq_unit = (1.0 + r).cumprod()
    eq_usd = cap * eq_unit

    n = len(r)
    years = n / 252.0
    end_eq = float(eq_usd.iloc[-1])
    total_ret = end_eq / cap - 1.0
    cagr = (end_eq / cap) ** (1.0 / years) - 1.0 if years > 0 else float("nan")
    dd = eq_usd / eq_usd.cummax() - 1.0
    max_dd = float(dd.min())
    sd = float(r.std(ddof=1)) if n > 1 else float("nan")
    sharpe = float(r.mean() / sd * np.sqrt(252.0)) if sd > 1e-12 else float("nan")
    vol_ann = sd * np.sqrt(252.0) if np.isfinite(sd) else float("nan")

    spy_r = spy_df["ret"].astype(float)
    spy_r.index = pd.to_datetime(spy_r.index).tz_localize(None)
    rho_spy = float(pd.DataFrame({"rip": r, "spy": spy_r}).dropna().corr().iloc[0, 1])

    prefix = Path(out_prefix).expanduser().resolve()
    prefix.parent.mkdir(parents=True, exist_ok=True)
    daily_path = Path(f"{prefix}_daily.csv")
    meta_path = Path(f"{prefix}_meta.json")
    metrics_path = Path(f"{prefix}_metrics.txt")

    out_df = pd.DataFrame(
        {
            "date": r.index.strftime("%Y-%m-%d"),
            "daily_ret": r.values,
            "daily_pnl_usd": pnl.values,
            "equity_unit": eq_unit.values,
            "equity_usd": eq_usd.values,
        }
    )
    out_df.to_csv(daily_path, index=False)

    cmd = (
        f"cd {_REPO} && PYTHONUNBUFFERED=1 .venv/bin/python "
        f"RenTech/strategy_stack/run_sp500_sell_rip_standard.py "
        f"--universe {uni} --start {args.start} --yahoo-period {args.yahoo_period} "
        f"--top-n {args.top_n} --hold-days {args.hold_days} "
        f"--pct-rise-min {args.pct_rise_min}"
    )
    if args.end.strip():
        cmd += f" --end {args.end}"
    if args.spy_bear_only:
        cmd += " --spy-bear-only"

    token = "sp500_sell_rip" if uni == "sp500" else "russell3000_sell_rip"
    meta = {
        "strategy": token,
        "token": token,
        "universe": uni,
        "signal": f"pct_rise>={args.pct_rise_min:.2%},close<SMA200",
        "direction": "short",
        "hold_days": int(args.hold_days),
        "top_n": int(args.top_n),
        "spy_bear_only": bool(args.spy_bear_only),
        "yahoo_period": args.yahoo_period,
        "start": str(r.index.min().date()) if len(r) else args.start,
        "end": str(r.index.max().date()) if len(r) else args.end,
        "n_sessions": int(n),
        "capital": cap,
        "ending_equity_usd": round(end_eq, 2),
        "total_return_pct": round(total_ret * 100.0, 4),
        "cagr_pct": round(cagr * 100.0, 4),
        "sharpe": round(sharpe, 4),
        "max_dd_pct": round(max_dd * 100.0, 4),
        "vol_ann_pct": round(vol_ann * 100.0, 4),
        "rho_spy": round(rho_spy, 4),
        "command": cmd,
        "daily_csv": str(daily_path),
    }
    with open(meta_path, "w") as f:
        json.dump(meta, f, indent=2)

    metrics_txt = (
        f"=== Sell-the-Rip Backtest ===\n"
        f"Universe  : {uni}  ({n_loaded} loaded → {len(equity_dict)} after history filter)\n"
        f"Window    : {meta['start']} → {meta['end']}  ({n} sessions, {years:.2f}y)\n"
        f"Capital   : ${cap:,.0f}\n"
        f"Signal    : prior-day return >= +{args.pct_rise_min:.0%}, close < SMA200\n"
        f"Rank      : highest ATR/price\n"
        f"Top N     : {args.top_n}  |  Hold : {args.hold_days}d\n"
        f"SPY bear gate: {args.spy_bear_only}\n"
        f"\n"
        f"Ending equity : ${end_eq:>12,.2f}\n"
        f"Total return  : {total_ret*100:>8.2f}%\n"
        f"CAGR          : {cagr*100:>8.2f}%\n"
        f"Sharpe        : {sharpe:>8.3f}\n"
        f"Max drawdown  : {max_dd*100:>8.2f}%\n"
        f"Ann. vol      : {vol_ann*100:>8.2f}%\n"
        f"ρ(SPY)        : {rho_spy:>8.3f}\n"
        f"\nCommand:\n  {cmd}\n"
    )
    with open(metrics_path, "w") as f:
        f.write(metrics_txt)

    print(metrics_txt)

    # Yearly returns
    yr_rows = []
    eq_cur = cap
    for yr, g in out_df.groupby(pd.to_datetime(out_df["date"]).dt.year, sort=True):
        p = float(g["daily_pnl_usd"].sum())
        end_y = eq_cur + p
        ret_y = (end_y / eq_cur - 1.0) * 100.0 if eq_cur > 0 else 0.0
        yr_rows.append({"year": int(yr), "return_pct": round(ret_y, 2), "pnl_usd": round(p, 0)})
        eq_cur = end_y
    yr_df = pd.DataFrame(yr_rows)
    print("=== Return by year (chained equity on $100k) ===")
    print(yr_df.to_string(index=False))
    yr_df.to_csv(Path(f"{prefix}_yearly.csv"), index=False)


if __name__ == "__main__":
    main()