ftdt-quant-lab/strategies/kalman_pairs/trading_system.py

"""
Kalman Pairs Trading System — Production Orchestrator.

Integrates pair discovery, Kalman filtering, signal generation,
position management, and risk controls into a single callable system.

Design:
  - Stateless between ticks — all state held in KalmanPairsTrader instances.
  - Multi-pair: manages N independent pairs simultaneously.
  - Risk overlay: per-pair stop-loss, max position, max drawdown.
  - Capital allocation: equal-weight or volatility-weighted.
  - Clean interface compatible with live node + backtester.

Usage (live):
    system = KalmanPairsTradingSystem(config)
    system.initialize(price_data)
    for tick in price_stream:
        signals = system.step(tick)

Usage (backtest):
    system = KalmanPairsTradingSystem(config)
    results = system.run_backtest(price_data)
"""

from __future__ import annotations

import numpy as np
from dataclasses import dataclass, field
from typing import Optional
from pathlib import Path
import json
import time

# Internal imports
from .kalman_filter import KalmanPairsTrader
from .pair_discovery import discover_pairs


# ═══════════════════════ Config ═════════════════════════════

@dataclass
class KalmanPairsConfig:
    """Configuration for the Kalman Pairs Trading System."""

    # ── Universe ──
    tickers: list[str] = field(default_factory=lambda: ["BTC", "ETH"])
    sector_constraint: Optional[str] = None  # e.g., None = any, "BTC" = BTC-only pairs

    # ── Pair Discovery ──
    min_half_life: float = 1.0
    max_half_life: float = 20.0
    max_pairs: int = 5
    coint_sig_level: float = 0.05

    # ── Kalman Filter ──
    transition_covariance: float = 1e-4
    observation_covariance: float = 1e-2
    warmup_bars: int = 50

    # ── Trading ──
    z_entry: float = 2.0
    z_exit: float = 0.5
    z_stop: float = 4.0
    trade_size_usd: float = 100.0      # Notional per leg
    max_position_per_pair: int = 1     # Max 1 unit long/short at a time

    # ── Risk ──
    max_drawdown_pct: float = 0.15     # Stop trading if equity drops > 15%
    max_daily_trades: int = 50         # Circuit breaker

    # ── Backtest ──
    transaction_cost_bps: float = 2.5  # 2.5 bps = 0.025% per leg (taker)
    initial_capital: float = 10000.0

    @classmethod
    def from_yaml(cls, path: str | Path) -> "KalmanPairsConfig":
        """Load from YAML. Falls back to defaults if YAML not available."""
        import yaml  # may not be installed
        with open(path) as f:
            data = yaml.safe_load(f)
        return cls(**data.get("kalman_pairs", data))

    @classmethod
    def from_dict(cls, d: dict) -> "KalmanPairsConfig":
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})


# ═══════════════════════ System ═════════════════════════════

class KalmanPairsTradingSystem:
    """
    Production Kalman Pairs Trading System.

    Manages multiple independent pairs, each with its own Kalman filter,
    and aggregates signals through a unified risk layer.
    """

    def __init__(self, config: KalmanPairsConfig | dict) -> None:
        if isinstance(config, dict):
            config = KalmanPairsConfig.from_dict(config)
        self.config = config

        # Active pair traders
        self.traders: dict[tuple[str, str], KalmanPairsTrader] = {}
        self.pair_info: dict[tuple[str, str], dict] = {}

        # Equity tracking
        self.capital = config.initial_capital
        self.peak_capital = config.initial_capital
        self.equity_curve: list[dict] = []
        self.daily_trades: int = 0
        self.daily_reset_time: float = time.time()

        # Trade log
        self.trades: list[dict] = []

    def initialize(self, price_data: dict[str, np.ndarray]) -> list[dict]:
        """
        Discover pairs and initialize Kalman traders.

        Args:
            price_data: {ticker: np.array of prices}

        Returns:
            List of discovered pair info dicts.
        """
        pairs = discover_pairs(
            price_data,
            sector_constraint=self.config.sector_constraint,
            min_half_life=self.config.min_half_life,
            max_half_life=self.config.max_half_life,
            sig_level=self.config.coint_sig_level,
        )

        # Take top N pairs by half-life (fastest mean reversion)
        pairs = pairs[: self.config.max_pairs]

        for p in pairs:
            key = p["pair"]
            self.pair_info[key] = p

            trader = KalmanPairsTrader(
                transition_covariance=self.config.transition_covariance,
                observation_covariance=self.config.observation_covariance,
                z_entry=self.config.z_entry,
                z_exit=self.config.z_exit,
                z_stop=self.config.z_stop,
                warmup_bars=self.config.warmup_bars,
            )
            self.traders[key] = trader

        return pairs

    def step(self, prices: dict[str, float]) -> dict:
        """
        Process one bar update for all active pairs.

        Args:
            prices: {ticker: current_price} for this bar.

        Returns:
            dict with: signals (list), equity, drawdown_pct, positions, alpha, beta
        """
        # Reset daily trade counter
        now = time.time()
        if now - self.daily_reset_time > 86400:
            self.daily_trades = 0
            self.daily_reset_time = now

        signals = []
        total_pnl = 0.0

        for key, trader in self.traders.items():
            t1, t2 = key
            if t1 not in prices or t2 not in prices:
                continue

            X = prices[t1]  # independent
            Y = prices[t2]  # dependent

            result = trader.step(X, Y)

            if result["signal"] != 0 and self.daily_trades < self.config.max_daily_trades:
                # Apply risk checks
                if self._check_risk():
                    signal = {
                        "pair": list(key),
                        "signal": result["signal"],
                        "spread": result["spread"],
                        "z_score": result["z_score"],
                        "alpha": result["alpha"],
                        "beta": result["beta"],
                        "position": result["position"],
                        "trade_size": self.config.trade_size_usd,
                    }
                    signals.append(signal)
                    self.daily_trades += 1

        # Update equity (simplified — full PnL in backtester)
        total_equity = self.capital + total_pnl
        self.peak_capital = max(self.peak_capital, total_equity)
        dd_pct = (self.peak_capital - total_equity) / self.peak_capital if self.peak_capital > 0 else 0.0

        self.equity_curve.append({
            "t": now,
            "equity": round(total_equity, 2),
            "dd": round(dd_pct, 4),
        })

        return {
            "signals": signals,
            "equity": round(total_equity, 2),
            "drawdown_pct": round(dd_pct, 4),
            "positions": {str(k): t.position for k, t in self.traders.items()},
            "alpha": {str(k): t.alpha for k, t in self.traders.items()},
            "beta": {str(k): t.beta for k, t in self.traders.items()},
        }

    def _check_risk(self) -> bool:
        """Return False if any risk limit is breached."""
        if self.peak_capital > 0:
            dd = (self.peak_capital - self.capital) / self.peak_capital
            if dd > self.config.max_drawdown_pct:
                return False
        return True

    def get_state(self) -> dict:
        """Return current system state for monitoring/dashboard."""
        return {
            "capital": round(self.capital, 2),
            "peak_capital": round(self.peak_capital, 2),
            "drawdown_pct": round(
                (self.peak_capital - self.capital) / self.peak_capital * 100
                if self.peak_capital > 0 else 0, 2
            ),
            "active_pairs": len(self.traders),
            "daily_trades": self.daily_trades,
            "positions": {str(k): t.position for k, t in self.traders.items()},
            "alpha": {str(k): round(t.alpha, 6) for k, t in self.traders.items()},
            "beta": {str(k): round(t.beta, 6) for k, t in self.traders.items()},
        }