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

"""做 T 策略的底仓、待确认委托和实际成交记录。"""

from __future__ import annotations

import json
import logging as log
import math
from dataclasses import asdict, dataclass, field
from pathlib import Path
from time import time

from sdk import DealItem, PositionItem

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_order_id: str = ""
    sell_qty: int = 0
    sell_price: float = 0.0
    buy_order_id: str = ""
    base_order_id: str = ""
    buy_qty: int = 0
    buy_cost: float = 0.0


@dataclass(slots=True)
class PendingOrder:
    order_id: str
    code: str
    kind: str  # base：底仓；sell：做 T 卖出；buy：做 T 买回
    qty: int
    trade_date: str
    submit_at: float = field(default_factory=time)


class TState:
    """交易逻辑串行更新；JSON 保存底仓、待确认委托及成交历史。"""

    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.items: dict[str, TStateItem] = {}
        self.pending: dict[str, PendingOrder] = {}
        self.records: list[dict] = []
        self._load()

    @classmethod
    def for_strategy(
        cls, data_dir: str | Path, strategy: str, account_id: str
    ) -> TState:
        return cls(Path(data_dir) / f"{strategy}_{account_id}_state.json")

    def busy(self, code: str) -> bool:
        return any(order.code == code for order in self.pending.values())

    def new_order(self, order: PendingOrder) -> None:
        """下单前落盘；请求超时不能当作失败删除，等待后续委托确认。"""
        if self.busy(order.code):
            raise ValueError(f"{order.code} 已有待确认委托")
        self.pending[order.order_id] = order
        try:
            self.save()
        except Exception:
            del self.pending[order.order_id]
            raise

    def reconcile(
        self, positions: list[PositionItem], deals: list[DealItem], today: str
    ) -> None:
        """先按实际成交记账，再接管未知持仓；不覆盖已记录的底仓成本。"""
        # 同一本地委托可能有多笔成交；按成交编号去重后合并数量和金额。
        by_id: dict[str, dict[str, DealItem]] = {}
        for deal in deals:
            if deal.local_order_id and deal.id:
                by_id.setdefault(deal.local_order_id, {})[deal.id] = deal
        for order_id, pending in list(self.pending.items()):
            side = "SELL" if pending.kind == "sell" else "BUY"
            rows = [
                row
                for row in by_id.get(order_id, {}).values()
                if row.code == pending.code and row.side == side
            ]
            if not rows:
                log.warning("[ZT 状态] 成交暂未查到，保留待确认：%s", order_id)
                continue
            qty = sum(row.volume for row in rows)
            # 成交未达到计划数量时继续等待，防止后续成交到达后重复记账。
            if qty != pending.qty:
                continue
            amounts = [
                row.amount if row.amount > 0 else row.price * row.volume
                for row in rows
                if row.volume > 0
            ]
            if any(not math.isfinite(amount) or amount <= 0 for amount in amounts):
                continue
            amount = sum(amounts)
            cost = amount / qty if qty else 0.0
            item = self.items.setdefault(pending.code, TStateItem(pending.code))
            if pending.kind == "base":
                item.base_order_id = order_id
                item.base_qty, item.base_cost = qty, cost
            elif pending.kind == "sell":
                item.trade_date = today  # 跨日成交也占用确认当天的一轮。
                item.sell_order_id = order_id
                item.sell_qty, item.sell_price = qty, cost
                item.buy_qty, item.buy_cost = 0, 0.0
                item.phase = SOLD if qty else READY
            else:
                total = item.buy_qty + qty
                item.buy_cost = (
                    (item.buy_qty * item.buy_cost + amount) / total if total else 0.0
                )
                item.buy_qty = total
                item.buy_order_id = order_id
                item.phase = DONE if total >= item.sell_qty else SOLD
                if item.phase == DONE:
                    item.trade_date = today
            # 记录真实成交编号，重启后仍可核对本次状态变更的来源。
            self.records.append(
                {
                    **asdict(pending),
                    "confirmed_date": today,
                    "filled_qty": qty,
                    "filled_cost": cost,
                    "amount": amount,
                    "deal_ids": [row.id for row in rows],
                }
            )
            del self.pending[order_id]

        for position in positions:
            code = position.stock_code
            if position.volume <= 0 or self.busy(code):
                continue
            if (
                code not in self.items
                and math.isfinite(position.open_price)
                and position.open_price > 0
            ):
                self.items[code] = TStateItem(
                    code, position.volume, position.open_price
                )
                self.records.append(
                    {
                        "kind": "import",
                        "code": code,
                        "date": today,
                        "filled_qty": position.volume,
                        "filled_cost": position.open_price,
                    }
                )
                log.warning(
                    "[ZT 底仓] 首次接管 %s，使用当前均价，无法还原历史成本", code
                )

        for item in self.items.values():
            # 未买回的轮次跨日继续，不删除零持仓的做 T 债务。
            if (
                item.trade_date != today
                and item.phase == DONE
                and not self.busy(item.code)
            ):
                item.phase, item.trade_date = READY, ""
                item.sell_qty = item.buy_qty = 0
                item.sell_price = item.buy_cost = 0.0
                item.sell_order_id = item.buy_order_id = ""
        self.save()

    def save(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        temporary = self.path.with_suffix(self.path.suffix + ".tmp")
        payload = {
            "items": {key: asdict(item) for key, item in self.items.items()},
            "pending": {key: asdict(item) for key, item in self.pending.items()},
            "records": self.records,
        }
        temporary.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + "\n",
            encoding="utf-8",
        )
        temporary.replace(self.path)

    def _load(self) -> None:
        if not self.path.is_file():
            return
        raw = json.loads(self.path.read_text(encoding="utf-8"))
        self.items = {
            code: TStateItem(**item) for code, item in raw["items"].items()
        }
        self.pending = {
            key: PendingOrder(**item) for key, item in raw["pending"].items()
        }
        self.records = raw["records"]