"""VOL09 — EWMA vol-forecast sizing (RiskMetrics vs rolling)

HYPOTHESIS: Use EWMA (RiskMetrics lambda=0.94) to forecast next-bar realized vol
instead of a simple rolling window. Size a long-only position proportionally to
target_vol / ewma_vol_forecast. Compare to simple rolling baseline.

Strategy:
  - Long-only on BTC/ETH (crypto trends upward, short adds drawdown)
  - Trend direction: TSMOM (1-3-6 month blend), flat if negative
  - Sizing: target_vol / ewma_vol_forecast (capped at leverage_cap)
  - EWMA lambda = 0.94 (RiskMetrics standard) vs rolling 30d baseline
  - Config grid: (lambda, target_vol) x 2 options each = 4 combinations
"""

import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np


def ewma_vol(returns: np.ndarray, lam: float, bars_per_year: float) -> np.ndarray:
    """Compute EWMA variance forecast (RiskMetrics style), return annualized vol.
    sigma2[0] = returns[0]^2
    sigma2[i] = lambda * sigma2[i-1] + (1-lambda) * r[i-1]^2  (causal: use r[i-1])

    This is the one-step-ahead forecast: sigma2[i] is the forecast for bar i
    using returns up to r[i-1]. Fully causal.
    """
    n = len(returns)
    sigma2 = np.zeros(n)
    # Initialize with first return squared
    if n > 0:
        sigma2[0] = returns[0] ** 2 if returns[0] != 0 else 1e-6
    for i in range(1, n):
        sigma2[i] = lam * sigma2[i - 1] + (1 - lam) * returns[i - 1] ** 2
    # Annualize: daily vol = sqrt(sigma2), annualized = daily_vol * sqrt(bars_per_year)
    vol = np.sqrt(np.maximum(sigma2, 1e-12)) * np.sqrt(bars_per_year)
    return vol


def tsmom_direction(df, bpd: int) -> np.ndarray:
    """Multi-horizon TSMOM signal (1-3-6 month blend), long-only (0 or 1)."""
    c = df["close"].values
    n = len(c)
    d = np.zeros(n)
    for months in (1, 3, 6):
        h = int(months * 30 * bpd)
        s = np.zeros(n)
        if h < n:
            s[h:] = np.sign(c[h:] / c[:-h] - 1.0)
        d += np.nan_to_num(s)
    # Long-only: clip direction to [0, 1]
    return np.clip(np.sign(d), 0, None)


def make_ewma_target(lam: float, target_vol: float, leverage_cap: float = 2.0):
    """Factory: returns a target_fn(df) for EWMA-vol-sized TSMOM."""
    def target_fn(df):
        c = df["close"].values.astype(float)
        bpd = al.bars_per_day(df)
        bpy = bpd * 365.25

        r = al.simple_returns(c)

        # Causal EWMA vol forecast
        vol_forecast = ewma_vol(r, lam, bpy)

        # TSMOM direction (long-only)
        direction = tsmom_direction(df, bpd)

        # Vol-targeted sizing
        scal = np.where(
            (vol_forecast > 0) & np.isfinite(vol_forecast),
            target_vol / vol_forecast,
            0.0
        )
        tgt = np.clip(direction * scal, 0.0, leverage_cap)
        tgt[~np.isfinite(tgt)] = 0.0
        return tgt
    return target_fn


def make_rolling_target(vol_win_days: int, target_vol: float, leverage_cap: float = 2.0):
    """Baseline: simple rolling vol sizing (same TSMOM direction)."""
    def target_fn(df):
        c = df["close"].values.astype(float)
        bpd = al.bars_per_day(df)
        bpy = bpd * 365.25

        r = al.simple_returns(c)

        # Rolling realized vol
        vol = al.realized_vol(r, max(2, vol_win_days * bpd), bpy)

        # TSMOM direction
        direction = tsmom_direction(df, bpd)

        scal = np.where((vol > 0) & np.isfinite(vol), target_vol / vol, 0.0)
        tgt = np.clip(direction * scal, 0.0, leverage_cap)
        tgt[~np.isfinite(tgt)] = 0.0
        return tgt
    return target_fn


# ---- Internal grid: 4 configs, 2 TFs = 8 backtests (just within 6 per TF pair) ---
# We test EWMA lambda in {0.94, 0.97} x target_vol {0.20} = 2 EWMA configs
# + 1 rolling baseline, across TFs (1d, 12h) = total 6 runs
configs = [
    ("EWMA-lam0.94-tv20", make_ewma_target(lam=0.94, target_vol=0.20)),
    ("EWMA-lam0.97-tv20", make_ewma_target(lam=0.97, target_vol=0.20)),
    ("ROLLING-30d-tv20",  make_rolling_target(vol_win_days=30, target_vol=0.20)),
]

TFS = ("1d", "12h")

# Run all configs on 1d only first to pick best, then run best on both TFs
results = {}
for cfg_name, cfg_fn in configs:
    rep = al.study_weights(f"VOL09/{cfg_name}", cfg_fn, tfs=("1d",))
    best_cell = rep["cells"][0]  # only 1d
    results[cfg_name] = {
        "rep": rep,
        "min_full": best_cell["min_asset_full_sharpe"],
        "min_hold": best_cell["min_asset_holdout_sharpe"],
        "fee_ok": best_cell["fee_survives"],
        "fn": cfg_fn,
    }
    print(f"[1d] {cfg_name}: fullSh={best_cell['min_asset_full_sharpe']:+.3f} "
          f"holdSh={best_cell['min_asset_holdout_sharpe']:+.3f} feeOK={best_cell['fee_survives']}")

# Pick best config by hold-out Sharpe
best_name = max(results, key=lambda k: results[k]["min_hold"])
best_fn = results[best_name]["fn"]
print(f"\nBest config: {best_name}")

# Run best config on both TFs for final report
rep = al.study_weights(f"VOL09 [{best_name}]", best_fn, tfs=TFS)
print(al.fmt(rep))
print("JSON:", al.as_json(rep))