okx-codex-trader/scripts/stress_short_bias_swing.py

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass
from pathlib import Path

import pandas as pd

sys.path.append(str(Path(__file__).resolve().parent))

import search_short_bias_swing as base


PREFIX = "swing-stress"
FEES = (0.0004, 0.0008, 0.0010)
SLIPPAGES = (0.0, 0.0005, 0.0010)
FUNDING_8H = (-0.00005, 0.0, 0.00005)
FOCUS_FAMILY = "vol_expansion_short"
FOCUS_BAR = "4H"


@dataclass(frozen=True)
class Stress:
    fee_single_side: float
    slippage: float
    funding_8h: float

    @property
    def label(self) -> str:
        return f"fee{self.fee_single_side:.4f}-slip{self.slippage:.4f}-funding{self.funding_8h:.5f}"


def funding_events(entry_time: pd.Timestamp, exit_time: pd.Timestamp) -> int:
    hours = (exit_time - entry_time).total_seconds() / 3600
    return int(hours // 8)


def close_short(entry_price: float, exit_price: float, entry_time: pd.Timestamp, exit_time: pd.Timestamp, equity: float, stress: Stress) -> tuple[float, float, float]:
    funding_return = funding_events(entry_time, exit_time) * stress.funding_8h
    net_return = entry_price / exit_price - 1.0 - stress.fee_single_side * 2 + funding_return
    return max(0.0, equity * (1.0 + net_return)), net_return, funding_return


def run_strategy(strategy: base.Strategy, frame: pd.DataFrame, btc_frame: pd.DataFrame | None, stress: Stress) -> dict[str, object]:
    signals = base.signal_frame(strategy, frame, btc_frame)
    warmup = max(int(strategy.params.get(key, 0)) for key in ("slow", "entry", "exit", "atr", "vol_window", "btc_slow")) + 2
    equity = base.INITIAL_EQUITY
    position: dict[str, float | int | pd.Timestamp] | None = None
    pending_entry = False
    pending_exit = False
    trades: list[dict[str, object]] = []
    equity_curve: list[dict[str, object]] = []

    rows = list(frame.itertuples())
    for index in range(warmup, len(rows)):
        row = rows[index]
        ts = frame.index[index]
        if pending_exit and position is not None:
            exit_price = float(row.open) * (1.0 + stress.slippage)
            equity, net_return, funding_return = close_short(
                float(position["entry_price"]),
                exit_price,
                pd.Timestamp(position["entry_time"]),
                ts,
                equity,
                stress,
            )
            trades.append(
                {
                    "side": "Short",
                    "entry_time": pd.Timestamp(position["entry_time"]).strftime("%Y-%m-%d %H:%M"),
                    "exit_time": ts.strftime("%Y-%m-%d %H:%M"),
                    "entry_price": float(position["entry_price"]),
                    "exit_price": exit_price,
                    "return": net_return,
                    "funding_return": funding_return,
                    "hold_bars": index - int(position["entry_index"]),
                }
            )
            position = None
            pending_exit = False

        if pending_entry and position is None and equity > 0.0:
            atr = float(signals["atr"].iloc[index - 1])
            entry_price = float(row.open) * (1.0 - stress.slippage)
            position = {
                "entry_time": ts,
                "entry_index": index,
                "entry_price": entry_price,
                "stop_price": entry_price + atr * float(strategy.params["stop_atr"]),
                "take_price": max(0.01, entry_price - atr * float(strategy.params["take_atr"])),
            }
            pending_entry = False

        mark_equity = equity
        if position is not None:
            stop_hit = float(row.high) >= float(position["stop_price"])
            take_hit = float(row.low) <= float(position["take_price"])
            if stop_hit or take_hit:
                raw_exit = float(position["stop_price"] if stop_hit else position["take_price"])
                exit_price = raw_exit * (1.0 + stress.slippage)
                equity, net_return, funding_return = close_short(
                    float(position["entry_price"]),
                    exit_price,
                    pd.Timestamp(position["entry_time"]),
                    ts,
                    equity,
                    stress,
                )
                trades.append(
                    {
                        "side": "Short",
                        "entry_time": pd.Timestamp(position["entry_time"]).strftime("%Y-%m-%d %H:%M"),
                        "exit_time": ts.strftime("%Y-%m-%d %H:%M"),
                        "entry_price": float(position["entry_price"]),
                        "exit_price": exit_price,
                        "return": net_return,
                        "funding_return": funding_return,
                        "hold_bars": index - int(position["entry_index"]),
                    }
                )
                position = None
                mark_equity = equity

        if position is not None:
            funding_return = funding_events(pd.Timestamp(position["entry_time"]), ts) * stress.funding_8h
            gross_return = float(position["entry_price"]) / float(row.close) - 1.0
            mark_equity = max(0.0, equity * (1.0 + gross_return - stress.fee_single_side + funding_return))

        equity_curve.append({"ts": ts, "equity": mark_equity})
        if index == len(rows) - 1 or equity <= 0.0:
            continue

        if position is not None:
            held = index - int(position["entry_index"])
            if bool(signals["exit"].iloc[index]) or held >= int(strategy.params["max_hold"]):
                pending_exit = True
        elif bool(signals["entry"].iloc[index]):
            pending_entry = True

    if position is not None:
        last = rows[-1]
        ts = frame.index[-1]
        exit_price = float(last.close) * (1.0 + stress.slippage)
        equity, net_return, funding_return = close_short(
            float(position["entry_price"]),
            exit_price,
            pd.Timestamp(position["entry_time"]),
            ts,
            equity,
            stress,
        )
        trades.append(
            {
                "side": "Short",
                "entry_time": pd.Timestamp(position["entry_time"]).strftime("%Y-%m-%d %H:%M"),
                "exit_time": ts.strftime("%Y-%m-%d %H:%M"),
                "entry_price": float(position["entry_price"]),
                "exit_price": exit_price,
                "return": net_return,
                "funding_return": funding_return,
                "hold_bars": len(rows) - 1 - int(position["entry_index"]),
            }
        )
        equity_curve.append({"ts": ts, "equity": equity})

    return {"trades": trades, "equity_curve": equity_curve}


def load_candidate_names(path: Path) -> set[str]:
    candidates = pd.read_csv(path)
    if candidates.empty:
        raise ValueError(f"no candidates in {path}")
    return set(candidates["name"].astype(str))


def format_cell(value: object) -> str:
    if isinstance(value, float):
        return f"{value:.6g}"
    return str(value).replace("|", "\\|")


def markdown_table(frame: pd.DataFrame) -> str:
    rows = [list(frame.columns), ["---" for _ in frame.columns]]
    rows.extend(frame.astype(object).where(pd.notna(frame), "").values.tolist())
    return "\n".join("| " + " | ".join(format_cell(value) for value in row) + " |" for row in rows)


def focus_decision(focus: pd.DataFrame) -> str:
    if focus.empty:
        return "No. BTC/ETH 4H vol_expansion_short was not present in the stressed candidate set."
    worst = focus.sort_values("worst_stress_total_return").iloc[0]
    both_positive = bool((focus["worst_stress_total_return"] > 0.0).all() and (focus["worst_stress_return_1y"] > 0.0).all())
    if both_positive:
        return (
            "Yes, still worth continuing as a research candidate. "
            f"The weakest focused variant is {worst['name']} with worst stressed total return {worst['worst_stress_total_return']:.2%} "
            f"and worst stressed 1y return {worst['worst_stress_return_1y']:.2%}."
        )
    return (
        "No for live deployment; keep only as a research candidate. "
        f"The weakest focused variant is {worst['name']} with worst stressed total return {worst['worst_stress_total_return']:.2%} "
        f"and worst stressed 1y return {worst['worst_stress_return_1y']:.2%}."
    )


def markdown_report(command: str, paths: list[Path], totals: pd.DataFrame, focus: pd.DataFrame, summary: dict[str, object]) -> str:
    worst = totals.sort_values(["total_return", "return_1y"]).head(10)
    focus_view = focus.sort_values(["symbol", "fast", "slow", "stop_atr", "take_atr"])
    lines = [
        "# Short-Bias Swing Stress Test",
        "",
        f"Run command: `{command}`",
        "",
        "Output files:",
        *[f"- `{path}`" for path in paths],
        "",
        "Scope: strategies from `reports/short-bias/swing-qualified.csv`; original search script was imported but not modified.",
        "Stress grid: fee 0.04/0.08/0.10% single-side, slippage 0/0.05/0.10%, funding every 8h -0.005/0/+0.005%. Positive funding is short receive; negative funding is short pay.",
        "",
        f"Conclusion: {summary['decision']}",
        "",
        "## BTC/ETH 4H vol_expansion_short robustness",
        "",
        markdown_table(
            focus_view[
                [
                    "name",
                    "symbol",
                    "stress_cases",
                    "positive_cases",
                    "worst_stress_total_return",
                    "worst_stress_max_drawdown",
                    "worst_stress_calmar",
                    "worst_stress_return_3y",
                    "worst_stress_return_1y",
                    "worst_stress_return_6m",
                    "worst_stress_return_3m",
                    "worst_stress_fee_single_side",
                    "worst_stress_slippage",
                    "worst_stress_funding_8h",
                    "trades",
                ]
            ]
        ),
        "",
        "## Worst stressed cases",
        "",
        markdown_table(
            worst[
                [
                    "name",
                    "symbol",
                    "bar",
                    "family",
                    "fee_single_side",
                    "slippage",
                    "funding_8h",
                    "total_return",
                    "annualized_return",
                    "max_drawdown",
                    "calmar",
                    "trades",
                    "profit_factor",
                    "return_3y",
                    "return_1y",
                    "return_6m",
                    "return_3m",
                    "total_funding_return",
                ]
            ]
        ),
    ]
    return "\n".join(lines) + "\n"


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--years", type=float, default=base.YEARS)
    parser.add_argument("--candidates", type=Path, default=base.OUTPUT_DIR / "swing-qualified.csv")
    parser.add_argument("--output-dir", type=Path, default=base.OUTPUT_DIR)
    args = parser.parse_args()

    candidate_names = load_candidate_names(args.candidates)
    strategies = [strategy for strategy in base.build_strategies() if strategy.name in candidate_names]
    if not strategies:
        raise ValueError("no matching built-in strategies for candidate file")

    raw = {symbol: base.load_15m_frame(symbol, args.years) for symbol in base.SYMBOLS}
    data = {(symbol, bar): base.resample_frame(raw[symbol], bar) for symbol in base.SYMBOLS for bar in base.BARS}
    stresses = [Stress(fee, slippage, funding) for fee in FEES for slippage in SLIPPAGES for funding in FUNDING_8H]

    rows: list[dict[str, object]] = []
    for strategy_index, strategy in enumerate(strategies, start=1):
        btc_frame = data[("BTC-USDT-SWAP", strategy.bar)] if strategy.family == "btc_riskoff_eth" else None
        frame = data[(strategy.symbol, strategy.bar)]
        for stress in stresses:
            result = run_strategy(strategy, frame, btc_frame, stress)
            series = base.daily_equity(result)
            trades = list(result["trades"])
            rows.append(
                {
                    "name": strategy.name,
                    "symbol": strategy.symbol,
                    "bar": strategy.bar,
                    "family": strategy.family,
                    "first_day": series.index[0].strftime("%Y-%m-%d"),
                    "last_day": series.index[-1].strftime("%Y-%m-%d"),
                    "fee_single_side": stress.fee_single_side,
                    "slippage": stress.slippage,
                    "funding_8h": stress.funding_8h,
                    "total_funding_return": sum(float(trade["funding_return"]) for trade in trades),
                    **base.equity_metrics(series),
                    **base.trade_metrics(trades),
                    **base.horizon_returns(series),
                    **strategy.params,
                }
            )
        print(f"done {strategy_index}/{len(strategies)} {strategy.name}", flush=True)

    totals = pd.DataFrame(rows).sort_values(["name", "fee_single_side", "slippage", "funding_8h"])
    grouped = totals.groupby("name", as_index=False)
    robustness = grouped.agg(
        symbol=("symbol", "first"),
        bar=("bar", "first"),
        family=("family", "first"),
        stress_cases=("total_return", "size"),
        positive_cases=("total_return", lambda values: int((values > 0.0).sum())),
        worst_stress_total_return=("total_return", "min"),
        worst_stress_max_drawdown=("max_drawdown", "max"),
        worst_stress_calmar=("calmar", "min"),
        worst_stress_return_3y=("return_3y", "min"),
        worst_stress_return_1y=("return_1y", "min"),
        worst_stress_return_6m=("return_6m", "min"),
        worst_stress_return_3m=("return_3m", "min"),
        trades=("trades", "first"),
    )
    params = totals.drop_duplicates("name")[["name", "fast", "slow", "stop_atr", "take_atr"]]
    worst_cases = totals.loc[totals.groupby("name")["total_return"].idxmin()][["name", "fee_single_side", "slippage", "funding_8h"]]
    robustness = robustness.merge(params, on="name", how="left").merge(worst_cases, on="name", how="left", suffixes=("", "_worst"))
    robustness = robustness.rename(
        columns={
            "fee_single_side": "worst_stress_fee_single_side",
            "slippage": "worst_stress_slippage",
            "funding_8h": "worst_stress_funding_8h",
        }
    ).sort_values(["worst_stress_total_return", "worst_stress_return_1y"], ascending=[False, False])
    focus = robustness[
        (robustness["family"] == FOCUS_FAMILY)
        & (robustness["bar"] == FOCUS_BAR)
        & (robustness["symbol"].isin(["BTC-USDT-SWAP", "ETH-USDT-SWAP"]))
    ].copy()
    decision = focus_decision(focus)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    totals_path = args.output_dir / f"{PREFIX}-totals.csv"
    robustness_path = args.output_dir / f"{PREFIX}-robustness.csv"
    focus_path = args.output_dir / f"{PREFIX}-vol-expansion-4h.csv"
    summary_path = args.output_dir / f"{PREFIX}-summary.json"
    report_path = args.output_dir / f"{PREFIX}-report.md"
    paths = [totals_path, robustness_path, focus_path, summary_path, report_path]

    totals.to_csv(totals_path, index=False)
    robustness.to_csv(robustness_path, index=False)
    focus.to_csv(focus_path, index=False)
    summary: dict[str, object] = {
        "years_requested": args.years,
        "candidate_file": str(args.candidates),
        "strategy_count": len(strategies),
        "stress_case_count": len(stresses),
        "row_count": len(totals),
        "focus_count": len(focus),
        "decision": decision,
        "output_files": [str(path) for path in paths],
    }
    summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    command = f"rtk .venv/bin/python {Path(__file__).as_posix()} --years {args.years}"
    report_path.write_text(markdown_report(command, paths, totals, focus, summary), encoding="utf-8")

    print(focus.to_string(index=False, formatters={col: format_cell for col in focus.columns}))
    print(f"wrote={report_path}")
    return 0


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