"""
How much return can risk-managed leverage add to the US sleeves, and at what
drawdown? A small account can borrow (Reg-T margin, at most 2x) where a fund
with a mandate cannot; the question is whether scaling exposure by recent
volatility lets that leverage be taken without the tail that static leverage
brings.

    PYTHONPATH=. uv run python -m research.lab.levered_book

Input: daily sleeve returns from research/lab/us_final.py (and, when present,
short_horizon.py). Leverage is applied to the sleeve's return stream:

  exposure_t = clip(target / vol_hat_{t-1}, 0, cap)   (vol_hat from returns through t-1)
  r_t        = exposure_t * sleeve_t - max(exposure_t - 1, 0) * (T-bill + spread)
               - |exposure_t - exposure_{t-1}| * cost

Exposure only moves when it is more than ``band`` away from the current one
(fewer orders). Stock sleeves may go to 2x (Reg-T); the trend sleeve already
holds 3x ETFs, which brokers margin at 75-100%, so it is capped at 1x.

Vol estimators: realized over the last 126 days (Barroso & Santa-Clara 2015)
and EWMA with a 20-day half-life. Targets are round numbers, not fitted.
"""

from __future__ import annotations

import os
import warnings

import numpy as np
import pandas as pd

from research.lab.factor_scan import RESULTS
from research.lab.us_final import monthly_blend

warnings.filterwarnings("ignore")

SPAN = ("2011-01-03", "2026-09-24")
LONG = ("2006-09-27", "2026-09-24")
SPREAD = 0.015          # IBKR benchmark + 1.5% under $100k, +1% to $1M
COST = 0.001            # one-way, per unit of exposure changed
RF_PATH = os.path.join("data", "research", "ext", "fred_dtb3.csv")


def tbill(index: pd.DatetimeIndex) -> pd.Series:
    """Daily 3-month T-bill rate (annual, decimal) from FRED DTB3."""
    rf = pd.read_csv(RF_PATH, index_col=0, parse_dates=True).iloc[:, 0]
    rf = pd.to_numeric(rf, errors="coerce").ffill() / 100
    return rf.reindex(index.union(rf.index)).ffill().reindex(index)


def vol_hat(r: pd.Series, kind: str) -> pd.Series:
    if kind == "rv126":
        v = np.sqrt((r ** 2).rolling(126, min_periods=63).mean() * 252)
    else:
        v = np.sqrt((r ** 2).ewm(halflife=20, min_periods=63).mean() * 252)
    return v.shift(1)


def banded(target: pd.Series, band: float) -> pd.Series:
    out = np.empty(len(target))
    cur = np.nan
    for i, t in enumerate(target.to_numpy()):
        if np.isnan(t):
            out[i] = cur if not np.isnan(cur) else 1.0
            continue
        if np.isnan(cur) or abs(t - cur) > band * max(cur, 0.25):
            cur = t
        out[i] = cur
    return pd.Series(out, index=target.index)


def lever(r: pd.Series, rf: pd.Series, exposure: pd.Series) -> pd.Series:
    borrow = (exposure - 1).clip(lower=0) * (rf.reindex(r.index) + SPREAD) / 252
    trade = exposure.diff().abs().fillna(0.0) * COST
    return exposure * r - borrow - trade


def managed(r: pd.Series, rf: pd.Series, target: float, cap: float, kind: str = "rv126",
            band: float = 0.1) -> tuple[pd.Series, pd.Series]:
    exp = banded((target / vol_hat(r, kind)).clip(upper=cap), band)
    return lever(r, rf, exp), exp


def stats(r: pd.Series, rf: pd.Series) -> dict:
    r = r.dropna()
    e = (1 + r).cumprod()
    yrs = len(r) / 252
    yr = (1 + r).groupby(r.index.year).prod() - 1
    roll3 = e.pct_change(756).dropna()
    return {"cagr": e.iloc[-1] ** (1 / yrs) - 1,
            "sharpe": (r - rf.reindex(r.index) / 252).mean() / r.std() * np.sqrt(252),
            "vol": r.std() * np.sqrt(252), "mdd": (e / e.cummax() - 1).min(),
            "worst_yr": yr.min(), "yrs_below_20": int((yr.iloc[1:-1] < 0.20).sum()),
            "n_yrs": len(yr) - 2, "worst_3y": (1 + roll3.min()) ** (1 / 3) - 1}


def fmt(name: str, s: dict, exp: pd.Series | None = None) -> str:
    x = f" avg exp {exp.mean():.2f}" if exp is not None else ""
    return (f"{name:46s} cagr {s['cagr']*100:5.1f}% xsh {s['sharpe']:.2f} vol {s['vol']*100:4.1f}% "
            f"mdd {s['mdd']*100:6.1f}% worst yr {s['worst_yr']*100:6.1f}% worst 3y/yr {s['worst_3y']*100:5.1f}% "
            f"yrs<20% {s['yrs_below_20']}/{s['n_yrs']}{x}")


def main() -> None:
    R = pd.read_pickle(os.path.join(RESULTS, "us_final_sleeves.pkl"))
    extra = os.path.join(RESULTS, "short_horizon_sleeves.pkl")
    if os.path.exists(extra):
        R = R.join(pd.read_pickle(extra), how="left")
    rf = tbill(R.index)

    print("=== single sleeves, static vs vol-managed (2006- for stock sleeves, 2010- trend)")
    for c in ["momentum", "lowrisk", "trend", "spy"]:
        r = R[c].dropna()
        cap = 1.0 if c == "trend" else 2.0
        print(fmt(f"{c} as is", stats(r, rf)))
        for target in (0.15, 0.25):
            for kind in ("rv126", "ewma20"):
                m, e = managed(r, rf, target, cap, kind)
                print(fmt(f"  {c} vt{int(target*100)} {kind} cap {cap:.0f}x", stats(m, rf), e))

    span = R.loc[SPAN[0]:SPAN[1]].fillna(0.0)
    three = {"momentum": 1 / 3, "lowrisk": 1 / 3, "trend": 1 / 3}
    book = monthly_blend(span, three)
    print("\n=== 3-sleeve book 2011-2026: static leverage vs managed")
    for L in (1.0, 1.5, 2.0, 2.5):
        e = pd.Series(L, index=book.index)
        print(fmt(f"static {L:.1f}x", stats(lever(book, rf, e), rf), e))
    for target in (0.20, 0.25, 0.30, 0.35, 0.40):
        for kind in ("rv126", "ewma20"):
            m, e = managed(book, rf, target, 2.5, kind)
            print(fmt(f"book vt{int(target*100)} {kind} cap 2.5x", stats(m, rf), e))

    print("\n=== sleeve-level management, then equal capital (stock sleeves cap 2x, trend 1x)")
    for target in (0.15, 0.20, 0.25, 0.30):
        parts = {}
        for c in ("momentum", "lowrisk", "trend"):
            cap = 1.0 if c == "trend" else 2.0
            parts[c], _ = managed(R[c].dropna(), rf, target, cap)
        P = pd.DataFrame(parts).loc[SPAN[0]:SPAN[1]].fillna(0.0)
        print(fmt(f"sleeves vt{int(target*100)} rv126, 1/3 each", stats(monthly_blend(P, three), rf)))

    print("\n=== the same leverage through 2008 (momentum + lowrisk 50/50, 2006-; no trend sleeve before 2010)")
    two = monthly_blend(R.loc[LONG[0]:LONG[1]].fillna(0.0), {"momentum": .5, "lowrisk": .5})
    for L in (1.0, 2.0):
        e = pd.Series(L, index=two.index)
        print(fmt(f"static {L:.1f}x", stats(lever(two, rf, e), rf), e))
    for target in (0.25, 0.35):
        m, e = managed(two, rf, target, 2.5, "ewma20")
        print(fmt(f"vt{int(target*100)} ewma20 cap 2.5x", stats(m, rf), e))

    print("\n=== alpha left after hedging market beta (rolling 126d beta, futures hedge), 2011-")
    spy = R["spy"]
    hedged = {}
    for c in ["momentum", "lowrisk", "trend"] + [c for c in R.columns if c.endswith("e21")]:
        r = R[c]
        cov = (r * spy).rolling(126, min_periods=63).mean() - r.rolling(126, min_periods=63).mean() * spy.rolling(126, min_periods=63).mean()
        beta = (cov / spy.rolling(126, min_periods=63).var()).shift(1).clip(0, 2)
        hedged[c] = r - beta * (spy - rf / 252)
        print(fmt(f"hedged {c} (avg beta {beta.loc[SPAN[0]:].mean():.2f})", stats(hedged[c].loc[SPAN[0]:].dropna(), rf)))

    print("\n=== yearly, book vt30 ewma20 cap 2.5x vs 3-sleeve 1x vs SPY")
    m, e = managed(book, rf, 0.30, 2.5, "ewma20")
    yr = pd.DataFrame({"vt30": (1 + m).groupby(m.index.year).prod() - 1,
                       "1x": (1 + book).groupby(book.index.year).prod() - 1,
                       "spy": (1 + span["spy"]).groupby(span.index.year).prod() - 1,
                       "avg_exp": e.groupby(e.index.year).mean()})
    print((yr * [100, 100, 100, 1]).round(1).to_string())


if __name__ == "__main__":
    main()