"""STA08 — AR(1) residual reversion.

IDEA: Fit an expanding-window AR(1) on log returns. The AR(1) residual is
r[t] - (a0 + a1 * r[t-1]), where a0 and a1 are estimated causally from all
data up to t-1. Trade the mean-reversion of the residual: if residual is
positive (return exceeded AR(1) prediction) we expect reversion → short;
if negative → long.

Signal: z-score the residual over a rolling window, take the negative of it
as the continuous position (mean-reversion), then vol-target it.

Grid: 2 lookback windows for z-scoring (60, 120 bars), tested on 1d and 12h.
Total cells: 2 TFs × 2 params × 2 assets = 8 backtests — within limit.

We pick the best config by min-asset hold-out Sharpe.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np


def ar1_residual_target(df, zscore_win: int = 60) -> np.ndarray:
    """
    Causal AR(1) residual reversion target.

    At each bar i:
      - Use all returns r[0..i-1] to fit AR(1): regress r[t] on r[t-1]
        (expanding OLS — efficient via running sums)
      - Compute residual[i] = r[i] - (a0 + a1 * r[i-1])  (uses closed bar i)
      - Z-score the residual over last zscore_win bars
      - Position = -z (mean-reversion) → vol-targeted

    Minimum warmup: 30 bars for stable OLS + zscore_win bars for z-score.
    """
    c = df["close"].values.astype(float)
    n = len(c)
    r = al.log_returns(c)  # r[0]=0, r[i] = log(c[i]/c[i-1])

    # Expanding AR(1): for each bar i, estimate (a0, a1) from data up to i-1.
    # We need: sum(r), sum(r^2), sum(r_t * r_{t-1}), sum(r_{t-1}), sum(r_{t-1}^2)
    # for t in [1..i-1].
    # Then OLS: regress r_t ~ a0 + a1*r_{t-1}.
    # Normal equations:
    #   [n-1,   sum_r1 ] [a0]   [sum_r ]
    #   [sum_r1, sum_r1sq] [a1] = [sum_r_r1]
    # where sum_r1 = sum(r[t-1]), sum_r = sum(r[t]), etc.

    residuals = np.zeros(n)
    min_warmup = 30  # minimum bars to fit AR(1)

    # Running sums for expanding OLS (using pairs (r[t-1], r[t]) for t>=1)
    S_n = 0.0       # count of pairs
    S_x = 0.0       # sum of r[t-1]
    S_y = 0.0       # sum of r[t]
    S_xx = 0.0      # sum of r[t-1]^2
    S_xy = 0.0      # sum of r[t-1]*r[t]

    for i in range(1, n):
        # Update running sums with pair (r[i-1], r[i]) but we use data up to i-1
        # So at step i, we first compute residual using sums from [1..i-1],
        # then update sums to include pair for t=i.

        if S_n >= min_warmup:
            # Fit AR(1) from expanding window up to t=i-1
            denom = S_n * S_xx - S_x * S_x
            if abs(denom) > 1e-14:
                a1 = (S_n * S_xy - S_x * S_y) / denom
                a0 = (S_y - a1 * S_x) / S_n
            else:
                a0, a1 = 0.0, 0.0
            # Residual at bar i: actual r[i] minus AR(1) prediction
            pred = a0 + a1 * r[i - 1]
            residuals[i] = r[i] - pred
        # else: residuals[i] remains 0

        # Update running sums with the new observation pair (r[i-1], r[i])
        # This is data point for t=i: x=r[i-1], y=r[i]
        S_n += 1.0
        S_x += r[i - 1]
        S_y += r[i]
        S_xx += r[i - 1] ** 2
        S_xy += r[i - 1] * r[i]

    # Z-score the residual with rolling window
    z = al.zscore(residuals, zscore_win)

    # Mean-reversion: negative of z-score
    direction = -z
    direction = np.nan_to_num(direction, nan=0.0)

    # Vol-target to 20% annualized, cap at 2x leverage
    target = al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
    return target


def make_target(zscore_win: int):
    return lambda df: ar1_residual_target(df, zscore_win=zscore_win)


if __name__ == "__main__":
    # Small internal grid: 2 z-score windows × 2 TFs = 4 cells per config
    # Pick best by min-asset holdout Sharpe
    configs = [
        {"zscore_win": 60,  "label": "z60"},
        {"zscore_win": 120, "label": "z120"},
    ]
    tfs = ("1d", "12h")

    best_rep = None
    best_score = -9.0

    for cfg in configs:
        zw = cfg["zscore_win"]
        rep = al.study_weights(
            f"STA08-AR1resid-z{zw}",
            make_target(zw),
            tfs=tfs,
        )
        score = rep["verdict"].get("best_holdout_sharpe", -9.0)
        if score > best_score:
            best_score = score
            best_rep = rep
        # Print intermediate for debug
        print(f"\n--- Config z{zw} ---")
        print(al.fmt(rep))

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