ftdt-quant-lab/framework/data.py

"""
Hyperliquid data provider — historical candles, orderbook snapshots, and WebSocket streams.

Fetches OHLCV candles from Hyperliquid info API (candleSnapshot) and
provides them as pandas DataFrames (for VectorBT) and NT Bar objects
(for NautilusTrader backtesting).

WebSocket support: real-time orderbook, trades, mark prices via Hyperliquid WS.
"""
from __future__ import annotations

import asyncio
import json
import logging
import time
from datetime import datetime, timezone
from typing import AsyncIterator, Callable

import numpy as np
import pandas as pd
import requests

from nautilus_trader.model.data import Bar, BarSpecification, BarType
from nautilus_trader.model.enums import BarAggregation, PriceType
from nautilus_trader.model.identifiers import InstrumentId
from nautilus_trader.model.objects import Price, Quantity

logger = logging.getLogger(__name__)

TESTNET_API = "https://api.hyperliquid-testnet.xyz/info"
MAINNET_API = "https://api.hyperliquid.xyz/info"

WS_TESTNET = "wss://api.hyperliquid-testnet.xyz/ws"
WS_MAINNET = "wss://api.hyperliquid.xyz/ws"

INTERVAL_MAP: dict[str, str] = {
    "1m": "1m", "5m": "5m", "15m": "15m", "30m": "30m",
    "1h": "1h", "4h": "4h", "8h": "8h", "1d": "1d",
    "1w": "1w",
}

INTERVAL_TO_SECONDS: dict[str, int] = {
    "1m": 60, "5m": 300, "15m": 900, "30m": 1800,
    "1h": 3600, "4h": 14400, "8h": 28800, "1d": 86400,
    "1w": 604800,
}


class HyperliquidDataProvider:
    """Fetches and manages Hyperliquid market data."""

    def __init__(self, testnet: bool = True):
        self._api_url = TESTNET_API if testnet else MAINNET_API
        self._ws_url = WS_TESTNET if testnet else WS_MAINNET
        self._testnet = testnet

    # ── Historical candles ──────────────────────────────────────

    def fetch_candles(
        self,
        coin: str,
        interval: str = "1h",
        start_ms: int | None = None,
        end_ms: int | None = None,
        limit: int = 5000,
    ) -> pd.DataFrame:
        """Fetch OHLCV candles from Hyperliquid info API.

        Returns DataFrame with columns: open, high, low, close, volume, timestamp.
        Timestamp is UTC datetime index.
        """
        hl_interval = INTERVAL_MAP.get(interval, interval)
        now = int(time.time() * 1000)
        payload = {
            "type": "candleSnapshot",
            "req": {
                "coin": coin.upper(),
                "interval": hl_interval,
                "startTime": start_ms or (now - limit * INTERVAL_TO_SECONDS.get(interval, 3600) * 1000),
                "endTime": end_ms or now,
            },
        }
        resp = requests.post(self._api_url, json=payload, timeout=30)
        resp.raise_for_status()
        candles = resp.json()

        if not candles:
            return pd.DataFrame(columns=["open", "high", "low", "close", "volume", "timestamp"])

        rows = []
        for c in candles:
            rows.append({
                "open": float(c["o"]),
                "high": float(c["h"]),
                "low": float(c["l"]),
                "close": float(c["c"]),
                "volume": float(c["v"]),
                "timestamp": datetime.fromtimestamp(c["t"] / 1000, tz=timezone.utc),
            })
        df = pd.DataFrame(rows)
        df.set_index("timestamp", inplace=True)
        df.sort_index(inplace=True)
        return df

    def fetch_multi_candles(
        self,
        coins: list[str],
        interval: str = "1h",
        limit: int = 5000,
    ) -> dict[str, pd.DataFrame]:
        """Fetch candles for multiple coins in parallel."""
        results = {}
        for coin in coins:
            try:
                results[coin] = self.fetch_candles(coin, interval=interval, limit=limit)
            except Exception as e:
                logger.warning("Failed to fetch %s candles: %s", coin, e)
        return results

    def to_nt_bars(
        self,
        df: pd.DataFrame,
        instrument_id: InstrumentId,
        step: int = 1,
        bar_aggregation: BarAggregation = BarAggregation.MINUTE,
        price_type: PriceType = PriceType.LAST,
    ) -> list[Bar]:
        """Convert a pandas DataFrame of candles to NautilusTrader Bar objects."""
        spec = BarSpecification(step, bar_aggregation, price_type)
        bar_type = BarType(instrument_id, spec)
        bars = []

        for idx, row in df.iterrows():
            ts_event = int(idx.timestamp() * 1e9)
            ts_init = ts_event
            bar = Bar(
                bar_type=bar_type,
                open=Price(row["open"], instrument_id.venue.precision or 2),
                high=Price(row["high"], instrument_id.venue.precision or 2),
                low=Price(row["low"], instrument_id.venue.precision or 2),
                close=Price(row["close"], instrument_id.venue.precision or 2),
                volume=Quantity(row["volume"], 0),
                ts_event=ts_event,
                ts_init=ts_init,
            )
            bars.append(bar)

        return bars

    # ── Orderbook snapshots ─────────────────────────────────────

    def fetch_orderbook(self, coin: str) -> dict:
        """Get current L2 orderbook snapshot."""
        resp = requests.post(self._api_url, json={"type": "l2Book", "coin": coin.upper()}, timeout=10)
        resp.raise_for_status()
        data = resp.json()
        bids = [[float(l["px"]), float(l["sz"])] for l in data["levels"][0]]
        asks = [[float(l["px"]), float(l["sz"])] for l in data["levels"][1]]
        return {
            "bids": bids,
            "asks": asks,
            "timestamp": time.time(),
        }

    def fetch_orderbook_df(self, coin: str) -> tuple[pd.DataFrame, pd.DataFrame]:
        """Get orderbook as bid/ask DataFrames."""
        ob = self.fetch_orderbook(coin)
        bids_df = pd.DataFrame(ob["bids"], columns=["price", "size"])
        asks_df = pd.DataFrame(ob["asks"], columns=["price", "size"])
        return bids_df, asks_df

    # ── Mark prices ─────────────────────────────────────────────

    def fetch_mark_prices(self) -> dict[str, float]:
        """Get current mark prices for all assets."""
        resp = requests.post(self._api_url, json={"type": "metaAndAssetCtxs"}, timeout=10)
        resp.raise_for_status()
        data = resp.json()
        if not isinstance(data, list) or len(data) < 2:
            return {}
        universe = data[0].get("universe", [])
        ctxs = data[1]
        prices = {}
        for i, u in enumerate(universe):
            if i < len(ctxs):
                prices[u["name"]] = float(ctxs[i].get("markPx", 0))
        return prices

    # ── WebSocket streaming ─────────────────────────────────────

    async def stream_orderbook(self, coin: str) -> AsyncIterator[dict]:
        """Stream L2 orderbook updates via Hyperliquid WebSocket."""
        try:
            import websockets
        except ImportError:
            logger.error("websockets not installed; pip install websockets")
            return

        subscribe_msg = json.dumps({"method": "subscribe", "subscription": {"type": "l2Book", "coin": coin.upper()}})
        while True:
            try:
                async with websockets.connect(self._ws_url) as ws:
                    await ws.send(subscribe_msg)
                    async for msg in ws:
                        yield json.loads(msg)
            except Exception as e:
                logger.warning("WebSocket error: %s (reconnecting)", e)
                await asyncio.sleep(1)

    async def stream_prices(self, coins: list[str]) -> AsyncIterator[dict[str, float]]:
        """Stream mark prices via polling fallback (1s interval).

        Hyperliquid WebSocket doesn't have a simple 'mark prices' stream,
        so we poll the REST API with async sleep.
        """
        while True:
            try:
                prices = self.fetch_mark_prices()
                yield {c: prices.get(c, 0) for c in coins}
            except Exception as e:
                logger.warning("Price poll error: %s", e)
            await asyncio.sleep(1)