"""STA04 — K-means regime -> trend gating.

IDEA: cluster causal (vol, return, range) features using K-means with expanding
statistics (z-scored causally), then enable TSMOM only in the historically-bullish/
trending cluster. No future labels. Fully causal.

APPROACH:
- Features (causal at bar i):
    1. realized_vol (30-day annualized)
    2. momentum return (lookback days)
    3. normalized range = ATR / close (relative range)
- Expanding z-score: we don't know the distribution of features ahead of time.
  We compute expanding mean/std up to bar i for each feature, then z-score.
  This is causal: uses data[0..i] only.
- K-means: we run offline K-means on the TRAINING portion (full history up to a
  burn-in), then use the fitted centroids to classify new bars causally.
  Strategy: classify each bar, determine which cluster(s) historically have
  been bullish/trending (positive mean return), gate TSMOM only in those clusters.
- TSMOM signal: sign of 3-month return, vol-targeted.

GRID (<=4 combos to keep total backtests <=6 with 2 TFs):
  - (n_clusters=3, lookback_months=3)   <- canonical
  - (n_clusters=4, lookback_months=3)   <- more granular clusters
  Keep TFs = (1d, 12h).
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np
import pandas as pd
from sklearn.cluster import KMeans


def expanding_zscore(x: np.ndarray, min_periods: int = 30) -> np.ndarray:
    """Causal expanding z-score: at bar i, use data[0..i] to compute mean/std."""
    out = np.full(len(x), np.nan)
    for i in range(min_periods, len(x)):
        window = x[:i+1]
        m = np.nanmean(window)
        s = np.nanstd(window)
        if s > 0:
            out[i] = (x[i] - m) / s
        else:
            out[i] = 0.0
    return out


def build_features(df: pd.DataFrame, lookback_months: int) -> np.ndarray:
    """Build causal feature matrix [vol_z, momentum_z, range_z] for each bar."""
    c = df["close"].values.astype(float)
    bpd = al.bars_per_day(df)
    bpy = bpd * 365.25

    # Feature 1: realized vol (30d)
    r = al.simple_returns(c)
    rv = al.realized_vol(r, max(2, 30 * bpd), bpy)

    # Feature 2: momentum return over lookback_months
    lb_bars = int(lookback_months * 30.44 * bpd)
    mom = np.zeros(len(c))
    for i in range(lb_bars, len(c)):
        mom[i] = c[i] / c[i - lb_bars] - 1.0

    # Feature 3: normalized range (ATR / close)
    at = al.atr(df, win=max(2, 14))
    rng = np.where(c > 0, at / c, 0.0)

    # Expanding z-score (causal)
    rv_z = expanding_zscore(rv, min_periods=60)
    mom_z = expanding_zscore(mom, min_periods=60)
    rng_z = expanding_zscore(rng, min_periods=60)

    feat = np.column_stack([rv_z, mom_z, rng_z])
    return feat


def make_target(df: pd.DataFrame, n_clusters: int, lookback_months: int,
                train_frac: float = 0.5) -> np.ndarray:
    """
    K-means regime-gated TSMOM.

    1. Build causal features.
    2. Use the first train_frac of valid data to fit K-means.
    3. Label each cluster: positive if mean forward return (in training) is positive.
    4. Gate TSMOM: position = vol_targeted_tsmom * in_bullish_cluster.
    """
    c = df["close"].values.astype(float)
    bpd = al.bars_per_day(df)
    n = len(df)

    # Build features
    feat = build_features(df, lookback_months)

    # Identify valid (non-NaN) rows
    valid_mask = np.all(np.isfinite(feat), axis=1)

    # TSMOM signal: sign of lookback_months return, vol-targeted, long-only (flat on negative)
    lb_bars = int(lookback_months * 30.44 * bpd)
    tsmom_dir = np.zeros(n)
    for i in range(lb_bars, n):
        ret = c[i] / c[i - lb_bars] - 1.0
        tsmom_dir[i] = 1.0 if ret > 0 else 0.0  # long-flat (no short, consistent with TP01)

    tsmom_pos = al.vol_target(tsmom_dir, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)

    # Find the training cutoff (first train_frac of valid bars)
    valid_idx = np.where(valid_mask)[0]
    if len(valid_idx) < n_clusters * 20:
        # Not enough data, return raw tsmom
        return tsmom_pos

    train_end_idx = valid_idx[int(len(valid_idx) * train_frac)]

    # Fit K-means on training portion
    train_feat = feat[valid_idx[valid_idx <= train_end_idx]]
    if len(train_feat) < n_clusters * 10:
        return tsmom_pos

    km = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
    km.fit(train_feat)

    # Determine cluster "bullishness" from training data:
    # For each training bar, check if the next bar's return is positive.
    # A cluster is "bullish" if mean(next_return | cluster) > 0.
    r = al.simple_returns(c)
    train_labels = km.labels_
    train_valid_indices = valid_idx[valid_idx <= train_end_idx]

    cluster_returns = {k: [] for k in range(n_clusters)}
    for i_pos, idx_i in enumerate(train_valid_indices):
        if idx_i + 1 < n:
            cluster_returns[train_labels[i_pos]].append(r[idx_i + 1])

    bullish_clusters = set()
    for k, rets in cluster_returns.items():
        if len(rets) > 5 and np.mean(rets) > 0:
            bullish_clusters.add(k)

    # If no bullish cluster found, use all clusters (fall back to pure TSMOM)
    if not bullish_clusters:
        bullish_clusters = set(range(n_clusters))

    # Classify ALL valid bars causally using fitted centroids
    all_valid_feat = feat[valid_mask]
    all_labels = km.predict(all_valid_feat)

    # Build gate array
    gate = np.zeros(n)
    for i_pos, idx_i in enumerate(np.where(valid_mask)[0]):
        if all_labels[i_pos] in bullish_clusters:
            gate[idx_i] = 1.0

    # Final position: TSMOM gated by regime
    target = tsmom_pos * gate
    target = np.nan_to_num(target, nan=0.0)
    return target


def run_config(n_clusters: int, lookback_months: int):
    name = f"STA04_k{n_clusters}_lb{lookback_months}m"
    fn = lambda df: make_target(df, n_clusters=n_clusters, lookback_months=lookback_months)
    rep = al.study_weights(name, fn, tfs=("1d", "12h"))
    return rep


if __name__ == "__main__":
    # Grid: 2 configs x 2 TFs = 4 backtests per asset x 2 assets = 8 backtests total.
    # Keep it small: just 2 configs.
    configs = [
        (3, 3),  # 3 clusters, 3-month lookback
        (4, 3),  # 4 clusters, 3-month lookback
    ]

    best_rep = None
    best_score = -999.0

    for n_clusters, lookback_months in configs:
        print(f"\n{'='*60}")
        print(f"CONFIG: n_clusters={n_clusters}, lookback_months={lookback_months}")
        print('='*60)
        rep = run_config(n_clusters, lookback_months)
        print(al.fmt(rep))
        print("JSON:", al.as_json(rep))

        score = rep.get("verdict", {}).get("best_holdout_sharpe", -999.0) or -999.0
        if score > best_score:
            best_score = score
            best_rep = rep

    print("\n" + "="*60)
    print("BEST CONFIG:")
    print("="*60)
    print(al.fmt(best_rep))
    print("JSON:", al.as_json(best_rep))