okx-codex-trader/scripts/run_bb_squeeze_executor.py

from __future__ import annotations

import argparse
import json
import os
import sys
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path

import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from okx_codex_trader.config import load_config
from okx_codex_trader.live_execution import (
    TargetPosition,
    current_position_from_okx,
    plan_position_delta,
    render_market_order_bodies,
)
from okx_codex_trader.okx_client import OkxClient


ROOT = Path(__file__).resolve().parents[1]
STATE_DIR = ROOT / "var" / "bb-squeeze-executor"
STATE_FILE = "runtime-state.json"
EVENTS_FILE = "events.jsonl"
SYMBOL = "ETH-USDT-SWAP"
BAR = "15m"
LEVERAGE = 3
CANDLES_PATH = ROOT / "data" / "okx-candles" / SYMBOL / f"{BAR}.csv"

BAND_LENGTH = 48
BANDWIDTH_LOOKBACK = 960
BANDWIDTH_QUANTILE = 0.25
STOP_LOSS_PCT = 0.01
ETH_VOL_CAP = 0.006
COOLDOWN_BARS = 24


@dataclass(frozen=True)
class StrategyState:
    last_candle_ts: int | None
    active_side: str | None
    entry_price: float | None
    entry_candle_ts: int | None
    cooldown_until_ts: int | None


EMPTY_STATE = StrategyState(None, None, None, None, None)


def now_iso() -> str:
    return datetime.now(UTC).isoformat(timespec="seconds").replace("+00:00", "Z")


def load_state(path: Path) -> StrategyState:
    if not path.exists():
        return EMPTY_STATE
    payload = json.loads(path.read_text(encoding="utf-8"))
    return StrategyState(
        last_candle_ts=payload["last_candle_ts"],
        active_side=payload["active_side"],
        entry_price=payload["entry_price"],
        entry_candle_ts=payload["entry_candle_ts"],
        cooldown_until_ts=payload["cooldown_until_ts"],
    )


def save_state(path: Path, state: StrategyState) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(asdict(state), indent=2, sort_keys=True) + "\n", encoding="utf-8")


def append_event(state_dir: Path, payload: dict[str, object]) -> None:
    state_dir.mkdir(parents=True, exist_ok=True)
    with (state_dir / EVENTS_FILE).open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(payload, sort_keys=True) + "\n")


def load_frame() -> pd.DataFrame:
    frame = pd.read_csv(CANDLES_PATH)
    frame["time"] = pd.to_datetime(frame["ts"], unit="ms", utc=True)
    return frame.sort_values("ts").drop_duplicates("ts", keep="last").reset_index(drop=True)


def signal_from_frame(frame: pd.DataFrame, state: StrategyState) -> tuple[StrategyState, dict[str, object]]:
    if len(frame) < BANDWIDTH_LOOKBACK + 3:
        raise ValueError("not enough candles")
    close = frame["close"].astype(float)
    middle = close.rolling(BAND_LENGTH).mean()
    stdev = close.rolling(BAND_LENGTH).std(ddof=0)
    upper = middle + (2.0 * stdev)
    lower = middle - (2.0 * stdev)
    bandwidth = (upper - lower) / middle
    threshold = bandwidth.rolling(BANDWIDTH_LOOKBACK).quantile(BANDWIDTH_QUANTILE)
    eth_vol = close.pct_change().rolling(96).std(ddof=0)
    decision_index = len(frame) - 2
    row = frame.iloc[decision_index]
    candle_ts = int(row["ts"])
    candle_time = pd.Timestamp(row["time"]).isoformat().replace("+00:00", "Z")

    indicators = {
        "eth_close": float(row["close"]),
        "middle": float(middle.iloc[decision_index]),
        "upper": float(upper.iloc[decision_index]),
        "lower": float(lower.iloc[decision_index]),
        "bandwidth": float(bandwidth.iloc[decision_index]),
        "bandwidth_threshold": float(threshold.iloc[decision_index]),
        "eth_vol_96": float(eth_vol.iloc[decision_index]),
    }

    if state.last_candle_ts is not None and candle_ts <= state.last_candle_ts:
        return state, {
            "decision_candle_ts": candle_ts,
            "decision_candle_time": candle_time,
            "signal": "state_replay",
            "target_side": state.active_side or "flat",
            "indicators": indicators,
        }

    next_state = StrategyState(candle_ts, state.active_side, state.entry_price, state.entry_candle_ts, state.cooldown_until_ts)
    signal = "hold"
    target_side = state.active_side or "flat"

    if state.active_side is not None:
        entry_price = float(state.entry_price)
        stop = entry_price * (1.0 - STOP_LOSS_PCT if state.active_side == "long" else 1.0 + STOP_LOSS_PCT)
        stop_hit = (state.active_side == "long" and float(row["low"]) <= stop) or (state.active_side == "short" and float(row["high"]) >= stop)
        middle_exit = (state.active_side == "long" and float(row["close"]) < indicators["middle"]) or (
            state.active_side == "short" and float(row["close"]) > indicators["middle"]
        )
        if stop_hit or middle_exit:
            signal = "exit_stop" if stop_hit else "exit_middle"
            target_side = "flat"
            next_state = StrategyState(
                candle_ts,
                None,
                None,
                None,
                candle_ts + (COOLDOWN_BARS * 900_000),
            )
    else:
        cooldown_ok = state.cooldown_until_ts is None or candle_ts >= state.cooldown_until_ts
        compressed = indicators["bandwidth"] <= indicators["bandwidth_threshold"]
        vol_ok = indicators["eth_vol_96"] <= ETH_VOL_CAP
        if cooldown_ok and compressed and vol_ok and float(row["close"]) > indicators["upper"]:
            signal = "entry_long"
            target_side = "long"
            next_state = StrategyState(candle_ts, "long", float(row["close"]), candle_ts, state.cooldown_until_ts)
        elif cooldown_ok and compressed and vol_ok and float(row["close"]) < indicators["lower"]:
            signal = "entry_short"
            target_side = "short"
            next_state = StrategyState(candle_ts, "short", float(row["close"]), candle_ts, state.cooldown_until_ts)

    return next_state, {
        "decision_candle_ts": candle_ts,
        "decision_candle_time": candle_time,
        "signal": signal,
        "target_side": target_side,
        "indicators": indicators,
        "params": {
            "band_length": BAND_LENGTH,
            "bandwidth_lookback": BANDWIDTH_LOOKBACK,
            "bandwidth_quantile": BANDWIDTH_QUANTILE,
            "stop_loss_pct": STOP_LOSS_PCT,
            "eth_vol_cap": ETH_VOL_CAP,
            "cooldown_bars": COOLDOWN_BARS,
            "side_mode": "both",
        },
    }


def account_current_position(client: OkxClient, margin_per_unit_usdt: float) -> tuple[TargetPosition, dict[str, object]]:
    positions = client.get_positions(SYMBOL)
    metadata = client.get_instrument_meta(SYMBOL)
    mark_price = client.get_last_price(SYMBOL)
    current = current_position_from_okx(
        positions=positions,
        mark_price=mark_price,
        metadata=metadata,
        leverage=LEVERAGE,
        margin_per_unit_usdt=margin_per_unit_usdt,
    )
    return current, {
        "positions": [asdict(position) for position in positions],
        "instrument_meta": asdict(metadata),
        "mark_price": mark_price,
    }


def target_position(signal: dict[str, object], current: TargetPosition) -> TargetPosition:
    side = str(signal["target_side"])
    if side == "flat":
        return TargetPosition(side="flat", unit=0.0, known=True, reason=str(signal["signal"]))
    if current.known and current.side == side and current.unit > 0.0:
        return TargetPosition(side=side, unit=current.unit, known=True, reason="keep existing same-side position")
    return TargetPosition(side=side, unit=1.0, known=True, reason=str(signal["signal"]))


def risk_arg(value: float | None, env_name: str) -> float:
    if value is not None:
        return value
    raw = os.environ.get(env_name)
    if raw is None or raw == "":
        raise ValueError(f"{env_name} is required")
    return float(raw)


def run_once(
    *,
    state_dir: Path,
    margin_per_unit_usdt: float,
    max_new_margin_usdt: float,
    max_total_margin_usdt: float,
    submit_live: bool,
) -> dict[str, object]:
    state_dir.mkdir(parents=True, exist_ok=True)
    state_path = state_dir / STATE_FILE
    previous_state = load_state(state_path)
    next_state, signal = signal_from_frame(load_frame(), previous_state)
    client = OkxClient(load_config())
    current, account = account_current_position(client, margin_per_unit_usdt)
    target = target_position(signal, current)
    plan = plan_position_delta(current, target)
    orders = ()
    if current.known and target.known:
        from okx_codex_trader.models import InstrumentMeta

        meta = account["instrument_meta"]
        orders = render_market_order_bodies(
            plan=plan,
            symbol=SYMBOL,
            mark_price=float(account["mark_price"]),
            metadata=InstrumentMeta(ct_val=float(meta["ct_val"]), lot_sz=float(meta["lot_sz"]), min_sz=float(meta["min_sz"])),
            leverage=LEVERAGE,
            margin_per_unit_usdt=margin_per_unit_usdt,
            max_new_margin_usdt=max_new_margin_usdt,
            max_total_margin_usdt=max_total_margin_usdt,
            client_order_id_prefix=f"bbsq-{signal['decision_candle_ts']}",
        )
    snapshot = {
        "created_at": now_iso(),
        "mode": "bb_squeeze_live_executor" if submit_live else "bb_squeeze_dry_run_executor",
        "orders_submitted": 0,
        "strategy": "bb-squeeze-l48-bw960-q0.25-sl0.01-tpnone-both-none-vc0.006-ddnone-cd24",
        "previous_state": asdict(previous_state),
        "next_state": asdict(next_state),
        "signal": signal,
        "current_position": asdict(current),
        "target_position": asdict(target),
        "execution_plan": {
            "current": asdict(plan.current),
            "target": asdict(plan.target),
            "actions": [asdict(action) for action in plan.actions],
        },
        "rendered_orders": [asdict(order) for order in orders],
        "account": account,
        "risk_limits": {
            "submit_enabled": submit_live,
            "margin_per_unit_usdt": margin_per_unit_usdt,
            "max_new_margin_usdt": max_new_margin_usdt,
            "max_total_margin_usdt": max_total_margin_usdt,
        },
    }
    if submit_live:
        client.ensure_hedge_mode()
        submitted = []
        try:
            for rendered in orders:
                body = rendered.body
                client.set_leverage(symbol=SYMBOL, leverage=LEVERAGE, pos_side=body["posSide"])
                submitted.append(asdict(client.submit_market_order_body(body)))
        except ValueError as exc:
            snapshot["orders_submitted"] = len(submitted)
            snapshot["submitted_orders"] = submitted
            snapshot["execution_error"] = str(exc)
            append_event(state_dir, snapshot)
            raise
        snapshot["orders_submitted"] = len(submitted)
        snapshot["submitted_orders"] = submitted
    save_state(state_path, next_state)
    append_event(state_dir, snapshot)
    return snapshot


def main() -> int:
    parser = argparse.ArgumentParser(description="Run BB squeeze live executor.")
    parser.add_argument("--state-dir", type=Path, default=STATE_DIR)
    parser.add_argument("--margin-per-unit-usdt", type=float)
    parser.add_argument("--max-new-margin-usdt", type=float)
    parser.add_argument("--max-total-margin-usdt", type=float)
    parser.add_argument("--submit-live", action="store_true")
    parser.add_argument("--confirm-live", action="store_true")
    args = parser.parse_args()
    if args.submit_live != args.confirm_live:
        raise ValueError("--submit-live and --confirm-live must be used together")
    snapshot = run_once(
        state_dir=args.state_dir,
        margin_per_unit_usdt=risk_arg(args.margin_per_unit_usdt, "ETH_NEXTGEN_MARGIN_PER_UNIT_USDT"),
        max_new_margin_usdt=risk_arg(args.max_new_margin_usdt, "ETH_NEXTGEN_MAX_NEW_MARGIN_USDT"),
        max_total_margin_usdt=risk_arg(args.max_total_margin_usdt, "ETH_NEXTGEN_MAX_TOTAL_MARGIN_USDT"),
        submit_live=args.submit_live,
    )
    print(json.dumps(snapshot, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())