115 lines
3.0 KiB
Python

from __future__ import annotations
from dataclasses import dataclass, field
from typing import List, Dict, Any
from .events import BookEvent
from .orders import MarketBuy, MarketSell
from .fees import calculate_okx_fee
@dataclass
class PortfolioState:
cash_usd: float
coin: float = 0.0
entry_price: float | None = None
entry_time_ms: int | None = None
fees_paid: float = 0.0
equity_curve: list[tuple[int, float]] = field(default_factory=list)
def process_event(
state: PortfolioState,
event: BookEvent,
orders: list,
stoploss_pct: float,
is_maker: bool,
) -> list[Dict[str, Any]]:
"""Apply orders at top-of-book and update stop-loss logic.
Returns log entries for any executions that occurred.
"""
logs: list[Dict[str, Any]] = []
# Evaluate stop-loss for an open long position
if state.coin > 0 and state.entry_price is not None:
threshold = state.entry_price * (1.0 - stoploss_pct)
if event.best_bid <= threshold:
# Exit at the worse of best_bid and threshold to emulate slippage
exit_price = min(event.best_bid, threshold)
gross = state.coin * exit_price
fee = calculate_okx_fee(gross, is_maker=is_maker)
state.cash_usd = gross - fee
state.fees_paid += fee
pnl = (exit_price - state.entry_price) / state.entry_price
logs.append({
"type": "stop_loss",
"time": event.timestamp_ms,
"price": exit_price,
"usd": state.cash_usd,
"coin": 0.0,
"pnl": pnl,
"fee": fee,
})
state.coin = 0.0
state.entry_price = None
state.entry_time_ms = None
# Apply incoming orders
for order in orders:
if isinstance(order, MarketBuy):
if state.cash_usd <= 0:
continue
# Execute at best ask
price = event.best_ask
notional = min(order.usd_notional, state.cash_usd)
qty = notional / price if price > 0 else 0.0
fee = calculate_okx_fee(notional, is_maker=is_maker)
# Deduct both notional and fee from cash
state.cash_usd -= (notional + fee)
state.fees_paid += fee
state.coin += qty
state.entry_price = price
state.entry_time_ms = event.timestamp_ms
logs.append({
"type": "buy",
"time": event.timestamp_ms,
"price": price,
"usd": state.cash_usd,
"coin": state.coin,
"fee": fee,
})
elif isinstance(order, MarketSell):
if state.coin <= 0:
continue
price = event.best_bid
qty = min(order.amount, state.coin)
gross = qty * price
fee = calculate_okx_fee(gross, is_maker=is_maker)
state.cash_usd += (gross - fee)
state.fees_paid += fee
state.coin -= qty
pnl = 0.0
if state.entry_price:
pnl = (price - state.entry_price) / state.entry_price
logs.append({
"type": "sell",
"time": event.timestamp_ms,
"price": price,
"usd": state.cash_usd,
"coin": state.coin,
"pnl": pnl,
"fee": fee,
})
if state.coin <= 0:
state.entry_price = None
state.entry_time_ms = None
# Track equity at each event using best bid for mark-to-market
mark_price = event.best_bid
equity = state.cash_usd + state.coin * mark_price
state.equity_curve.append((event.timestamp_ms, equity))
return logs