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 OrderItem, PositionItem

READY, SOLD, DONE = "READY", "SOLD", "DONE"
TERMINAL_STATUSES = {"53", "54", "56", "57"}


@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], orders: list[OrderItem], today: str
    ) -> None:
        """先按实际成交记账，再接管未知持仓；不覆盖已记录的底仓成本。"""
        # 同一本地委托可能拆单；按券商订单号去重，数量齐全且全部结束才记账。
        by_id: dict[str, dict[str, OrderItem]] = {}
        for order in orders:
            if order.local_order_id:
                by_id.setdefault(order.local_order_id, {})[order.id] = order
        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
            if sum(row.volume for row in rows) != pending.qty:
                continue
            if any(row.status not in TERMINAL_STATUSES for row in rows):
                continue
            if any(
                row.status == "56" and row.traded_volume != row.volume for row in rows
            ):
                continue
            qty = sum(row.traded_volume for row in rows)
            amounts = [
                (
                    row.trade_amount
                    if row.trade_amount > 0
                    else row.trade_price * row.traded_volume
                )
                for row in rows
                if row.traded_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,
                    "statuses": [row.status 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"))
        # 兼容原 ZT 文件，保留原底仓成本；没有额外版本字段。
        self.items = {
            code: TStateItem(**item) for code, item in raw.get("items", raw).items()
        }
        self.pending = {
            key: PendingOrder(**item) for key, item in raw.get("pending", {}).items()
        }
        self.records = raw.get("records", [])
        if "items" not in raw:
            # 旧记录只有提交行情价，不把它伪装成真实成交历史。
            self.records.append({"kind": "legacy_import", "items": raw})
            for item in self.items.values():
                if item.phase in {"SELLING", "BUYING", SOLD}:
                    raise ValueError(
                        f"[ZT 状态] {item.code} 旧做 T 轮次未结束，需先核对成交再迁移"
                    )