"""STA02 — Walk-forward Logistic Regression on TA features (1d).

Idea: a logistic classifier is periodically re-fit on features
{rsi, zscore_price, momentum, realized_vol} all computed causally.
Predict P(next bar up) -> long if P > 0.5, else flat (long-only, no short).

Causal contract
---------------
At decision bar d (close[d] known):
  - features use data up to and including close[d]
  - we predict: will close[d+1] > close[d] ?
  - target[d] = position held during bar d+1
  - altlib eval_weights shifts by 1 for us -> no double shift

Feature construction (all using data <= close[d]):
  - rsi_14:     RSI(14) at bar d
  - zscore_20:  (close[d] - sma_20[d]) / std_20[d]
  - mom_10:     log(close[d] / close[d-10]) (10-bar momentum)
  - rvol_20:    realized annualized vol, 20-bar window

Training label:
  - y[k] = 1 if close[k+1] > close[k], else 0
  - Train on (X[k], y[k]) for k in [warmup .. d-1]

Grid (4 configs x 1 TF = 4 total backtests <= 6 limit):
  - min_train_years: 1.0 or 2.0
  - C (inverse regularization): 0.1 or 1.0

Best config by min(BTC, ETH) hold-out Sharpe.
"""

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")

from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler


def logistic_target(df, min_train_years: float = 1.0, C: float = 1.0) -> np.ndarray:
    """
    Walk-forward logistic on {rsi14, zscore20, mom10, rvol20}.
    Returns vol-targeted position array (target[i] decided at close[i]).
    """
    c = df["close"].values.astype(float)
    n = len(c)
    bpy = al.bars_per_year(df)
    bpd = al.bars_per_day(df)

    # --- build features (all causal at bar i) ---
    # RSI 14
    feat_rsi = al.rsi(c, win=14)

    # Z-score of close over 20-bar window
    feat_zsc = al.zscore(c, win=20)

    # 10-bar log-momentum: log(close[i] / close[i-10])
    # Using lag=10 bars; only valid for i >= 10
    feat_mom = np.full(n, np.nan)
    lag = 10
    feat_mom[lag:] = np.log(c[lag:] / c[:-lag])

    # Realized annualized vol (20-bar)
    r = al.simple_returns(c)
    feat_rvol = al.realized_vol(r, win=20, bars_per_year=bpy)

    # Stack into feature matrix [n x 4]
    X_all = np.column_stack([feat_rsi, feat_zsc, feat_mom, feat_rvol])

    # Label: 1 if next bar close > current close, else 0
    # y[i] = 1 if close[i+1] > close[i] — shape (n,), y[n-1] is undefined
    y_all = np.zeros(n, dtype=float)
    y_all[:-1] = (c[1:] > c[:-1]).astype(float)

    min_train_bars = int(min_train_years * bpy)
    # Need at least warmup + lags for first valid sample
    first_valid = max(20, lag)  # 20 for zscore/rvol, 10 for mom
    # first training sample k: k >= first_valid AND feature X[k] fully defined
    # first prediction at bar d: d >= first_valid + min_train_bars
    first_pred = first_valid + min_train_bars

    # Refit quarterly
    refit_every = max(1, int(bpy / 4))

    direction = np.zeros(n, dtype=float)
    last_refit = -refit_every  # force first refit
    model = LogisticRegression(C=C, solver="lbfgs", max_iter=500,
                                random_state=42, class_weight="balanced")
    scaler = StandardScaler()
    trained = False

    for d in range(first_pred, n - 1):
        if d - last_refit >= refit_every:
            # Build training set: samples k in [first_valid, d-1] (predict y[k] from X[k])
            # X[k] causal (uses data <= close[k]), y[k] requires close[k+1] (NOT at k, at k+1)
            # So the last valid training sample is k = d-1 (we know close[d] = close[(d-1)+1])
            k_start = first_valid
            k_end = d  # exclusive, so training on [k_start, d-1]

            if k_end - k_start < 30:
                continue

            X_tr = X_all[k_start:k_end]
            y_tr = y_all[k_start:k_end]

            # Drop rows with NaN features
            valid_mask = np.all(np.isfinite(X_tr), axis=1) & np.isfinite(y_tr)
            if valid_mask.sum() < 20:
                continue
            X_tr = X_tr[valid_mask]
            y_tr = y_tr[valid_mask]

            # Check both classes present
            if len(np.unique(y_tr)) < 2:
                continue

            try:
                scaler.fit(X_tr)
                X_tr_scaled = scaler.transform(X_tr)
                model.fit(X_tr_scaled, y_tr)
                trained = True
                last_refit = d
            except Exception:
                continue

        if not trained:
            continue

        # Predict at bar d: features X_all[d]
        x_d = X_all[d]
        if not np.all(np.isfinite(x_d)):
            continue

        x_scaled = scaler.transform(x_d.reshape(1, -1))
        prob_up = model.predict_proba(x_scaled)[0]
        # class order: model.classes_ = [0, 1]
        idx_up = list(model.classes_).index(1) if 1 in model.classes_ else 1
        p_up = prob_up[idx_up]

        # Long if P(up) > 0.5, else flat (long-only, no short)
        direction[d] = 1.0 if p_up > 0.5 else 0.0

    # Vol-target the direction signal
    target = al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
    return target


def run_grid():
    configs = [
        dict(min_train_years=1.0, C=0.1),
        dict(min_train_years=1.0, C=1.0),
        dict(min_train_years=2.0, C=0.1),
        dict(min_train_years=2.0, C=1.0),
    ]

    best_rep = None
    best_holdout = -999.0

    for cfg in configs:
        name = f"STA02(train={cfg['min_train_years']}y,C={cfg['C']})"
        print(f"\n--- Running {name} ---")
        rep = al.study_weights(
            name,
            lambda df, c=cfg: logistic_target(df, **c),
            tfs=("1d",)
        )
        print(al.fmt(rep))

        min_hold = rep["verdict"].get("best_holdout_sharpe", -999.0)
        if min_hold > best_holdout:
            best_holdout = min_hold
            best_rep = rep
            best_rep["_cfg"] = cfg

    return best_rep


if __name__ == "__main__":
    best = run_grid()
    print("\n\n=== BEST CONFIG ===")
    print(al.fmt(best))
    print("JSON:", al.as_json(best))