okx-codex-trader/okx_codex_trader/backtest.py

from okx_codex_trader.models import BacktestResult, BacktestTrade, Candle
from okx_codex_trader.strategy import simple_moving_average


def run_backtest(candles: list[Candle], leverage: int) -> BacktestResult:
    if leverage is True or leverage is False or not isinstance(leverage, int) or not 1 <= leverage <= 3:
        raise ValueError("leverage is invalid")

    fast = simple_moving_average(candles, 10)
    slow = simple_moving_average(candles, 20)

    initial_equity = 10_000.0
    equity = initial_equity
    trades: list[BacktestTrade] = []
    wins = 0
    peak_equity = initial_equity
    max_drawdown = 0.0
    position: dict[str, float | str] | None = None

    for index in range(1, len(candles) - 1):
        if position is not None:
            entry_price = float(position["entry_price"])
            margin_used = float(position["margin_used"])
            if position["direction"] == "long":
                price_return = (candles[index].close - entry_price) / entry_price
            else:
                price_return = (entry_price - candles[index].close) / entry_price
            marked_equity = margin_used + (margin_used * leverage * price_return)
            if marked_equity > peak_equity:
                peak_equity = marked_equity
            drawdown = (peak_equity - marked_equity) / peak_equity
            if drawdown > max_drawdown:
                max_drawdown = drawdown

        if fast[index] is None or slow[index] is None:
            continue

        signal: str | None = None
        if fast[index - 1] is not None and slow[index - 1] is not None:
            if fast[index - 1] <= slow[index - 1] and fast[index] > slow[index]:
                signal = "long"
            elif fast[index - 1] >= slow[index - 1] and fast[index] < slow[index]:
                signal = "short"

        if signal is None:
            continue

        execution_price = candles[index + 1].open

        if position is not None and position["direction"] != signal:
            entry_price = float(position["entry_price"])
            margin_used = float(position["margin_used"])
            if position["direction"] == "long":
                price_return = (execution_price - entry_price) / entry_price
            else:
                price_return = (entry_price - execution_price) / entry_price

            ending_equity = margin_used + (margin_used * leverage * price_return)
            trades.append(
                BacktestTrade(
                    direction=str(position["direction"]),
                    entry_price=entry_price,
                    exit_price=execution_price,
                    margin_used=margin_used,
                    ending_equity=ending_equity,
                )
            )
            equity = ending_equity
            if ending_equity > float(position["margin_used"]):
                wins += 1
            if equity > peak_equity:
                peak_equity = equity
            drawdown = (peak_equity - equity) / peak_equity
            if drawdown > max_drawdown:
                max_drawdown = drawdown
            position = None

        if position is None:
            position = {
                "direction": signal,
                "entry_price": execution_price,
                "margin_used": equity,
            }

    trade_count = len(trades)
    win_rate = wins / trade_count if trade_count else 0.0
    ending_equity = equity
    if position is not None:
        entry_price = float(position["entry_price"])
        margin_used = float(position["margin_used"])
        if position["direction"] == "long":
            price_return = (candles[-1].close - entry_price) / entry_price
        else:
            price_return = (entry_price - candles[-1].close) / entry_price
        ending_equity = margin_used + (margin_used * leverage * price_return)
        if ending_equity > peak_equity:
            peak_equity = ending_equity
        drawdown = (peak_equity - ending_equity) / peak_equity
        if drawdown > max_drawdown:
            max_drawdown = drawdown

    return BacktestResult(
        initial_equity=initial_equity,
        ending_equity=ending_equity,
        total_return=(ending_equity - initial_equity) / initial_equity,
        max_drawdown=max_drawdown,
        win_rate=win_rate,
        trade_count=trade_count,
        trades=trades,
    )