"""TRD04 — Supertrend(period, multiplier)
Classic ATR-band trend flip: long when price above supertrend line, short/flat below.
Grid: (period, mult) in [(10,3),(14,3),(10,2),(14,2)] — 4 configs x 2 TFs x 2 assets = 16 backtests.
Style: continuous weights (vol-targeted, long-flat).
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np
import pandas as pd


def supertrend_direction(df: pd.DataFrame, period: int = 10, mult: float = 3.0) -> np.ndarray:
    """Compute Supertrend and return causal direction in {0, 1}.
    Long (1) when close > supertrend, flat (0) otherwise.
    The Supertrend uses ATR-based bands and flips only when price crosses the band.
    Causal: at bar i we use data up to and including close[i].
    """
    h = df["high"].values.astype(float)
    l = df["low"].values.astype(float)
    c = df["close"].values.astype(float)
    n = len(c)

    # ATR via EWM (causal, same as al.atr)
    a = al.atr(df, period)

    hl2 = (h + l) / 2.0
    upper = hl2 + mult * a
    lower = hl2 - mult * a

    # Final upper/lower bands (adjusted to not widen against trend)
    final_upper = upper.copy()
    final_lower = lower.copy()
    direction = np.zeros(n, dtype=float)  # 1 = uptrend (long), 0 = downtrend (flat)

    # Warm-up: first bar
    final_upper[0] = upper[0]
    final_lower[0] = lower[0]
    direction[0] = 1.0 if c[0] > hl2[0] else 0.0

    for i in range(1, n):
        # Tighten upper: new upper only replaces if lower than previous (or if prev close was above)
        if upper[i] < final_upper[i-1] or c[i-1] > final_upper[i-1]:
            final_upper[i] = upper[i]
        else:
            final_upper[i] = final_upper[i-1]

        # Tighten lower: new lower only replaces if higher than previous (or if prev close was below)
        if lower[i] > final_lower[i-1] or c[i-1] < final_lower[i-1]:
            final_lower[i] = lower[i]
        else:
            final_lower[i] = final_lower[i-1]

        # Determine direction (trend)
        prev_dir = direction[i-1]
        if prev_dir == 0.0:  # was downtrend (flat)
            if c[i] > final_upper[i]:
                direction[i] = 1.0  # flip to uptrend
            else:
                direction[i] = 0.0  # stay flat
        else:  # was uptrend
            if c[i] < final_lower[i]:
                direction[i] = 0.0  # flip to downtrend (flat)
            else:
                direction[i] = 1.0  # stay in uptrend

    return direction


def make_target(period: int, mult: float):
    """Returns a target_fn(df) that computes vol-targeted Supertrend weights."""
    def target_fn(df: pd.DataFrame) -> np.ndarray:
        direction = supertrend_direction(df, period=period, mult=mult)
        # vol-targeted: scale by realized vol, cap at 2x leverage, long-flat only
        return al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
    return target_fn


# Small internal grid: 4 param sets
GRID = [
    (10, 3.0),
    (14, 3.0),
    (10, 2.0),
    (14, 2.0),
]

TFS = ("1d", "12h")

# Run each config on both TFs
best_rep = None
best_score = -999.0

print("=== TRD04: Supertrend Grid Search ===")
for period, mult in GRID:
    label = f"TRD04-ST({period},{mult})"
    fn = make_target(period, mult)
    rep = al.study_weights(label, fn, tfs=TFS)
    v = rep["verdict"]
    score = v.get("best_holdout_sharpe", -9)
    print(al.fmt(rep))
    print()
    if score > best_score:
        best_score = score
        best_rep = rep
        best_period = period
        best_mult = mult

print("\n" + "="*60)
print(f"BEST CONFIG: period={best_period}, mult={best_mult}")
print(al.fmt(best_rep))
print("JSON:", al.as_json(best_rep))