src/microstructure/collector.py

"""Kalshi order-book tick collector.

Connects to the Kalshi market-data websocket, subscribes to the
``orderbook_delta`` channel for the most liquid currently-tradeable markets,
reconstructs each book, and appends a microstructure feature row on every
update to date-partitioned Parquet files under ``data/ticks/``.

This is the data-collection stage of the microstructure study: level-2 tick
history is not downloadable, so it has to be accumulated live. Run it on an
always-on machine — every row written is a future training example.

    python -m src.microstructure.collector

Paths are relative to the current directory; run it from the repo root.
"""

import asyncio
import json
import logging
import time
from datetime import datetime, timezone
from pathlib import Path

import aiohttp
import pandas as pd
import websockets

from src.microstructure.orderbook import KalshiOrderBook

logger = logging.getLogger(__name__)

KALSHI_REST = "https://api.elections.kalshi.com/trade-api/v2"
KALSHI_WS = "wss://api.elections.kalshi.com/trade-api/ws/v2"
DATA_DIR = Path("data/ticks")

N_MARKETS = 30          # number of top markets to track
FLUSH_SECONDS = 60      # write a Parquet file at least this often...
FLUSH_ROWS = 5000       # ...or once this many rows are buffered


async def discover_liquid_markets(session, limit=N_MARKETS):
    """Return the most liquid currently-open Kalshi market tickers.

    The default ``/markets`` feed leads with same-day multi-leg combo markets
    (``KXMVE*``); ``min_close_ts`` skips those. Markets are then ranked by 24h
    volume so the collector spends its websocket budget on books that move.
    """
    min_close_ts = int(time.time()) + 86400
    ranked = []
    cursor = None
    for _ in range(4):
        params = {"limit": 200, "min_close_ts": min_close_ts}
        if cursor:
            params["cursor"] = cursor
        async with session.get(
            f"{KALSHI_REST}/markets", params=params,
            timeout=aiohttp.ClientTimeout(total=15),
        ) as resp:
            data = await resp.json()
        for m in data.get("markets", []):
            if m["ticker"].startswith("KXMVE"):
                continue
            try:
                bid = float(m.get("yes_bid_dollars"))
                ask = float(m.get("yes_ask_dollars"))
            except (TypeError, ValueError):
                continue
            if bid <= 0 or ask <= 0:
                continue
            ranked.append((float(m.get("volume_24h_fp") or 0), m["ticker"]))
        cursor = data.get("cursor")
        if not cursor:
            break
    ranked.sort(reverse=True)
    return [ticker for _, ticker in ranked[:limit]]


class TickCollector:
    def __init__(self, n_markets=N_MARKETS):
        self.n_markets = n_markets
        self.books: dict[str, KalshiOrderBook] = {}
        self.buffer: list[dict] = []
        self.rows_written = 0
        self._last_seq = None
        self._last_flush = time.monotonic()

    def _apply_message(self, msg) -> bool:
        """Fold one parsed websocket message into the books.

        Returns ``True`` when a sequence gap is detected and the caller should
        reconnect (Kalshi will replay fresh snapshots). Touches only in-memory
        state, so it is unit-testable without a live socket.
        """
        mtype = msg.get("type")
        if mtype not in ("orderbook_snapshot", "orderbook_delta"):
            return False  # 'subscribed' / 'ok' / 'error' control frames

        body = msg.get("msg", {})
        ticker = body.get("market_ticker") or body.get("market_id")
        seq = msg.get("seq")

        # A non-consecutive seq means we missed an update — the local book is
        # now wrong. Bail out so the connection is rebuilt from a snapshot.
        if (mtype == "orderbook_delta" and self._last_seq is not None
                and seq is not None and seq != self._last_seq + 1):
            logger.warning("seq gap %s -> %s; resyncing", self._last_seq, seq)
            return True
        self._last_seq = seq

        book = self.books.setdefault(ticker, KalshiOrderBook(ticker))
        if mtype == "orderbook_snapshot":
            book.apply_snapshot(body.get("yes", []), body.get("no", []))
        else:
            book.apply_delta(body["price"], body["delta"], body["side"])

        feats = book.features()
        if feats is not None:
            feats["ts"] = datetime.now(timezone.utc).isoformat()
            feats["seq"] = seq
            self.buffer.append(feats)
        return False

    def _maybe_flush(self):
        due = (len(self.buffer) >= FLUSH_ROWS
               or time.monotonic() - self._last_flush >= FLUSH_SECONDS)
        if due:
            self.flush()

    def flush(self):
        """Write the buffered rows to a fresh Parquet file and clear the buffer."""
        self._last_flush = time.monotonic()
        if not self.buffer:
            return
        now = datetime.now(timezone.utc)
        out_dir = DATA_DIR / f"date={now:%Y-%m-%d}"
        out_dir.mkdir(parents=True, exist_ok=True)
        path = out_dir / f"kalshi_{now:%H%M%S_%f}.parquet"
        pd.DataFrame(self.buffer).to_parquet(path, index=False)
        self.rows_written += len(self.buffer)
        logger.info("flushed %d rows -> %s (total %d)",
                    len(self.buffer), path, self.rows_written)
        self.buffer.clear()

    async def _run_once(self, tickers):
        """One websocket session: subscribe and consume until it drops."""
        self._last_seq = None
        async with websockets.connect(KALSHI_WS, ping_interval=10) as ws:
            await ws.send(json.dumps({
                "id": 1,
                "cmd": "subscribe",
                "params": {
                    "channels": ["orderbook_delta"],
                    "market_tickers": tickers,
                },
            }))
            logger.info("subscribed to %d markets", len(tickers))
            async for raw in ws:
                if self._apply_message(json.loads(raw)):
                    return  # seq gap -> let the outer loop reconnect
                self._maybe_flush()

    async def run(self):
        async with aiohttp.ClientSession() as session:
            tickers = await discover_liquid_markets(session, self.n_markets)
        if not tickers:
            logger.error("no liquid markets found; nothing to collect")
            return
        logger.info("tracking %d markets: %s", len(tickers), ", ".join(tickers))

        backoff = 1
        while True:
            try:
                await self._run_once(tickers)
                backoff = 1
            except (websockets.ConnectionClosed, OSError) as e:
                logger.warning("connection lost: %s", e)
                backoff = min(backoff * 2, 30)
            except Exception:
                logger.exception("unexpected collector error")
                backoff = min(backoff * 2, 30)
            self.flush()              # keep buffered rows across the gap
            self.books.clear()        # stale — a reconnect replays snapshots
            await asyncio.sleep(backoff)


def main():
    from src import config  # noqa: F401  (configures root logging)
    logging.getLogger("websockets").setLevel(logging.WARNING)
    collector = TickCollector()
    try:
        asyncio.run(collector.run())
    except KeyboardInterrupt:
        collector.flush()
        logger.info("collector stopped; %d rows written total",
                    collector.rows_written)


if __name__ == "__main__":
    main()