#!/usr/bin/env python3
"""Does an ML vol-nowcaster (proxy for CfC/Mamba2) improve the risk layer's de-lever timing vs EMA?

Target: next-20d realized vol of the book's combination. Forecasters: EMA (baseline, what the book
uses) vs gradient-boosting (walk-forward, leak-free) on vol features. Compare (1) OOS vol-forecast
skill, (2) downstream: vol-targeting with ML-vol vs EMA-vol -> Sharpe / maxDD / realized-vol stability.
If GB ~ EMA downstream, the heavier CfC/Mamba2 won't help (more capacity = more overfit). Vol clusters,
so this is the fair shot.
"""
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from multistrat_book import build_streams, volnorm  # noqa: E402
from multistrat_book_v2 import edge_decay_trust  # noqa: E402
from sklearn.ensemble import HistGradientBoostingRegressor  # noqa: E402

TV = 0.10


def ema(x, span):
    a = 2 / (span + 1); out = np.zeros_like(x); m = 0.0
    for t in range(len(x)):
        m = a * x[t] + (1 - a) * m; out[t] = m
    return out


def main():
    streams, days, year = build_streams()
    names = list(streams)
    M = np.column_stack([volnorm(streams[n]) for n in names])
    TH = np.column_stack([edge_decay_trust(M[:, i]) for i in range(len(names))])
    tw = TH / np.maximum(TH.sum(1, keepdims=True), 1e-9)
    combo = np.nansum(tw * np.nan_to_num(M), axis=1)
    T = len(combo)

    # realized next-20d vol (target) and features (all causal)
    H = 20
    fwd = np.full(T, np.nan)
    for t in range(T - H):
        fwd[t] = combo[t + 1:t + 1 + H].std() * math.sqrt(252)
    sq = combo ** 2
    f_ema20 = np.sqrt(ema(sq, 20) * 252)            # EMA vol (the baseline the book uses)
    feats = np.column_stack([
        np.sqrt(ema(sq, 10) * 252), np.sqrt(ema(sq, 20) * 252), np.sqrt(ema(sq, 60) * 252),
        np.abs(combo), np.array([combo[max(0, t - 5):t].std() for t in range(T)]) * math.sqrt(252),
        np.array([combo[max(0, t - 60):t].std() for t in range(T)]) * math.sqrt(252)])

    # walk-forward GB vol forecast
    gb_fc = np.full(T, np.nan); INIT, STEP = 252, 60
    for s in range(INIT, T - H, STEP):
        e = min(s + STEP, T - H)
        tr = np.arange(63, s)
        ok = np.isfinite(fwd[tr]) & np.isfinite(feats[tr]).all(1)
        gb = HistGradientBoostingRegressor(max_depth=3, max_iter=150, learning_rate=0.05, min_samples_leaf=30).fit(feats[tr][ok], fwd[tr][ok])
        gb_fc[s:e] = gb.predict(feats[s:e])

    valid = np.arange(INIT, T - H)
    v = valid[np.isfinite(fwd[valid]) & np.isfinite(gb_fc[valid])]
    print(f"VOL-NOWCASTER TEST — {T} days, OOS {len(v)}")
    print(f"  OOS vol-forecast corr w/ realized:  EMA {np.corrcoef(f_ema20[v], fwd[v])[0,1]:+.2f}   GB {np.corrcoef(gb_fc[v], fwd[v])[0,1]:+.2f}")
    print(f"  OOS RMSE (lower=better):            EMA {np.sqrt(np.mean((f_ema20[v]-fwd[v])**2)):.3f}   GB {np.sqrt(np.mean((gb_fc[v]-fwd[v])**2)):.3f}")

    # downstream: vol-target with each forecast (leverage = TV/forecast, capped 1x, floor 0.3)
    def book(fc):
        L = np.clip(TV / (fc + 1e-9), 0.3, 1.0)
        r = np.zeros(T); r[1:] = L[:-1] * combo[1:]
        rr = r[v[0]:v[-1]]; rr = rr[np.isfinite(rr)]
        ann = rr.mean() * 252; vol = rr.std() * math.sqrt(252)
        eq = np.cumprod(1 + rr); dd = float((eq / np.maximum.accumulate(eq) - 1).min())
        return ann / vol if vol > 0 else float("nan"), dd, vol
    se, dde, ve = book(f_ema20); sg, ddg, vg = book(gb_fc)
    print(f"  downstream book:  EMA-vol  Sharpe {se:+.2f}  maxDD {100*dde:+.1f}%  realizedVol {100*ve:.1f}%")
    print(f"                    GB-vol   Sharpe {sg:+.2f}  maxDD {100*ddg:+.1f}%  realizedVol {100*vg:.1f}%")
    print("\n  VERDICT: GB-vol downstream Sharpe/maxDD meaningfully > EMA-vol = an ML nowcaster helps the risk")
    print("  layer (then CfC/Mamba2 worth it). If ~equal = EMA already captures vol-clustering; keep it simple.")


if __name__ == "__main__":
    main()