#!/usr/bin/env python3
"""
**Time-Series Momentum (TSMOM) / Managed Futures** backtest.

This is the academic strategy first published by Moskowitz, Ooi, and Pedersen (2012)
and the basis of most systematic CTA / managed-futures funds.

Core rule — for each asset each month:
  - Compute the trailing 12-minus-1-month return (skip most-recent month to avoid
    short-term reversal noise, i.e., use the close 252 trading days ago vs 21 days ago)
  - If the return is positive  → go LONG  that asset
  - If the return is negative  → go SHORT that asset
  - Scale each position so it contributes equal annualized volatility to the portfolio
    (vol-normalized sizing, a.k.a. risk-parity across assets)
  - Rebalance monthly

Universe (8 ETFs across 4 asset class buckets):
  Equity:        SPY (US large cap), EFA (intl developed), EEM (EM equities)
  Fixed income:  TLT (20yr US Treasury), IEF (7-10yr)
  Commodities:   GLD (gold), DBC (broad commodity index)
  FX/Dollar:     UUP (US dollar basket)

Why this works in multiple environments:
  - 2022 (rate-rising bear market): short SPY, short TLT, short EEM, long DBC, long UUP
  - 2020 COVID crash+rally: short equity early, long bonds, then flip to long equity
  - 2008: short equity/credit, long bonds/gold
  - Bull markets: long equity, long credit, short commodities when appropriate
  The key: it does NOT predict direction — it just rides whatever is already trending.

Example::

    cd /Users/robzingale/trading_bot
    PYTHONUNBUFFERED=1 .venv/bin/python \\
        RenTech/strategy_stack/run_tsmom_managed_futures.py \\
        --start 2016-01-04 --capital 100000
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import yfinance as yf

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

UNIVERSE = {
    "SPY":  "US equities (S&P 500)",
    "EFA":  "Intl developed equities",
    "EEM":  "Emerging market equities",
    "TLT":  "US 20yr Treasury bonds",
    "IEF":  "US 7-10yr Treasury bonds",
    "GLD":  "Gold",
    "DBC":  "Broad commodity index",
    "UUP":  "US Dollar basket",
}

# ─── signal parameters ───────────────────────────────────────────────────────
# Blended signal lookbacks: professional CTAs mix 3/6/12-month horizons equally.
# Each tuple: (lookback_bars, skip_bars). Skip avoids short-term reversal.
SIGNAL_LOOKBACKS = [
    ( 63,  5),   # 3-month trend, skip 1 week
    (126, 10),   # 6-month trend, skip 2 weeks
    (252, 21),   # 12-month trend, skip 1 month
]
VOL_WINDOW       =  63   # ~3 months realized vol for position sizing
TARGET_ASSET_VOL = 0.15  # each asset targets 15% ann vol (standard CTA sizing)
PORT_VOL_TARGET  = 0.12  # portfolio-level annual vol target
# ─────────────────────────────────────────────────────────────────────────────


def download_prices(tickers: list[str], start: str, end: str) -> pd.DataFrame:
    raw = yf.download(tickers, start=start, end=end, auto_adjust=True, progress=False)
    close = raw["Close"].copy()
    close.index = pd.to_datetime(close.index).tz_localize(None)
    close = close.sort_index()
    # Forward-fill gaps ≤ 5 days (holidays / ETF inception holes)
    close = close.ffill(limit=5)
    return close


def run_tsmom(
    start: str,
    end: str,
    capital: float,
    *,
    out_prefix: Path,
    verbose: bool = True,
) -> dict:
    # ── 1. Download prices ────────────────────────────────────────────────────
    # Fetch enough history for warm-up (max lookback + vol window + buffer)
    fetch_start = (pd.Timestamp(start) - pd.DateOffset(years=2)).strftime("%Y-%m-%d")
    prices = download_prices(list(UNIVERSE.keys()), fetch_start, end)
    tickers = [t for t in UNIVERSE if t in prices.columns]
    prices = prices[tickers]

    # ── 2. Daily returns ──────────────────────────────────────────────────────
    rets = prices.pct_change()

    # ── 3. Monthly rebalance dates ────────────────────────────────────────────
    rebal_dates = prices.resample("BME").last().index  # business month-end

    # ── 4. Build daily weight matrix (forward-filled from monthly rebalance) ──
    max_lookback = max(lb for lb, _ in SIGNAL_LOOKBACKS)
    N = len(tickers)
    weights = pd.DataFrame(0.0, index=prices.index, columns=tickers)

    prev_weights = pd.Series(0.0, index=tickers)

    for rd in rebal_dates:
        if rd not in prices.index:
            prior = prices.index[prices.index <= rd]
            if len(prior) == 0:
                continue
            rd = prior[-1]

        pos = rd
        i_now = prices.index.get_loc(pos)
        # Need enough history for the longest lookback + vol window
        if i_now < max_lookback + VOL_WINDOW:
            continue

        # ── Blend signals across 3/6/12-month lookbacks ──────────────────────
        # Each lookback contributes an equal-weighted vote to the direction signal.
        # The final signal is the average of signed scores across horizons.
        signal_votes = pd.Series(0.0, index=tickers)
        valid_lookbacks = 0
        for lb, skip in SIGNAL_LOOKBACKS:
            i_past = i_now - lb
            i_skip = i_now - skip
            if i_past < 0 or i_skip < 0:
                continue
            price_now  = prices.iloc[i_skip]   # recent price (skip short-term reversal)
            price_past = prices.iloc[i_past]   # price lb bars ago
            ret_signal = (price_now / price_past.replace(0, np.nan) - 1.0).fillna(0.0)
            # Normalize signal: |signal| is proportional to trend strength
            # Use sign for direction, weight equally across horizons
            signal_votes += ret_signal / (ret_signal.abs().mean() + 1e-8)
            valid_lookbacks += 1

        if valid_lookbacks == 0:
            continue
        signal_votes /= valid_lookbacks
        direction = np.sign(signal_votes)  # +1 long, -1 short, 0 flat

        # Vol-normalize: scale each asset to TARGET_ASSET_VOL / realized_vol
        vol_window_data = rets.iloc[max(0, i_now - VOL_WINDOW): i_now]
        realized_vol = vol_window_data.std(ddof=1) * np.sqrt(252)
        realized_vol = realized_vol.replace(0, np.nan).fillna(0.20)  # floor at 20%

        asset_weight = direction * (TARGET_ASSET_VOL / realized_vol)

        # Portfolio-level vol scaling
        # Approximate portfolio vol (assuming equal correlation = 0 for simplicity,
        # then re-scale to PORT_VOL_TARGET)
        gross_vol_estimate = (asset_weight.abs() * realized_vol).sum()
        if gross_vol_estimate > 0:
            port_scale = PORT_VOL_TARGET / gross_vol_estimate
            # Cap total gross exposure at 3× (leverage limit)
            port_scale = min(port_scale, 3.0 / asset_weight.abs().sum() if asset_weight.abs().sum() > 0 else 1.0)
            asset_weight *= port_scale

        prev_weights = asset_weight.fillna(0.0)

        # Apply to all days until next rebalance (will be overwritten by next loop)
        next_rebal_dates = rebal_dates[rebal_dates > rd]
        next_rd = next_rebal_dates[0] if len(next_rebal_dates) else prices.index[-1]
        mask = (prices.index > rd) & (prices.index <= next_rd)
        weights.loc[mask] = prev_weights.values

    # ── 5. Trim to backtest window and compute PnL ────────────────────────────
    weights = weights.shift(1).fillna(0.0)  # execute next day (no lookahead)
    start_ts = pd.Timestamp(start)
    mask = (prices.index >= start_ts)
    if end:
        mask &= (prices.index <= pd.Timestamp(end))

    w = weights.loc[mask]
    r = rets.loc[mask].fillna(0.0)
    portfolio_ret = (w * r).sum(axis=1)

    # ── 6. Metrics ────────────────────────────────────────────────────────────
    cap = float(capital)
    eq = cap * (1.0 + portfolio_ret).cumprod()
    n = len(portfolio_ret)
    years = n / 252.0
    end_eq = float(eq.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 / eq.cummax() - 1.0
    max_dd = float(dd.min())
    sd = float(portfolio_ret.std(ddof=1))
    sharpe = float(portfolio_ret.mean() / sd * np.sqrt(252)) if sd > 1e-12 else float("nan")
    vol_ann = sd * np.sqrt(252)

    # Correlation to SPY
    spy_r = rets["SPY"].loc[mask].fillna(0.0) if "SPY" in rets.columns else pd.Series(0.0, index=r.index)
    rho_spy = float(pd.DataFrame({"p": portfolio_ret, "spy": spy_r}).dropna().corr().iloc[0, 1])

    # ── 7. Yearly returns ──────────────────────────────────────────────────────
    yr_rows = []
    eq_cur = cap
    for yr, g in portfolio_ret.groupby(portfolio_ret.index.year):
        p = float((g * eq_cur).sum()) if False else float(eq_cur * ((1+g).prod()-1))
        end_y = eq_cur * float((1+g).prod())
        ret_y = float((1+g).prod()-1)*100
        yr_rows.append({"year": int(yr), "return_pct": round(ret_y, 2), "pnl_usd": round(end_y - eq_cur, 0)})
        eq_cur = end_y
    yr_df = pd.DataFrame(yr_rows)

    # Asset-level attribution for 2022
    w22 = w[w.index.year == 2022]
    r22 = r[r.index.year == 2022]
    asset_pnl_2022 = {}
    for t in tickers:
        asset_pnl_2022[t] = float(((w22[t] * r22[t]) * cap).sum())

    # ── 8. Save outputs ────────────────────────────────────────────────────────
    out_prefix = Path(out_prefix)
    out_prefix.parent.mkdir(parents=True, exist_ok=True)

    # Margin estimate (portfolio margin rates: 15% long, 30% short).
    # TSMOM takes ETF long/short positions; no options premium at risk.
    gross_long_notional = (w.clip(lower=0) * cap).sum(axis=1)
    gross_short_notional = ((-w).clip(lower=0) * cap).sum(axis=1)
    margin_usd = gross_long_notional * 0.15 + gross_short_notional * 0.30

    daily_df = pd.DataFrame({
        "date": portfolio_ret.index.strftime("%Y-%m-%d"),
        "daily_ret": portfolio_ret.values,
        "daily_pnl_usd": (portfolio_ret * cap).values,
        "equity_usd": eq.values,
        "gross_long_notional_usd": gross_long_notional.values,
        "gross_short_notional_usd": gross_short_notional.values,
        "margin_usd": margin_usd.values,
    })
    daily_df.to_csv(f"{out_prefix}_daily.csv", index=False)
    yr_df.to_csv(f"{out_prefix}_yearly.csv", index=False)

    # Per-asset weight snapshot (last bar of each year)
    weight_snap = {}
    for yr, g in w.groupby(w.index.year):
        last = g.iloc[-1]
        weight_snap[str(yr)] = {t: round(float(last[t]), 3) for t in tickers}

    cmd = (
        f"cd {_REPO} && PYTHONUNBUFFERED=1 .venv/bin/python "
        f"RenTech/strategy_stack/run_tsmom_managed_futures.py "
        f"--start {start} --capital {int(capital)}"
    )
    if end:
        cmd += f" --end {end}"

    meta = {
        "strategy": "tsmom_managed_futures",
        "universe": list(UNIVERSE.keys()),
        "signal": "12-minus-1-month return, direction = sign(signal)",
        "sizing": f"vol-norm per asset ({TARGET_ASSET_VOL:.0%} target), portfolio vol target {PORT_VOL_TARGET:.0%}",
        "rebalance": "monthly (business month-end)",
        "start": str(portfolio_ret.index.min().date()),
        "end": str(portfolio_ret.index.max().date()),
        "n_sessions": n,
        "capital": cap,
        "ending_equity_usd": round(end_eq, 2),
        "total_return_pct": round(total_ret * 100, 4),
        "cagr_pct": round(cagr * 100, 4),
        "sharpe": round(sharpe, 4),
        "max_dd_pct": round(max_dd * 100, 4),
        "vol_ann_pct": round(vol_ann * 100, 4),
        "rho_spy": round(rho_spy, 4),
        "asset_pnl_2022_usd": {k: round(v, 0) for k, v in asset_pnl_2022.items()},
        "weight_snapshots_year_end": weight_snap,
        "command": cmd,
    }
    with open(f"{out_prefix}_meta.json", "w") as f:
        json.dump(meta, f, indent=2)

    metrics_txt = f"""=== Time-Series Momentum (Managed Futures) Backtest ===
Universe   : {', '.join(tickers)}
Window     : {meta['start']} → {meta['end']}  ({n} sessions, {years:.2f}y)
Capital    : ${cap:,.0f}
Signal     : 12-minus-1 month trailing return → long/short
Sizing     : vol-normalized (each asset ~{TARGET_ASSET_VOL:.0%} ann vol)
Rebalance  : monthly

Ending equity  : ${end_eq:>12,.2f}
Total return   : {total_ret*100:>8.2f}%
CAGR           : {cagr*100:>8.2f}%
Sharpe         : {sharpe:>8.3f}
Max drawdown   : {max_dd*100:>8.2f}%
Ann. vol       : {vol_ann*100:>8.2f}%
ρ(SPY)         : {rho_spy:>8.3f}

Command:
  {cmd}
"""
    with open(f"{out_prefix}_metrics.txt", "w") as f:
        f.write(metrics_txt)

    if verbose:
        print(metrics_txt)
        print("=== Return by year ===")
        print(yr_df.to_string(index=False))
        print()
        print("=== 2022 asset-level P&L attribution (${:,.0f} notional) ===".format(cap))
        for t, pnl in sorted(asset_pnl_2022.items(), key=lambda x: -x[1]):
            desc = UNIVERSE.get(t, "")
            print(f"  {t:<5}  {pnl:>+9,.0f}  {desc}")
        print()
        print("=== Year-end position snapshot (weight = gross fraction of capital) ===")
        for yr, snaps in weight_snap.items():
            positions = [(t, v) for t, v in snaps.items() if abs(v) > 0.005]
            pos_str = "  ".join(f"{t}:{v:+.2f}" for t, v in positions)
            print(f"  {yr}: {pos_str}")

    return meta


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--start", default="2016-01-04")
    ap.add_argument("--end", default="")
    ap.add_argument("--capital", type=float, default=100_000.0)
    ap.add_argument("--out-prefix", type=Path, default=LOGS / "tsmom_managed_futures")
    args = ap.parse_args()

    run_tsmom(
        start=args.start,
        end=args.end or "",
        capital=args.capital,
        out_prefix=args.out_prefix,
        verbose=True,
    )


if __name__ == "__main__":
    main()
