#!/usr/bin/env python3
"""
**Expanded TSMOM / Managed Futures** — 24-asset universe.

Same academic signal as ``run_tsmom_managed_futures.py`` (Moskowitz, Ooi, Pedersen 2012):
  12-minus-1-month blended 3/6/12m signal → long/short → vol-normalized sizing.

Extended universe (6 buckets, 24 ETF proxies):
  Equities (5):  SPY, EFA, EEM, IWM, EWJ
  Bonds (4):     TLT, IEF, SHY, EMB (EM debt)
  FX (5):        UUP (USD), FXE (EUR), FXY (JPY), FXA (AUD), FXB (GBP)
  Commodities (4): GLD, SLV, DBC, USO
  Real assets (3): VNQ (US REIT), PDBC (commodity), CPER (copper)
  Rates/Inflation (3): TIP (inflation), FLOT (float-rate), HYG (HY bonds)

Why this helps weak years:
  2015: short EEM/commodity, long USD — all trended strongly
  2019: long bonds (TLT +14%), short USD (faded) — misses if 8-asset book is short bonds
  2022: short TLT/EEM/SPY, long UUP/CPER/USO — classic managed-futures year

Example::

    cd /Users/robzingale/trading_bot && PYTHONUNBUFFERED=1 \\
      .venv/bin/python RenTech/strategy_stack/run_tsmom_expanded.py \\
      --start 2011-01-03 --end 2025-12-31 --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: dict[str, str] = {
    # Equities
    "SPY":  "US large-cap equities",
    "EFA":  "Intl developed equities",
    "EEM":  "Emerging market equities",
    "IWM":  "US small-cap equities",
    "EWJ":  "Japan equities",
    # Bonds
    "TLT":  "US 20yr Treasury",
    "IEF":  "US 7-10yr Treasury",
    "SHY":  "US 1-3yr Treasury",
    "EMB":  "EM sovereign USD bonds",
    # FX (via ETFs — UUP is USD long; others are quoted vs USD so move opposite USD)
    "UUP":  "US Dollar basket (DXY proxy)",
    "FXE":  "Euro (EUR/USD proxy)",
    "FXY":  "Japanese Yen (JPY/USD proxy)",
    "FXA":  "Australian Dollar (AUD/USD proxy)",
    "FXB":  "British Pound (GBP/USD proxy)",
    # Commodities
    "GLD":  "Gold",
    "SLV":  "Silver",
    "DBC":  "Broad commodity index",
    "USO":  "Crude oil",
    # Real assets / inflation
    "VNQ":  "US REITs",
    "TIP":  "US TIPS (inflation-linked Treasuries)",
    "HYG":  "US High Yield bonds",
    "PDBC": "Diversified commodity (tax-efficient DBC)",
    "CPER": "Copper ETF",
    "IAU":  "Gold (iShares, longer history backup)",
}

SIGNAL_LOOKBACKS = [
    ( 63,  5),
    (126, 10),
    (252, 21),
]
VOL_WINDOW       = 63
TARGET_ASSET_VOL = 0.15
PORT_VOL_TARGET  = 0.12
MAX_GROSS_LEVER  = 3.0


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)
    if isinstance(raw.columns, pd.MultiIndex):
        close = raw["Close"].copy()
    else:
        close = raw[["Close"]].copy() if "Close" in raw.columns else raw.copy()
    close.index = pd.to_datetime(close.index).tz_localize(None)
    close = close.sort_index().ffill(limit=5)
    # Drop tickers with < 2 years of history
    close = close.loc[:, close.notna().sum() >= 500]
    return close


def run_tsmom_expanded(
    start: str,
    end: str,
    capital: float,
    *,
    out_prefix: Path,
    verbose: bool = True,
) -> dict:
    fetch_start = (pd.Timestamp(start) - pd.DateOffset(years=2)).strftime("%Y-%m-%d")
    fetch_end = end or pd.Timestamp.today().strftime("%Y-%m-%d")
    prices = download_prices(list(UNIVERSE.keys()), fetch_start, fetch_end)
    # IAU is backup for GLD — drop if GLD present
    if "GLD" in prices.columns and "IAU" in prices.columns:
        prices = prices.drop(columns=["IAU"])
    # PDBC/CPER may have short history — keep if available, drop otherwise
    tickers = [t for t in prices.columns if t in UNIVERSE]
    prices = prices[tickers]

    rets = prices.pct_change()
    rebal_dates = prices.resample("BME").last().index
    max_lookback = max(lb for lb, _ in SIGNAL_LOOKBACKS)
    weights = pd.DataFrame(0.0, index=prices.index, columns=tickers)

    for rd in rebal_dates:
        if rd not in prices.index:
            prior = prices.index[prices.index <= rd]
            if not len(prior):
                continue
            rd = prior[-1]
        i_now = prices.index.get_loc(rd)
        if i_now < max_lookback + VOL_WINDOW:
            continue

        signal_votes = pd.Series(0.0, index=tickers)
        valid_lb = 0
        for lb, skip in SIGNAL_LOOKBACKS:
            i_past, i_skip = i_now - lb, i_now - skip
            if i_past < 0 or i_skip < 0:
                continue
            p_now  = prices.iloc[i_skip]
            p_past = prices.iloc[i_past]
            r_sig  = (p_now / p_past.replace(0, np.nan) - 1.0).fillna(0.0)
            denom  = r_sig.abs().mean() + 1e-8
            signal_votes += r_sig / denom
            valid_lb += 1
        if valid_lb == 0:
            continue
        signal_votes /= valid_lb
        direction = np.sign(signal_votes)

        vol_data = rets.iloc[max(0, i_now - VOL_WINDOW): i_now]
        realized_vol = vol_data.std(ddof=1) * np.sqrt(252)
        realized_vol = realized_vol.replace(0, np.nan).fillna(0.20)
        asset_weight = direction * (TARGET_ASSET_VOL / realized_vol)

        gross_vol_est = (asset_weight.abs() * realized_vol).sum()
        if gross_vol_est > 0:
            port_scale = PORT_VOL_TARGET / gross_vol_est
            gross = asset_weight.abs().sum()
            if gross > 0:
                port_scale = min(port_scale, MAX_GROSS_LEVER / gross)
            asset_weight *= port_scale

        asset_weight = asset_weight.fillna(0.0)
        next_rdates = rebal_dates[rebal_dates > rd]
        next_rd = next_rdates[0] if len(next_rdates) else prices.index[-1]
        mask = (prices.index > rd) & (prices.index <= next_rd)
        weights.loc[mask] = asset_weight.values

    weights = weights.shift(1).fillna(0.0)
    t0, t1 = pd.Timestamp(start), pd.Timestamp(fetch_end)
    mask = (prices.index >= t0) & (prices.index <= t1)
    w = weights.loc[mask]
    r = rets.loc[mask].fillna(0.0)
    portfolio_ret = (w * r).sum(axis=1)

    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)

    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])

    yr_rows = []
    eq_cur = cap
    for yr, g in portfolio_ret.groupby(portfolio_ret.index.year):
        ret_y = float((1 + g).prod() - 1) * 100
        end_y = eq_cur * (1 + ret_y / 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 weak years
    attribution: dict[str, dict] = {}
    for check_yr in [2015, 2016, 2018, 2019, 2022]:
        wy = w[w.index.year == check_yr]
        ry = r[r.index.year == check_yr]
        if wy.empty:
            continue
        attribution[str(check_yr)] = {
            t: round(float(((wy[t] * ry[t]) * cap).sum()), 0)
            for t in tickers if t in wy.columns
        }

    gross_long = (w.clip(lower=0) * cap).sum(axis=1)
    gross_short = ((-w).clip(lower=0) * cap).sum(axis=1)
    margin_usd = gross_long * 0.15 + gross_short * 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.values,
        "gross_short_notional_usd":gross_short.values,
        "margin_usd":              margin_usd.values,
    })

    out_prefix = Path(out_prefix)
    out_prefix.parent.mkdir(parents=True, exist_ok=True)
    daily_df.to_csv(f"{out_prefix}_daily.csv", index=False)
    yr_df.to_csv(f"{out_prefix}_yearly.csv", index=False)

    meta = {
        "strategy": "tsmom_expanded",
        "universe": tickers,
        "n_assets": len(tickers),
        "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),
        "weak_year_attribution": attribution,
        "yearly_csv": f"{out_prefix}_yearly.csv",
        "daily_csv": f"{out_prefix}_daily.csv",
    }
    with open(f"{out_prefix}_meta.json", "w") as fh:
        json.dump(meta, fh, indent=2)

    if verbose:
        print(f"=== TSMOM Expanded ({len(tickers)} assets) ===")
        print(f"Window: {meta['start']} → {meta['end']}  ({n} sessions)")
        print(f"Return {total_ret*100:.1f}%  CAGR {cagr*100:.1f}%  Sharpe {sharpe:.2f}  MaxDD {max_dd*100:.1f}%  ρ(SPY) {rho_spy:.2f}")
        print("\nYearly returns:")
        print(yr_df.to_string(index=False))
        print("\nAttribution in weak years:")
        for yr_str, attr in attribution.items():
            top = sorted(attr.items(), key=lambda x: -abs(x[1]))[:5]
            top_str = "  ".join(f"{t}:{v:+,.0f}" for t, v in top)
            yr_total = sum(attr.values())
            print(f"  {yr_str}: total ${yr_total:+,.0f}  |  {top_str}")

    return meta


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

    run_tsmom_expanded(
        start=args.start,
        end=args.end,
        capital=float(args.capital),
        out_prefix=args.out_prefix,
        verbose=True,
    )


if __name__ == "__main__":
    main()
