okx-codex-trader/okx_codex_trader/donchian_report.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import pandas as pd

from okx_codex_trader.models import Candle
from okx_codex_trader.sampled_report import (
    SegmentResult,
    generate_sampled_report,
    mark_to_market as _mark_to_market,
    trade_equity as _trade_equity,
)


DONCHIAN_STRATEGY_DESCRIPTION = (
    "Donchian channel breakout, previous-window close breakout entries at next open, "
    "opposite-channel exits at next open, intrabar 1.0% stop-loss."
)


@dataclass(frozen=True)
class DonchianConfig:
    entry_window: int = 20
    exit_window: int = 10
    stop_loss_pct: float = 0.01
    initial_equity: float = 10_000.0


def _format_ts(ts: int) -> str:
    return pd.to_datetime(ts, unit="ms", utc=True).strftime("%Y-%m-%d %H:%M")


def run_donchian_segment(
    *,
    candles: list[Candle],
    leverage: int,
    warmup_bars: int,
    config: DonchianConfig = DonchianConfig(),
) -> SegmentResult:
    highs = pd.Series([candle.high for candle in candles], dtype=float)
    lows = pd.Series([candle.low for candle in candles], dtype=float)
    entry_high = highs.shift(1).rolling(config.entry_window).max().tolist()
    entry_low = lows.shift(1).rolling(config.entry_window).min().tolist()
    exit_high = highs.shift(1).rolling(config.exit_window).max().tolist()
    exit_low = lows.shift(1).rolling(config.exit_window).min().tolist()

    equity = config.initial_equity
    ending_equity = equity
    peak_equity = equity
    max_drawdown = 0.0
    wins = 0
    trades: list[dict[str, object]] = []
    entries: list[dict[str, object]] = []
    exits: list[dict[str, object]] = []
    equity_curve: list[dict[str, float | int]] = []
    position: dict[str, object] | None = None
    pending_entry_side: str | None = None
    pending_exit = False

    for index in range(warmup_bars, len(candles)):
        candle = candles[index]

        if pending_exit and position is not None:
            exit_price = candle.open
            exit_equity = _trade_equity(
                side=str(position["side"]),
                margin_used=float(position["margin_used"]),
                entry_price=float(position["entry_price"]),
                exit_price=exit_price,
                leverage=leverage,
            )
            trades.append(
                {
                    "side": "Long" if position["side"] == "long" else "Short",
                    "entry_time": _format_ts(int(position["entry_time"])),
                    "exit_time": _format_ts(candle.ts),
                    "entry_price": round(float(position["entry_price"]), 4),
                    "exit_price": round(exit_price, 4),
                    "pnl": round(exit_equity - float(position["margin_used"]), 4),
                    "return_pct": round(
                        (exit_equity - float(position["margin_used"])) / float(position["margin_used"]) * 100,
                        4,
                    ),
                }
            )
            exits.append({"ts": candle.ts, "price": exit_price, "side": position["side"]})
            if exit_equity > float(position["margin_used"]):
                wins += 1
            equity = exit_equity
            position = None
            pending_exit = False

        if pending_entry_side is not None and position is None:
            entry_price = candle.open
            margin_used = equity
            stop_multiplier = 1 - config.stop_loss_pct if pending_entry_side == "long" else 1 + config.stop_loss_pct
            position = {
                "side": pending_entry_side,
                "entry_time": candle.ts,
                "entry_price": entry_price,
                "entry_index": index,
                "margin_used": margin_used,
                "stop_price": entry_price * stop_multiplier,
            }
            entries.append({"ts": candle.ts, "price": entry_price, "side": pending_entry_side})
            pending_entry_side = None

        current_equity = equity
        if position is not None:
            stop_hit = (
                position["side"] == "long" and candle.low <= float(position["stop_price"])
            ) or (
                position["side"] == "short" and candle.high >= float(position["stop_price"])
            )
            if stop_hit:
                if position["side"] == "long" and candle.open < float(position["stop_price"]):
                    exit_price = candle.open
                elif position["side"] == "short" and candle.open > float(position["stop_price"]):
                    exit_price = candle.open
                else:
                    exit_price = float(position["stop_price"])
                exit_equity = _trade_equity(
                    side=str(position["side"]),
                    margin_used=float(position["margin_used"]),
                    entry_price=float(position["entry_price"]),
                    exit_price=exit_price,
                    leverage=leverage,
                )
                trades.append(
                    {
                        "side": "Long" if position["side"] == "long" else "Short",
                        "entry_time": _format_ts(int(position["entry_time"])),
                        "exit_time": _format_ts(candle.ts),
                        "entry_price": round(float(position["entry_price"]), 4),
                        "exit_price": round(exit_price, 4),
                        "pnl": round(exit_equity - float(position["margin_used"]), 4),
                        "return_pct": round(
                            (exit_equity - float(position["margin_used"])) / float(position["margin_used"]) * 100,
                            4,
                        ),
                    }
                )
                exits.append({"ts": candle.ts, "price": exit_price, "side": position["side"]})
                if exit_equity > float(position["margin_used"]):
                    wins += 1
                equity = exit_equity
                current_equity = exit_equity
                position = None
                if current_equity > peak_equity:
                    peak_equity = current_equity
                max_drawdown = max(max_drawdown, (peak_equity - current_equity) / peak_equity)
                equity_curve.append({"ts": candle.ts, "equity": current_equity, "close": candle.close})
                ending_equity = current_equity
                continue

        if position is not None:
            current_equity = _mark_to_market(
                side=str(position["side"]),
                margin_used=float(position["margin_used"]),
                entry_price=float(position["entry_price"]),
                mark_price=candle.close,
                leverage=leverage,
            )

        if current_equity > peak_equity:
            peak_equity = current_equity
        max_drawdown = max(max_drawdown, (peak_equity - current_equity) / peak_equity)
        equity_curve.append({"ts": candle.ts, "equity": current_equity, "close": candle.close})
        ending_equity = current_equity

        if index == len(candles) - 1:
            continue

        if position is not None:
            exit_signal = (
                position["side"] == "long" and exit_low[index] == exit_low[index] and candle.close < float(exit_low[index])
            ) or (
                position["side"] == "short" and exit_high[index] == exit_high[index] and candle.close > float(exit_high[index])
            )
            if exit_signal:
                pending_exit = True
            continue

        if entry_high[index] == entry_high[index] and candle.close > float(entry_high[index]):
            pending_entry_side = "long"
        elif entry_low[index] == entry_low[index] and candle.close < float(entry_low[index]):
            pending_entry_side = "short"

    trade_count = len(trades)
    return SegmentResult(
        trade_count=trade_count,
        total_return=(ending_equity - config.initial_equity) / config.initial_equity,
        win_rate=(wins / trade_count) if trade_count else 0.0,
        max_drawdown=max_drawdown,
        trades=trades,
        open_position=position,
        candles=candles[warmup_bars:],
        equity_curve=equity_curve,
        entries=entries,
        exits=exits,
    )


def generate_donchian_sampled_report(
    *,
    candles: list[Candle],
    leverage: int,
    output_file: Path,
    symbol: str,
    bar: str,
    segments: int,
    window_size: int,
    entry_window: int = DonchianConfig.entry_window,
    exit_window: int = DonchianConfig.exit_window,
    stop_loss_pct: float = DonchianConfig.stop_loss_pct,
) -> dict[str, object]:
    config = DonchianConfig(
        entry_window=entry_window,
        exit_window=exit_window,
        stop_loss_pct=stop_loss_pct,
    )
    return generate_sampled_report(
        candles=candles,
        leverage=leverage,
        output_file=output_file,
        symbol=symbol,
        bar=bar,
        segments=segments,
        window_size=window_size,
        report_title="Donchian Sampled Report",
        strategy_label="Donchian",
        strategy_description=DONCHIAN_STRATEGY_DESCRIPTION,
        strategy_params={
            "entry_window": config.entry_window,
            "exit_window": config.exit_window,
            "stop_loss_pct": config.stop_loss_pct,
        },
        run_segment=lambda *, candles, leverage, warmup_bars: run_donchian_segment(
            candles=candles,
            leverage=leverage,
            warmup_bars=warmup_bars,
            config=config,
        ),
        warmup_bars=max(config.entry_window, config.exit_window),
    )