#!/usr/bin/env python3
"""
Canonical **equity buy-the-dip** sleeve (S&P 500 or Russell 3000; top experiment variant).

**Signal:** prior close-to-close return ≤ −3% (``pct_drop``), close > SMA200.
**Hold:** 10 trading days · **Top N** names by normalized ATR · equal weight (default **5** S&P 500, **10** Russell 3000).

Replaces the legacy SPDR **sector RSI** dip in the Best Ideas stack; see
``buy_the_dip_experiment_sp500_20160104.md``.

Example::

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

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 BuyTheDipSleeve

LOGS = _REPO / "RenTech" / "data" / "logs"
DEFAULT_OUT_PREFIX = {
    "sp500": LOGS / "sp500_dip_standard",
    "sp100": LOGS / "sp100_dip_standard",
    "russell3000": LOGS / "russell3000_dip_standard",
}
DEFAULT_TOP_N_BY_UNIVERSE: dict[str, int] = {
    "sp500": 5,
    "sp100": 10,
    "russell3000": 10,
}


def _filter_equity_by_history(
    equity_dict: dict[str, pd.DataFrame],
    min_first_date: pd.Timestamp,
) -> dict[str, pd.DataFrame]:
    """Keep tickers with daily history stretching back to ``min_first_date``."""
    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", "sp100", "russell3000"),
        help="Stock universe (sp100 ≈ Nasdaq-100 proxy for CrackingMarkets article)",
    )
    ap.add_argument("--start", default="2016-01-04")
    ap.add_argument("--end", default="")
    ap.add_argument(
        "--yahoo-period",
        default="max",
        help="Yahoo lookback (default max — needed for SMA200 before 2016 starts)",
    )
    ap.add_argument(
        "--top-n",
        type=int,
        default=None,
        help="Max new dip entries per day (default 5 for sp500, 10 for russell3000)",
    )
    ap.add_argument("--hold-days", type=int, default=10, help="Hold trading days (default 10)")
    ap.add_argument(
        "--pct-drop-min",
        type=float,
        default=0.03,
        help="Enter 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 list)")
    ap.add_argument(
        "--min-history-from",
        default="",
        help="Drop names whose first bar is after this date (default: --start minus SMA warmup)",
    )
    ap.add_argument(
        "--refresh-cache",
        action="store_true",
        help="Ignore fresh disk cache and re-download Yahoo bars",
    )
    ap.add_argument(
        "--spy-bull-only",
        action="store_true",
        help="Only open new dips when SPY > SMA200 (off by default; experiment showed worse)",
    )
    ap.add_argument(
        "--relative-spy-min",
        type=float,
        default=0.0,
        metavar="F",
        help="Require same-day return to underperform SPY by at least F (e.g. 0.02)",
    )
    ap.add_argument(
        "--rank-by",
        default="atr_norm",
        choices=("atr_norm", "rsi", "rel_underperf"),
        help="Rank new dip candidates (default atr_norm; rel_underperf for vs-SPY)",
    )
    ap.add_argument(
        "--hedge-spy",
        action="store_true",
        help="Subtract SPY return × gross long exposure (market-neutral excess sleeve)",
    )
    ap.add_argument(
        "--max-holdings",
        type=int,
        default=None,
        metavar="N",
        help="Max names held at once across overlapping positions (default unlimited)",
    )
    ap.add_argument(
        "--execution-style",
        default="legacy",
        choices=("legacy", "cracking_markets"),
        help="legacy=close entry/fixed hold; cracking_markets=limit entry + multi-rule exits",
    )
    ap.add_argument(
        "--limit-atr-mult",
        type=float,
        default=0.9,
        help="CrackingMarkets limit offset below signal close (× ATR5; default 0.9)",
    )
    ap.add_argument(
        "--profit-atr-mult",
        type=float,
        default=0.5,
        help="CrackingMarkets profit target (× ATR5 at signal; default 0.5)",
    )
    ap.add_argument(
        "--no-exit-prior-high",
        action="store_true",
        help="Disable CrackingMarkets exit when close > prior session high",
    )
    args = ap.parse_args()
    uni = str(args.universe).lower()
    top_n = int(args.top_n) if args.top_n is not None else int(DEFAULT_TOP_N_BY_UNIVERSE[uni])
    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:
        # SMA200 + ATR/RSI warmup (~1 calendar year buffer before --start).
        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 in ("sp500", "sp100") 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)
    )
    exec_style = str(args.execution_style).lower()
    cm = exec_style == "cracking_markets"
    eng = BuyTheDipSleeve(
        signal_mode="pct_drop",
        pct_drop_min=float(args.pct_drop_min),
        hold_trading_days=int(args.hold_days),
        only_when_spy_bull=bool(args.spy_bull_only),
        relative_spy_min=float(args.relative_spy_min) if not cm else 0.0,
        rank_by=str(args.rank_by),
        hedge_spy=bool(args.hedge_spy) if not cm else False,
        max_concurrent=int(args.max_holdings) if args.max_holdings is not None else None,
        execution_style=exec_style,
        limit_atr_mult=float(args.limit_atr_mult),
        profit_atr_mult=float(args.profit_atr_mult),
        exit_on_prior_high=not bool(args.no_exit_prior_high),
        cracking_markets_pure=cm,
    )
    daily_ret = eng.generate_returns(
        equity_dict,
        top_n=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({"dip": 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 = 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.to_csv(daily_path, index=False)

    cmd = (
        f"cd {_REPO} && PYTHONUNBUFFERED=1 .venv/bin/python "
        f"RenTech/strategy_stack/run_sp500_dip_standard.py "
        f"--universe {uni} --start {args.start} --yahoo-period {args.yahoo_period} "
        f"--top-n {top_n} --hold-days {args.hold_days} "
        f"--pct-drop-min {args.pct_drop_min}"
    )
    if args.end.strip():
        cmd += f" --end {args.end}"
    if args.spy_bull_only:
        cmd += " --spy-bull-only"
    if args.max_holdings is not None:
        cmd += f" --max-holdings {int(args.max_holdings)}"
    if cm:
        cmd += (
            f" --execution-style cracking_markets --limit-atr-mult {args.limit_atr_mult}"
            f" --profit-atr-mult {args.profit_atr_mult}"
        )

    token = {
        "sp500": "sp500_dip",
        "sp100": "sp100_dip",
        "russell3000": "russell3000_dip",
    }[uni]
    meta = {
        "strategy": token,
        "token": token,
        "portfolio_status": "approved_for_total_portfolio" if uni == "sp500" else "optional_sleeve",
        "replaces": "sector_dip" if uni == "sp500" else None,
        "engine": "BuyTheDipSleeve",
        "universe": uni,
        "n_tickers": len(equity_dict),
        "n_tickers_loaded_before_filter": int(n_loaded),
        "min_first_bar_required": str(min_first.date()),
        "signal_mode": "pct_drop",
        "pct_drop_min": float(args.pct_drop_min),
        "relative_spy_min": float(args.relative_spy_min),
        "rank_by": str(args.rank_by),
        "hedge_spy": bool(args.hedge_spy),
        "hold_trading_days": int(args.hold_days),
        "top_n": int(top_n),
        "max_concurrent": int(args.max_holdings) if args.max_holdings is not None else None,
        "execution_style": exec_style,
        "limit_atr_mult": float(args.limit_atr_mult),
        "profit_atr_mult": float(args.profit_atr_mult),
        "exit_on_prior_high": not bool(args.no_exit_prior_high),
        "cracking_markets_pure": cm,
        "dip_in_uptrend": True,
        "only_when_spy_bull": bool(args.spy_bull_only),
        "capital_usd": cap,
        "yahoo_period": args.yahoo_period,
        "start": str(r.index.min().date()),
        "end": str(r.index.max().date()),
        "n_trading_days": int(n),
        "total_return_pct": round(100.0 * total_ret, 4),
        "cagr_pct": round(100.0 * cagr, 4),
        "sharpe_daily": round(sharpe, 4),
        "vol_annual_pct": round(100.0 * vol_ann, 4),
        "max_drawdown_pct": round(100.0 * max_dd, 4),
        "corr_vs_spy_daily": round(rho_spy, 4),
        "experiment_ref": "RenTech/data/logs/buy_the_dip_experiment_sp500_20160104.md",
        "command": cmd,
        "daily_csv": str(daily_path),
    }
    meta_path.write_text(json.dumps(meta, indent=2) + "\n")

    metrics_path.write_text(
        f"""# {uni} pct-drop dip ({meta['start']} -> {meta['end']})

Variant: drop>={args.pct_drop_min:.0%}, hold {args.hold_days}d, top {top_n} ({uni}, n={len(equity_dict)})

Command:
{cmd}

Headline ($ {cap:,.0f} notional, daily_pnl_usd = daily_ret * capital):
  Total return: {meta['total_return_pct']:.2f}%
  CAGR: {meta['cagr_pct']:.2f}%
  Sharpe: {meta['sharpe_daily']:.3f}
  Max DD: {meta['max_drawdown_pct']:.2f}%
  Corr vs SPY: {meta['corr_vs_spy_daily']:.3f}

Artifacts:
  {daily_path}
  {meta_path}
"""
    )

    print(metrics_path.read_text(), flush=True)
    print(f"Wrote {daily_path}", flush=True)


if __name__ == "__main__":
    main()