"""STA07 — Online SGD Logistic Regression (next-bar sign prediction)
Hypothesis: An online logistic classifier (sklearn SGDClassifier with partial_fit) is
updated bar-by-bar using causal features and predicts the sign of the NEXT bar's return.
The prediction confidence (decision_function score) is used as a continuous position
(long if positive score, short/flat if negative — but long-only via clip to [0,1]).

Features (all causal at bar i):
  - short EMA vs long EMA ratio (trend)
  - RSI(14) normalized to [-1,1]
  - z-score of close over 20 bars
  - realized vol ratio (fast / slow) as regime indicator
  - log return of last bar (momentum/mean-reversion signal)
  - ATR normalized (relative volatility)

The label for bar i is: sign(close[i+1] / close[i] - 1)
  -> at decision time i we don't have i+1 yet, but we use PAST labels to train.
  -> Specifically, we do partial_fit at bar i using features[i-1] and label[i-1]
     (the actual outcome that just resolved), then predict at bar i using features[i].
  -> This is fully causal: model at bar i trained only on history ending at close[i-1].

Grid: 2 warmup periods (60 / 120 bars) × 2 TFs (1d / 12h) = 4 total cells (<=6 limit).
Best config selected by min_asset_holdout_sharpe across all cells.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np
from sklearn.linear_model import SGDClassifier
from sklearn.preprocessing import StandardScaler


def online_sgd_logistic_target(df: "pd.DataFrame", warmup: int = 60) -> np.ndarray:
    """
    Online SGD logistic regression updated each bar.

    Causality:
      At bar i:
        1. We receive outcome from bar i-1 (sign of return from close[i-2] to close[i-1]).
        2. We do partial_fit(features[i-1], label[i-1]) — update model.
        3. We predict at features[i] -> continuous score via decision_function.
        4. Position = clip(score, 0, 1) to stay long-flat, then vol-target.

    The model is never trained on data beyond close[i-1] when producing the position for
    bar i+1 (altlib shifts pos by 1 internally). So there is no look-ahead.
    """
    c = df["close"].values.astype(float)
    n = len(c)

    # --- Causal features computed once vectorially ---
    r = al.log_returns(c)
    ema_fast = al.ema(c, 10)
    ema_slow = al.ema(c, 40)
    ema_ratio = np.where(ema_slow > 0, ema_fast / ema_slow - 1.0, 0.0)

    rsi14 = al.rsi(c, 14)
    rsi_norm = (rsi14 - 50.0) / 50.0  # normalize to [-1, 1]

    zsc = al.zscore(c, 20)
    zsc = np.nan_to_num(zsc, nan=0.0)

    rv_fast = al.realized_vol(r, 5, al.bars_per_year(df))
    rv_slow = al.realized_vol(r, 20, al.bars_per_year(df))
    rv_ratio = np.where((rv_slow > 0) & np.isfinite(rv_slow) & np.isfinite(rv_fast),
                        rv_fast / rv_slow - 1.0, 0.0)

    atr14 = al.atr(df, 14)
    atr_norm = np.where(c > 0, atr14 / c, 0.0)

    # Feature matrix [n, 6]
    X = np.column_stack([
        ema_ratio,
        rsi_norm,
        zsc,
        rv_ratio,
        r,           # last bar return (known at bar i)
        atr_norm,
    ])
    X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0)

    # Labels: sign of NEXT return (for training only; not used in prediction)
    # label[i] = sign(r[i+1]): known at bar i+1, used to update model at bar i+1
    labels = np.sign(np.roll(r, -1))  # peek-ahead in labels array only
    # But we access labels[i-1] at bar i -> labels[i-1] = sign(r[i]) which is known at i
    # So: when we update at bar i, we use label[i-1] = sign(r[i-1+1]) = sign(r[i])
    # r[i] = log(close[i]/close[i-1]) — fully known at bar i. Causal. ✓

    # Online SGD Logistic
    clf = SGDClassifier(
        loss="log_loss",
        penalty="l2",
        alpha=1e-4,
        learning_rate="optimal",
        random_state=42,
        max_iter=1,
        warm_start=True,
    )

    scores = np.zeros(n)
    classes = np.array([-1, 1])

    for i in range(1, n):
        # Update model: use features[i-1] and label[i-1] (=sign(r[i]), known at i)
        label_i_minus_1 = int(np.sign(r[i]))  # sign of return from close[i-1] to close[i]
        if label_i_minus_1 == 0:
            label_i_minus_1 = 1  # tie-break: treat flat as up

        feat = X[i - 1].reshape(1, -1)

        # Only partial_fit after warmup — before that, accumulate without predicting
        try:
            clf.partial_fit(feat, [label_i_minus_1], classes=classes)
        except Exception:
            pass

        # Predict at bar i if model has been fitted (after warmup)
        if i >= warmup:
            try:
                score = clf.decision_function(X[i].reshape(1, -1))[0]
                scores[i] = score
            except Exception:
                scores[i] = 0.0
        else:
            scores[i] = 0.0

    # Convert decision score to long-flat position in [0, 1]
    # Use tanh to squash to (-1, 1), then clip to [0, 1] for long-flat
    pos_raw = np.tanh(scores)  # in (-1, 1)
    pos_lf = np.clip(pos_raw, 0.0, 1.0)  # long-flat

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


def make_target(warmup: int):
    def target_fn(df):
        return online_sgd_logistic_target(df, warmup=warmup)
    return target_fn


if __name__ == "__main__":
    configs = [
        ("warmup60", 60),
        ("warmup120", 120),
    ]

    results = []
    for label, warmup in configs:
        print(f"\n--- Running STA07 config: {label} ---")
        rep = al.study_weights(
            f"STA07-OnlineSGD-{label}",
            make_target(warmup),
            tfs=("1d", "12h"),
        )
        print(al.fmt(rep))
        print("JSON:", al.as_json(rep))
        results.append((label, warmup, rep))

    # Pick best config by best_holdout_sharpe from verdict
    best_label, best_warmup, best_rep = max(
        results,
        key=lambda x: x[2]["verdict"].get("best_holdout_sharpe", -99)
    )
    print(f"\n=== BEST CONFIG: {best_label} ===")
    print(al.fmt(best_rep))
    print("JSON:", al.as_json(best_rep))