"""STA03 — Random Forest direction (walk-forward, causal, long-flat).

Idea:
  Small RF (50 trees, max_depth 4) trained walk-forward on causal features decided at
  close[i-1]. Features: multi-period returns, RSI, vol ratio, trend signals (EMA crossovers).
  Predicts binary direction of next bar (1=up, 0=down/flat). Position = predicted probability
  of up, vol-targeted, long-flat only (clip to [0, leverage_cap]).

Walk-forward:
  - Train window: 252 bars (1 year of 1d data; ~252*8 for shorter TF but we stay 1d)
  - Retrain every 63 bars (quarterly)
  - Min 252 bars before first prediction; otherwise position=0

Causal guarantee:
  Feature for bar i uses returns/indicators up to close[i].
  Target for bar i is sign(close[i+1]/close[i] - 1) = r[i+1] sign.
  During training we shift: X[t], y[t] = direction of bar t+1.
  At prediction time we use X[i] -> predicted prob of next bar going up -> position[i].
  altlib eval_weights then holds position[i] during bar i+1 (the shift is done for us).
  No leak.

Grid (<=4 configs, total backtests <=6 since only 1d TF):
  A: train_win=252, retrain=63, n_estimators=50, max_depth=4
  B: train_win=365, retrain=63, n_estimators=50, max_depth=3
  C: train_win=252, retrain=21, n_estimators=50, max_depth=4  (monthly retrain)
  D: train_win=365, retrain=126, n_estimators=100, max_depth=4 (semi-annual retrain)

Pick best by min_asset_holdout_sharpe on 1d.
"""

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

try:
    from sklearn.ensemble import RandomForestClassifier
except ImportError:
    print("ERROR: scikit-learn not available")
    sys.exit(1)


def build_features(df):
    """Build a causal feature matrix. Feature at row i uses data up to close[i].
    Returns X array shape (N, n_features). First ~30 rows will have NaN -> handled."""
    c = df["close"].values.astype(float)
    N = len(c)

    # Returns at various horizons (causal: r[i] = close[i]/close[i-1] - 1)
    r = al.simple_returns(c)
    r1 = r                                       # 1-bar return
    r5 = np.zeros(N); r5[5:] = c[5:] / c[:-5] - 1   # 5-bar
    r10 = np.zeros(N); r10[10:] = c[10:] / c[:-10] - 1
    r21 = np.zeros(N); r21[21:] = c[21:] / c[:-21] - 1
    r63 = np.zeros(N); r63[63:] = c[63:] / c[:-63] - 1

    # RSI
    rsi14 = al.rsi(c, 14)

    # Vol ratio: short vol / long vol (vol regime)
    rv_short = al.realized_vol(r, 10, al.bars_per_year(df))
    rv_long = al.realized_vol(r, 30, al.bars_per_year(df))
    vol_ratio = np.where(rv_long > 0, rv_short / rv_long, 1.0)

    # EMA crossovers
    ema10 = al.ema(c, 10)
    ema21 = al.ema(c, 21)
    ema50 = al.ema(c, 50)
    cross_fast = (ema10 - ema21) / np.where(ema21 > 0, ema21, 1e-8)
    cross_slow = (ema21 - ema50) / np.where(ema50 > 0, ema50, 1e-8)

    # Z-score of price
    z21 = al.zscore(c, 21)
    z63 = al.zscore(c, 63)

    # ATR-normalized range (volatility clustering proxy)
    atr14 = al.atr(df, 14)
    atr_ratio = np.where(c > 0, atr14 / c, 0.0)

    X = np.column_stack([
        r1, r5, r10, r21, r63,
        rsi14,
        vol_ratio,
        cross_fast, cross_slow,
        z21, z63,
        atr_ratio,
    ])
    return X


def make_target_fn(train_win: int, retrain_every: int,
                   n_estimators: int, max_depth: int):
    """Return a target_fn(df) -> prob array in [0,1] for long-flat vol-targeted pos."""

    def target_fn(df):
        c = df["close"].values.astype(float)
        N = len(c)
        X = build_features(df)

        # Future direction: y[i] = 1 if close[i+1] > close[i], else 0
        # We train on (X[t], y[t]) where y[t] is known at t+1
        # At prediction time for bar i, we have X[i] and predict prob(up next bar)
        y = np.zeros(N, dtype=int)
        y[:-1] = (c[1:] > c[:-1]).astype(int)  # y[N-1] unknown, set 0 (unused)

        prob_up = np.zeros(N)
        last_retrain = -retrain_every  # force retrain at first opportunity
        clf = None

        for i in range(train_win, N):
            # Retrain if due
            if i - last_retrain >= retrain_every or clf is None:
                # Training data: indices [i-train_win .. i-1]
                # X_train[t] -> y_train[t] = direction of bar t+1
                # We use t from i-train_win to i-2 (y[i-1] = direction of bar i = known)
                start = i - train_win
                end = i - 1  # last sample where y is known (y[i-1] is direction of bar i = close[i]/close[i-1]-1)
                X_tr = X[start:end]
                y_tr = y[start:end]

                # Drop rows with NaN in features
                valid = np.all(np.isfinite(X_tr), axis=1)
                X_tr_v = X_tr[valid]
                y_tr_v = y_tr[valid]

                if len(X_tr_v) > 50 and len(np.unique(y_tr_v)) > 1:
                    clf = RandomForestClassifier(
                        n_estimators=n_estimators,
                        max_depth=max_depth,
                        random_state=42,
                        n_jobs=1,
                    )
                    clf.fit(X_tr_v, y_tr_v)
                    last_retrain = i
                else:
                    clf = None  # insufficient data

            # Predict probability for bar i
            if clf is not None and np.all(np.isfinite(X[i])):
                p = clf.predict_proba(X[i:i+1])
                # Find prob of class 1 (up)
                classes = list(clf.classes_)
                if 1 in classes:
                    prob_up[i] = p[0][classes.index(1)]
                else:
                    prob_up[i] = 0.0
            else:
                prob_up[i] = 0.5  # neutral when no model

        # Convert probability to direction signal: prob > 0.5 -> long, else flat
        # Use soft threshold: direction = 2*(prob_up - 0.5), clipped to [0,1]
        # This gives continuous [0,1] position proportional to confidence
        direction = np.clip(2 * (prob_up - 0.5), 0.0, 1.0)
        direction[:train_win] = 0.0  # no position before warmup

        # Apply vol targeting (long-flat, no short)
        pos = al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
        pos = np.clip(pos, 0.0, 2.0)  # long-flat
        return pos

    return target_fn


# Grid of configs
CONFIGS = [
    dict(name="A", train_win=252, retrain_every=63, n_estimators=50, max_depth=4),
    dict(name="B", train_win=365, retrain_every=63, n_estimators=50, max_depth=3),
    dict(name="C", train_win=252, retrain_every=21, n_estimators=50, max_depth=4),
    dict(name="D", train_win=365, retrain_every=126, n_estimators=100, max_depth=4),
]

print("STA03 — Random Forest direction (walk-forward, causal, long-flat)")
print(f"Grid: {len(CONFIGS)} configs on 1d only (total backtests = {len(CONFIGS)*2})")
print()

results = []
for cfg in CONFIGS:
    print(f"Config {cfg['name']}: train_win={cfg['train_win']}, "
          f"retrain={cfg['retrain_every']}, trees={cfg['n_estimators']}, depth={cfg['max_depth']}")
    fn = make_target_fn(
        train_win=cfg["train_win"],
        retrain_every=cfg["retrain_every"],
        n_estimators=cfg["n_estimators"],
        max_depth=cfg["max_depth"],
    )
    rep = al.study_weights(
        f"STA03-RF-{cfg['name']}",
        fn,
        tfs=("1d",),
    )
    print(al.fmt(rep))
    print()
    results.append((cfg, rep))

# Pick best by min_asset_holdout_sharpe
best_cfg, best_rep = max(
    results,
    key=lambda x: x[1]["verdict"].get("best_holdout_sharpe", -99)
)
print("=" * 60)
print(f"BEST CONFIG: {best_cfg['name']} "
      f"(train_win={best_cfg['train_win']}, retrain={best_cfg['retrain_every']}, "
      f"trees={best_cfg['n_estimators']}, depth={best_cfg['max_depth']})")
print()

# Re-label report as STA03 canonical
best_rep["name"] = "STA03"
print(al.fmt(best_rep))
print()
print("JSON:", al.as_json(best_rep))