"""STA01 — Ridge on lagged returns (1d only).

Walk-forward expanding-window Ridge regression that predicts next-bar return sign
from lagged log-returns (lags 1..10). Position = sign(prediction) vol-targeted.

Key causal rule: at bar i, we have log_return[i] = log(close[i]/close[i-1]).
We predict return[i+1], so we build features from lags 1..10 ending at lag 1
relative to i, meaning we use returns[i-1], returns[i-2], ..., returns[i-10].
This is strictly causal: no return from bar i is used in the feature vector for
the prediction that drives the position held during bar i+1.

The lib's eval_weights shift handles the final no-lookahead guarantee:
  target[i] -> position held during bar i+1.
We set target[i] = sign of prediction made at close[i] using lags ending at i-1.

Grid (<=4 sets, 1 TF -> 4 total backtests, well within 6 limit):
  - min_train_years: 1 or 2  (warm-up before first prediction)
  - alpha: 1.0 or 10.0       (ridge regularization)
Best config chosen by min(BTC,ETH) holdout Sharpe.
"""
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 Ridge

N_LAGS = 10  # lags 1..10 (i.e. features use returns[i-1]..returns[i-10])


def ridge_target(df, min_train_years: float = 2.0, alpha: float = 1.0) -> np.ndarray:
    """
    Walk-forward expanding-window Ridge: predict sign of next-bar log-return.
    Feature at bar i: [ret[i-1], ret[i-2], ..., ret[i-10]]  <- strictly causal.
    Output target[i] = vol-targeted position decided at bar i.
    """
    c = df["close"].values.astype(float)
    lr = al.log_returns(c)  # lr[k] = log(close[k]/close[k-1]), lr[0]=0

    n = len(lr)
    bpy = al.bars_per_year(df)
    min_train_bars = int(min_train_years * bpy) + N_LAGS

    # raw signal array (before vol targeting)
    direction = np.zeros(n, dtype=float)

    # Walk-forward: at each bar i, we have features built from lags 1..N_LAGS
    # i.e. X[i] = [lr[i-1], lr[i-2], ..., lr[i-N_LAGS]]
    # We predict lr[i+1] sign, so we train on (X[k], lr[k+1]) for all k < i
    # where we have N_LAGS lags available (k >= N_LAGS).
    # The first valid feature row is at k = N_LAGS (uses lr[N_LAGS-1]..lr[0]).
    # We need min_train_bars samples before making the first prediction.

    # Build full feature matrix: row k uses lr[k-1]..lr[k-N_LAGS]
    # valid for k >= N_LAGS
    # target for row k: lr[k] (we're predicting the return at bar k)
    # Training on pairs: (X[k], lr[k]) means we're predicting current bar return
    # from lagged features — used to predict what comes next.
    # Specifically: predict lr[i] using X[i] = [lr[i-1]..lr[i-N_LAGS]]
    # Position at bar i-1 (decided at close[i-1]) will hold during bar i.
    # So in altlib terms: target[i-1] = sign(predict lr[i]) via X[i] = [lr[i-1]..lr[i-N_LAGS]]
    # But X[i] uses lr[i-1] which is available at close[i-1].
    # Therefore: at close[i-1], we have lr[i-1]..lr[i-N_LAGS] -> predict lr[i] -> target[i-1].

    # Let's index: prediction at "decision bar" d means:
    #   features: [lr[d], lr[d-1], ..., lr[d-N_LAGS+1]]  (all available at close[d])
    #   prediction target: lr[d+1]
    #   train on (X[k], lr[k+1]) for k = N_LAGS-1 .. d-1
    #     X[k] = [lr[k], lr[k-1], ..., lr[k-N_LAGS+1]]
    # First prediction: d = min_train_bars - 1 (0-indexed), need d >= N_LAGS-1 and d-1 >= N_LAGS-1+1

    first_pred_d = max(N_LAGS, min_train_bars - 1)

    model = Ridge(alpha=alpha, fit_intercept=True)
    trained = False

    for d in range(first_pred_d, n - 1):
        # Build training set: samples k from (N_LAGS-1) to (d-1)
        # X[k] = [lr[k], lr[k-1], ..., lr[k-N_LAGS+1]], y[k] = lr[k+1]
        # We rebuild only when needed; for efficiency, fit incrementally isn't
        # trivial with sklearn, so we do a periodic refit every 'refit_every' bars
        # to keep runtime manageable.
        pass

    # Vectorized approach for speed: refit every refit_every bars
    refit_every = max(1, int(bpy / 4))  # quarterly refit

    last_refit = -refit_every  # force first refit

    for d in range(first_pred_d, n - 1):
        if d - last_refit >= refit_every:
            # Build full training set up to d-1
            # k ranges from N_LAGS-1 to d-1
            k_start = N_LAGS - 1
            k_end = d  # exclusive (train up to d-1 inclusive)
            if k_end - k_start < 10:
                continue
            # Build X matrix
            rows = k_end - k_start
            X_train = np.zeros((rows, N_LAGS))
            y_train = np.zeros(rows)
            for row_i, k in enumerate(range(k_start, k_end)):
                # X[k] = [lr[k], lr[k-1], ..., lr[k-N_LAGS+1]]
                X_train[row_i] = lr[k - N_LAGS + 1: k + 1][::-1]  # lag1=lr[k], lag10=lr[k-N_LAGS+1]
                y_train[row_i] = lr[k + 1]
            model.fit(X_train, y_train)
            trained = True
            last_refit = d

        if not trained:
            continue

        # Predict lr[d+1] using [lr[d], lr[d-1], ..., lr[d-N_LAGS+1]]
        x_pred = lr[d - N_LAGS + 1: d + 1][::-1].reshape(1, -1)
        pred = model.predict(x_pred)[0]
        direction[d] = np.sign(pred) if pred != 0 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, alpha=1.0),
        dict(min_train_years=1.0, alpha=10.0),
        dict(min_train_years=2.0, alpha=1.0),
        dict(min_train_years=2.0, alpha=10.0),
    ]

    best_rep = None
    best_holdout = -999.0

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

        # Extract min holdout Sharpe across assets/cells
        min_hold = rep["verdict"].get("best_holdout_sharpe", -999)
        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))