"""
Return sources outside S&P 500 stocks for the US book: multi-asset trend on
ETFs, and trend on bitcoin. Does either add to the three current sleeves
(momentum, low risk, TrendRider)?

    PYTHONPATH=. uv run python -m research.lab.fetch_long_history   # once (BTC)
    PYTHONPATH=. uv run python -m research.lab.diversifiers

Multi-asset trend (managed futures through ETFs): each asset held while above
its n-day average, sized to 10% vol (EWMA, 30-day half-life), gross capped at
1x, the rest in cash at the T-bill rate. ETFs from data/research/tr_long.pkl
(2008-) and data/us.csv (DBC).

Bitcoin trend: long BTC or cash, decided on closes through t-1, traded at the
close of t, 10 bp per unit traded. Tradable in a US brokerage account through
spot bitcoin ETFs (IBIT, FBTC; 2024-), before that only GBTC (OTC, premium /
discount) or futures ETFs (BITO, 2021-, roll cost). The backtest uses BTC-USD
on US trading days (Yahoo's daily close is 00:00 UTC, ~4-5 hours after the
US close), so it is a proxy for the ETF.
  sma{n}    hold while BTC > its n-day average
  tsmom{n}  hold while the n-day return is positive
  vt60      the same, sized to a 60% vol target (cap 1x)
The whole grid is printed; the median rule is the honest estimate.
"""

from __future__ import annotations

import os
import pickle
import warnings

import numpy as np
import pandas as pd

from research.lab.factor_scan import RESULTS
from research.lab.levered_book import fmt, stats, tbill
from research.lab.us_final import monthly_blend

warnings.filterwarnings("ignore")
SRC = os.path.join("data", "research", "long_history.pkl")
COST = 0.001
SPAN = ("2011-01-03", "2026-09-24")
BTC_BLOCKS = [("2015-06-01", "2017-12-31"), ("2018-01-01", "2020-12-31"), ("2021-01-01", "2023-12-31"),
              ("2024-01-01", "2026-09-25")]
THREE = {"momentum": 1 / 3, "lowrisk": 1 / 3, "trend": 1 / 3}


def multi_asset_trend(close: pd.DataFrame, rf: pd.Series, assets: list[str], n: int,
                      target: float = 0.10) -> tuple[pd.Series, pd.Series]:
    r = close.pct_change(fill_method=None)
    W = {}
    for a in assets:
        px = close[a]
        on = (px > px.rolling(n).mean()).astype(float)
        vol = r[a].ewm(halflife=30, min_periods=60).std() * np.sqrt(252)
        W[a] = (on * (target / vol).clip(upper=3.0)).shift(1)
    W = pd.DataFrame(W).fillna(0.0)
    W = W.div(W.sum(axis=1).clip(lower=1.0), axis=0)
    cash = 1 - W.sum(axis=1)
    ret = (W * r[assets].fillna(0.0)).sum(axis=1) + cash * rf.reindex(W.index) / 252 \
        - W.diff().abs().sum(axis=1) * COST
    return ret, W.sum(axis=1)


def btc_trend(px: pd.Series, rf: pd.Series, rule: str, n: int,
              vol_target: float | None = None) -> tuple[pd.Series, pd.Series]:
    r = px.pct_change(fill_method=None).fillna(0.0)
    on = px > (px.rolling(n).mean() if rule == "sma" else px.shift(n))
    expo = on.astype(float)
    if vol_target:
        vol = r.ewm(halflife=20, min_periods=40).std() * np.sqrt(252)
        expo = expo * (vol_target / vol).clip(upper=1.0).fillna(0.0)
    expo = expo.shift(1).fillna(0.0)
    out = expo * r + (1 - expo) * rf.reindex(r.index) / 252 - expo.diff().abs().fillna(0.0) * COST
    return out, expo


def with_book(book: pd.DataFrame, name: str, r: pd.Series, rf: pd.Series, start: str) -> None:
    span = book.loc[start:SPAN[1], ["momentum", "lowrisk", "trend"]].fillna(0.0)
    span[name] = r.reindex(span.index).fillna(0.0)
    m = (1 + span).resample("ME").prod() - 1
    print(f"monthly correlation of {name} with the sleeves:",
          m.corr()[name].drop(name).round(2).to_dict())
    print(fmt(f"3 sleeves {start[:4]}-", stats(monthly_blend(span, THREE), rf)))
    for w in (0.1, 0.25):
        mix = {k: v * (1 - w) for k, v in THREE.items()} | {name: w}
        print(fmt(f"3 sleeves + {int(w*100)}% {name}", stats(monthly_blend(span, mix), rf)))


def main() -> None:
    book = pd.read_pickle(os.path.join(RESULTS, "us_final_sleeves.pkl"))
    rf = tbill(book.index.union(pd.bdate_range("2014-01-01", "2026-09-30")))

    print("=== multi-asset trend on ETFs, 2010-03-")
    with open(os.path.join("data", "research", "tr_long.pkl"), "rb") as fh:
        etf = pickle.load(fh)["close"].ffill()
    etf = etf.join(pd.read_csv(os.path.join("data", "us.csv"), index_col=0, parse_dates=True)[["DBC"]], how="left")
    best = None
    for assets in (["SPY", "TLT", "GLD"], ["TQQQ", "TLT", "GLD"], ["SPY", "TLT", "GLD", "DBC"]):
        for n in (100, 200):
            r, g = multi_asset_trend(etf, rf, assets, n)
            r = r.loc["2010-03-01":]
            s = stats(r, rf)
            print(fmt(f"{'+'.join(assets)} ma{n}", s, g.loc["2010-03-01":]))
            if best is None or s["sharpe"] > best[1]:
                best = (r, s["sharpe"], f"{'+'.join(assets)} ma{n}")
    print(f"best of the grid: {best[2]} (optimistic: picked in sample)")
    with_book(book, "xtrend", best[0], rf, SPAN[0])

    if not os.path.exists(SRC):
        print("\nno data/research/long_history.pkl: run research.lab.fetch_long_history for the BTC part")
        return
    print("\n=== bitcoin trend, 2015-06-")
    with open(SRC, "rb") as fh:
        src = pickle.load(fh)
    days = src["close"]["SPY"].dropna().loc["2014-12-01":].index
    btc = src["close"]["BTC-USD"].reindex(days, method="ffill").dropna()
    hold = btc.pct_change().fillna(0.0).loc["2015-06-01":]
    print(fmt("hold BTC", stats(hold, rf)))
    rows = {}
    for rule in ("sma", "tsmom"):
        for n in (20, 50, 100, 150, 200):
            for vt in (None, 0.6):
                r, e = btc_trend(btc, rf, rule, n, vt)
                r = r.loc["2015-06-01":]
                key = f"{rule}{n}{' vt60' if vt else ''}"
                rows[key] = r
                bl = " ".join(f"{stats(r.loc[a:b], rf)['sharpe']:5.2f}" for a, b in BTC_BLOCKS)
                print(fmt(key, stats(r, rf), e.loc["2015-06-01":]), f"| blocks {bl}")
    R = pd.DataFrame(rows)
    sh = R.apply(lambda x: stats(x, rf)["sharpe"]).sort_values()
    med = sh.index[len(sh) // 2]
    print(f"median-Sharpe rule: {med} ({sh[med]:.2f}); range {sh.min():.2f}-{sh.max():.2f}")
    with_book(book, "btc", R[med], rf, "2015-06-01")
    with open(os.path.join(RESULTS, "diversifiers.pkl"), "wb") as fh:
        pickle.dump({"btc": R, "btc_median_rule": med, "xtrend": best[0]}, fh)


if __name__ == "__main__":
    main()