"""BRK02 — Donchian55 + Chandelier ATR trailing stop.

IDEA:
  - Enter LONG when close[i] breaks above the 55-bar Donchian high (prior bars, causal).
  - Exit (go flat) when close[i] falls below the Chandelier trailing stop:
      chandelier_stop = highest_close(22 bars, ending at i) - 3 * ATR(22, ending at i).
  - Position is 0 (flat) or 1 (long). No short. Vol-targeted to 20% with 2x cap.

Implementation (weights style, continuous position):
  - Donchian high computed on PRIOR bars (shift(1) already done by al.donchian).
  - Chandelier stop computed causally on current+prior bars:
      hc[i] = max(close[i-21..i])   -> rolling max of close, window=22
      atr22[i] = ATR(22 bars) at i
      stop[i] = hc[i] - 3 * atr22[i]
  - State machine:
      if flat and close[i] > donchian_high[i]: go long
      if long and close[i] < stop[i]: go flat

Params tried: (don_win=55, atr_win=22, atr_mult=3.0) — canonical
              (don_win=40, atr_win=22, atr_mult=2.5)  — tighter
Best picked by min_asset_holdout_sharpe.
"""

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 chandelier_signal(df: pd.DataFrame, don_win: int = 55,
                      atr_win: int = 22, atr_mult: float = 3.0) -> np.ndarray:
    """Return a direction array in {0, 1} (0=flat, 1=long) using Donchian+Chandelier.
    Causal: decision at i uses only data <= close[i]."""
    close = df["close"].values.astype(float)
    n = len(close)

    # Donchian upper channel: highest HIGH over prior don_win bars (shift(1) in al.donchian)
    don_high, _ = al.donchian(df, don_win)   # don_high[i] = max(high[i-don_win..i-1])

    # ATR(atr_win) — causal, uses bars up to and including i
    atr22 = al.atr(df, atr_win)

    # Highest CLOSE over trailing atr_win bars (inclusive of i) — causal
    highest_close = pd.Series(close).rolling(atr_win, min_periods=atr_win).max().values

    # Chandelier stop at i
    chandelier_stop = highest_close - atr_mult * atr22

    # State machine: flat=0, long=1
    pos = np.zeros(n, dtype=float)
    state = 0  # start flat
    for i in range(n):
        c = close[i]
        dh = don_high[i]
        cs = chandelier_stop[i]

        if state == 0:
            # Enter long if close breaks above prior Donchian high (valid only if dh is defined)
            if np.isfinite(dh) and c > dh:
                state = 1
        else:  # state == 1
            # Exit long if close drops below chandelier stop (and stop is defined)
            if np.isfinite(cs) and c < cs:
                state = 0

        pos[i] = float(state)

    return pos


def make_target(don_win: int = 55, atr_win: int = 22, atr_mult: float = 3.0):
    """Factory returning a vol-targeted weight function for a given param set."""
    def target_fn(df: pd.DataFrame) -> np.ndarray:
        direction = chandelier_signal(df, don_win=don_win, atr_win=atr_win, atr_mult=atr_mult)
        return al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
    return target_fn


# --- Small param grid (4 configurations, TFs: 1d and 12h -> 8 backtest cells total)
CONFIGS = [
    dict(don_win=55, atr_win=22, atr_mult=3.0, label="D55-A22-M3.0"),
    dict(don_win=40, atr_win=22, atr_mult=2.5, label="D40-A22-M2.5"),
    dict(don_win=55, atr_win=14, atr_mult=3.0, label="D55-A14-M3.0"),
    dict(don_win=70, atr_win=22, atr_mult=3.5, label="D70-A22-M3.5"),
]

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

best_rep = None
best_score = -999.0

for cfg in CONFIGS:
    lbl = cfg["label"]
    fn = make_target(don_win=cfg["don_win"], atr_win=cfg["atr_win"], atr_mult=cfg["atr_mult"])
    rep = al.study_weights(f"BRK02[{lbl}]", fn, tfs=TFS)
    score = rep["verdict"].get("best_holdout_sharpe", -9)
    print(f"Config {lbl}: grade={rep['verdict']['grade']} score(hold)={score:.3f}")
    if score > best_score:
        best_score = score
        best_rep = rep

# Rename best result to canonical BRK02
best_rep["name"] = "BRK02"
print()
print(al.fmt(best_rep))
print("JSON:", al.as_json(best_rep))