hk-weather-mkt/markets/trader.py

"""Trading execution engine for Polymarket weather markets.

Handles order placement, position sizing, and risk management
for the HK weather prediction market strategy.
"""

from datetime import datetime
from typing import Optional, Dict, List, Tuple
from dataclasses import dataclass, field

from .polymarket_client import PolymarketClient
from config import MIN_EDGE_BPS, MAX_POSITION_USDC, MIN_LIQUIDITY_USDC


@dataclass
class TradeSignal:
    """A trading signal from the strategy engine."""
    market_id: str
    condition_id: str
    question: str
    outcome_index: int
    outcome_label: str
    model_probability: float       # Our model-implied probability (0-100)
    market_probability: float      # Market-implied probability (0-100)
    edge_bps: float                # Edge in basis points
    recommended_size_usdc: float   # Kelly-recommended bet size
    max_size_usdc: float           # Maximum allowed position
    signal_type: str               # "buy_yes", "buy_no", "pass"


@dataclass
class ExecutionResult:
    """Result of a trade execution."""
    signal: TradeSignal
    success: bool
    order_id: Optional[str] = None
    filled_amount: float = 0.0
    avg_price: float = 0.0
    error: Optional[str] = None
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())


class Trader:
    """Execute trades based on strategy signals."""

    def __init__(
        self,
        client: PolymarketClient,
        private_key: str = "",
        funder_address: str = "",
        dry_run: bool = True,
    ):
        self.client = client
        self.private_key = private_key
        self.funder_address = funder_address
        self.dry_run = dry_run
        self.clob = None
        self.positions: Dict[str, float] = {}
        self.trade_history: List[ExecutionResult] = []

        if not dry_run and private_key:
            self._init_clob()

    def _init_clob(self):
        """Initialize CLOB client for live trading."""
        try:
            from py_clob_client.client import ClobClient
            from py_clob_client.clob_types import OrderArgs

            host = "https://clob.polymarket.com"
            chain_id = 137  # Polygon mainnet

            self.clob = ClobClient(
                host=host,
                key=self.private_key,
                chain_id=chain_id,
                funder=self.funder_address,
                signature_type=2,
            )
            print("CLOB client initialized for live trading")
        except Exception as e:
            print(f"CLOB init failed: {e}. Running in dry-run mode.")
            self.dry_run = True

    def execute_signal(self, signal: TradeSignal) -> ExecutionResult:
        """Execute a single trade signal."""
        if signal.signal_type == "pass":
            return ExecutionResult(
                signal=signal,
                success=True,
                note="No trade: edge below threshold",
            )

        # Get token ID
        token_id = self.client.get_clob_token_id(
            signal.condition_id, signal.outcome_index
        )

        if not token_id:
            return ExecutionResult(
                signal=signal,
                success=False,
                error="Could not get token ID",
            )

        # Calculate number of shares at size (each share = $1 if correct)
        price = signal.market_probability / 100.0
        size = min(signal.recommended_size_usdc, signal.max_size_usdc)

        if size < 1.0:
            return ExecutionResult(
                signal=signal,
                success=False,
                error=f"Size too small: ${size:.2f}",
            )

        if self.dry_run:
            return self._execute_dry_run(signal, token_id, size, price)
        else:
            return self._execute_live(signal, token_id, size, price)

    def _execute_dry_run(
        self, signal: TradeSignal, token_id: str, size: float, price: float
    ) -> ExecutionResult:
        """Simulate trade execution for testing."""
        result = ExecutionResult(
            signal=signal,
            success=True,
            order_id=f"DRY_RUN_{datetime.now().timestamp()}",
            filled_amount=size,
            avg_price=price,
        )
        self.trade_history.append(result)

        position_key = f"{signal.condition_id}_{signal.outcome_index}"
        self.positions[position_key] = self.positions.get(position_key, 0) + size

        print(f"  [DRY RUN] {signal.signal_type}: ${size:.2f} on '{signal.question}'"
              f" @ {price:.4f} (edge: {signal.edge_bps:.0f}bps)")

        return result

    def _execute_live(
        self, signal: TradeSignal, token_id: str, size: float, price: float
    ) -> ExecutionResult:
        """Execute real trade on Polymarket CLOB."""
        if not self.clob:
            return ExecutionResult(
                signal=signal,
                success=False,
                error="CLOB not initialized",
            )

        try:
            # Create a limit order (IOC to avoid partial fills on stale prices)
            order_args = {
                "token_id": token_id,
                "price": price,
                "size": size,
                "side": "BUY" if signal.signal_type == "buy_yes" else "SELL",
            }

            response = self.clob.create_and_post_order(
                order_args, orderType="GTC"
            )

            result = ExecutionResult(
                signal=signal,
                success=True,
                order_id=response.get("orderID", ""),
                filled_amount=float(response.get("filled_size", 0)),
                avg_price=float(response.get("avg_price", price)),
            )
            self.trade_history.append(result)

            print(f"  [LIVE] {signal.signal_type}: ${size:.2f} on '{signal.question}'"
                  f" @ {price:.4f} (edge: {signal.edge_bps:.0f}bps)")

            return result

        except Exception as e:
            return ExecutionResult(
                signal=signal,
                success=False,
                error=str(e),
            )

    def get_positions_summary(self) -> Dict:
        """Get summary of current positions and P&L."""
        total_bet = sum(self.positions.values())
        open_trades = len([t for t in self.trade_history if t.success])

        return {
            "total_positions_value_usdc": total_bet,
            "num_open_trades": open_trades,
            "num_markets": len(self.positions),
            "positions": self.positions,
            "dry_run": self.dry_run,
        }

    def cancel_all_orders(self):
        """Cancel all open orders. Only works in live mode."""
        if self.dry_run or not self.clob:
            print("Cannot cancel orders in dry-run mode")
            return

        try:
            self.clob.cancel_all()
            print("All orders cancelled")
        except Exception as e:
            print(f"Cancellation error: {e}")