big-qmt/py-client/strategy/zt/state.py

"""做 T 策略的持仓状态和逐笔实际成交记录。"""

import math
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

from sdk import DealItem, PositionItem
from libs.orderbook import OrderBook

READY, SOLD, DONE = "READY", "SOLD", "DONE"


@dataclass(slots=True)
class TStateItem:
    code: str
    base_qty: int = 0
    base_cost: float = 0.0
    trade_date: str = ""
    phase: str = READY
    sell_qty: int = 0
    sell_price: float = 0.0
    buy_qty: int = 0
    buy_cost: float = 0.0
    id: int = 0
    base_order_id: str = ''
    added_order_id: str = ''
    added_num: int = 0
    added_qty: int = 0
    added_cost: float = 0.0


class TState:
    """Apply actual executions immediately, atomically with their position changes."""

    def __init__(self, path: str | Path) -> None:
        self._store = OrderBook(path)
        self.path = self._store.path
        self._load()

    @staticmethod
    def _is_zt_deal(deal: DealItem) -> bool:
        return (
            deal.side == 'BUY' and deal.local_order_id.startswith(('zt-base-', 'zt-t-buy-'))
        ) or (
            deal.side == 'SELL' and deal.local_order_id.startswith('zt-t-sell-')
        )

    @staticmethod
    def _reset(item: TStateItem, date: str) -> bool:
        if item.phase == DONE and item.trade_date != date:
            item.phase, item.trade_date = READY, ''
            item.sell_qty = item.buy_qty = 0
            item.sell_price = item.buy_cost = 0.0
            return True
        return False

    @classmethod
    def _apply_t_deal(cls, item: TStateItem, deal: dict) -> None:
        """实时入账与重启恢复共用同一套做 T 轮次计算。"""
        cls._reset(item, deal['trade_date'])
        qty, amount = deal['volume'], deal['trade_amount']
        if str(deal['offset_flag']) in ('24', '49'):
            total = item.sell_qty + qty
            item.sell_price = (item.sell_qty * item.sell_price + amount) / total
            item.sell_qty = total
            item.phase = SOLD
        else:
            total = item.buy_qty + qty
            item.buy_cost = (item.buy_qty * item.buy_cost + amount) / total
            item.buy_qty = total
            item.phase = DONE if total >= item.sell_qty else SOLD
        item.trade_date = deal['trade_date']

    def reconcile(
        self, positions: list[PositionItem], deals: list[DealItem]
    ) -> None:
        """Deduplicate each fill; partial fills do not wait for order completion."""
        today = datetime.now().date().isoformat()
        seen = {row['order_sys_id'] for row in self.deals}
        rows = []
        for deal in deals:
            if not self._is_zt_deal(deal) or deal.order_sys_id in seen:
                continue
            try:
                row = self._store.deal_record(deal)
            except ValueError:
                continue
            rows.append(row)
            seen.add(deal.order_sys_id)
        rows.sort(key=lambda r: (r['trade_date'], r['trade_time']))
        modified = False
        try:
            # Snapshot includes these fills: subtract their net quantity before replay.
            net = {}
            for row in rows:
                net[row['stock_code']] = net.get(row['stock_code'], 0) + (
                    row['volume'] if str(row['offset_flag']) in ('23', '48') else -row['volume']
                )
            for position in positions:
                code = position.stock_code
                if code in self.items or position.volume <= 0:
                    continue
                if not math.isfinite(position.open_price) or position.open_price <= 0:
                    continue
                qty = max(0, position.volume - net.get(code, 0))
                self.items[code] = TStateItem(code, qty, position.open_price if qty else 0.0)
                modified = True

            # 全部卖出时快照可能已无该证券，按净卖出数量恢复待买回的底仓数量。
            for code, delta in net.items():
                if code not in self.items and delta < 0:
                    self.items[code] = TStateItem(code, -delta)

            for row in rows:
                code = row['stock_code']
                item = self.items.get(code)
                if item is None:
                    item = self.items[code] = TStateItem(code)
                self._reset(item, row['trade_date'])
                qty, amount = row['volume'], row['trade_amount']
                if row['order_local_id'].startswith('zt-base-'):
                    total = item.base_qty + qty
                    item.base_cost = (item.base_qty * item.base_cost + amount) / total
                    item.base_qty = total
                    item.base_order_id = row['order_local_id']
                else:
                    self._apply_t_deal(item, row)
                self.deals.append(row)
                modified = True

            for item in self.items.values():
                modified = self._reset(item, today) or modified
            if modified:
                self.save()
        except Exception:
            self._load()
            raise

    def save(self) -> None:
        try:
            self._store.save(
                {
                    code: {
                        'stock_code': item.code,
                        'volume': item.base_qty,
                        'open_price': item.base_cost,
                        'open_cost': item.base_qty * item.base_cost,
                    }
                    for code, item in self.items.items()
                },
                self.deals,
            )
        except Exception:
            self._load()
            raise

    def _load(self) -> None:
        self._store.load()
        self.items = {}
        for code, position in self._store.positions.items():
            self.items[code] = TStateItem(code, position['volume'], position['open_price'], id=position['id'])
        self.deals = [
            {key: value for key, value in deal.items() if key != 'id'}
            for deal in self._store.deals.values()
        ]
        # 轮次明细不占用持仓表字段，从已保存的逐笔成交重建。
        for deal in self.deals:
            item = self.items.get(deal['stock_code'])
            if item is None:
                continue
            local_order_id = deal['order_local_id']
            if local_order_id.startswith('zt-base-'):
                item.base_order_id = local_order_id
            else:
                self._apply_t_deal(item, deal)
        today = datetime.now().date().isoformat()
        for item in self.items.values():
            self._reset(item, today)