market_maker.py

"""High-frequency market making bot for Polymarket, using the Avellaneda-Stoikov model.

This bot implements the A-S model for optimal quoting, combined with a toxicity
filter to manage adverse selection risk.
"""

from __future__ import annotations

import logging
import signal
import time
from datetime import datetime, timezone
from decimal import Decimal, ROUND_DOWN

from config import BotConfig
from market_sources import MarketSpec
from py_clob_client_v2.exceptions import PolyApiException
from polymarket_adapter import FatalTradingError, PolymarketAdapter, Position, Quote, TokenConfig
from order_reconciler import OrderReconciler
from quote_engine import QuoteEngine, TargetOrder
from emergency_exit import EmergencyExitEngine
from data_recorder import DataRecorder

LOGGER = logging.getLogger(__name__)


class MarketMaker:
    def __init__(
        self,
        config: BotConfig,
        adapter: PolymarketAdapter,
        spec: MarketSpec | None = None,
    ) -> None:
        self.config = config
        self.adapter = adapter
        self.spec = spec
        self.quote_engine = QuoteEngine(config, adapter)
        self.emergency_exit_engine = EmergencyExitEngine(config)
        self.order_reconciler = OrderReconciler(self.config, self.adapter)
        self.data_recorder: DataRecorder | None = None
        self.should_stop = False
        self.yes: TokenConfig
        self.no: TokenConfig
        self.condition_id: str
        self.initial_market_exposure = Decimal("0")
        self.last_midpoints: dict[str, Decimal] = {}
        self.risk_off = False
        self.last_placed_orders: list[TargetOrder] = []
        self.dynamic_order_notional: Decimal = config.order_notional_usdc
        self.last_success_order_time: float = time.time()

    def run(self) -> None:
        self.bootstrap()
        failures = 0
        while not self.should_stop:
            try:
                self.tick()
                failures = 0
                time.sleep(self.next_sleep_seconds())
            except KeyboardInterrupt:
                self.should_stop = True
            except FatalTradingError:
                self.should_stop = True
                LOGGER.exception("Fatal trading error; stopping instead of retrying")
            except Exception:
                failures += 1
                sleep_for = min(30, 2 ** min(failures, 5))
                LOGGER.exception("Main loop error; reconnecting after %ss", sleep_for)
                time.sleep(sleep_for)

        self._cancel_working_orders()
        LOGGER.info("Stopped; old orders cancelled")
        if self.data_recorder:
            self.data_recorder.stop()

    def shutdown(self) -> None:
        self.should_stop = True
        self._cancel_working_orders()

    def next_sleep_seconds(self) -> float:
        return float(self.config.refresh_interval_seconds)

    def bootstrap(self) -> None:
        self.config.validate()
        if self.spec is None:
            self.yes, self.no, self.condition_id = self.adapter.resolve_tokens()
        else:
            self.yes, self.no, self.condition_id = self.adapter.resolve_market_spec(self.spec)
        self._install_signal_handlers()

        if self.data_recorder is None and self.condition_id:
            self.data_recorder = DataRecorder(self.condition_id)

        if self.config.cancel_on_start:
            self._cancel_working_orders()

        self.initial_market_exposure = self._current_market_exposure()

        LOGGER.info(
            "Starting Market Maker: condition_id=%s dry_run=%s order_type=%s post_only=%s order_notional=%s refresh=%ss",
            self.condition_id or "<token-only>",
            self.config.dry_run,
            self.config.order_type,
            self.config.post_only,
            self.config.order_notional_usdc,
            self.config.refresh_interval_seconds,
        )

    def tick(self) -> None:
        log_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "condition_id": self.condition_id,
            "reasons": [],
        }

        # --- Restore order notional gradually after failures ---
        if (
            self.dynamic_order_notional < self.config.order_notional_usdc
            and time.time() - self.last_success_order_time > 30
        ):
            new_notional = self.dynamic_order_notional + (
                self.config.order_notional_usdc - self.dynamic_order_notional
            ) * Decimal("0.1")
            self.dynamic_order_notional = min(new_notional, self.config.order_notional_usdc)
            log_entry["reasons"].append(f"gradual_notional_restore_to_{self.dynamic_order_notional:.2f}")
            LOGGER.info(
                "[STRATEGY] Gradually restoring order notional to $%.2f", self.dynamic_order_notional
            )

        # 1. Get market data
        tokens = (self.yes, self.no)
        snapshots = {token.token_id: self.adapter.get_order_book_snapshot(token) for token in tokens}
        quotes = {token_id: snapshot.quote for token_id, snapshot in snapshots.items()}
        positions = self.adapter.get_positions(tokens, quotes=quotes)

        # Log quotes and positions
        for token in tokens:
            pos = positions[token.token_id]
            q = quotes[token.token_id]
            log_entry[f"{token.outcome}_pos_size"] = pos.size
            log_entry[f"{token.outcome}_pos_avg_price"] = pos.avg_price
            log_entry[f"{token.outcome}_quote_bid"] = q.bid
            log_entry[f"{token.outcome}_quote_ask"] = q.ask

        # --- Ignore small "dust" positions ---
        for token in tokens:
            position = positions[token.token_id]
            if 0 < position.notional < 1:
                log_entry["reasons"].append(f"ignore_dust_position_{token.outcome}")
                positions[token.token_id] = Position(
                    token_id=position.token_id, size=Decimal("0"), avg_price=Decimal("0"), current_price=position.current_price
                )

        # 2. Calculate risk and PnL
        self._log_status_and_positions(tokens, positions)

        # 3. Calculate A-S model inputs
        seconds_left = self._get_seconds_to_expiry()
        time_to_expiry = Decimal(max(0, seconds_left)) / Decimal(365 * 24 * 3600)
        is_close_only_window = 0 < seconds_left <= self.config.close_only_hours_before_end * 3600
        log_entry["time_to_expiry_years"] = time_to_expiry

        # 4. Check global risk limits
        if self._is_risk_limit_breached(positions):
            log_entry["reasons"].append("global_risk_limit_breached")
            self._cancel_working_orders()
            if self.data_recorder:
                self.data_recorder.record(log_entry)
            return

        # 5. Main loop for each token (YES/NO)
        close_only = self.risk_off or is_close_only_window
        log_entry["close_only_mode"] = close_only
        if close_only:
            reason = "global_risk_off_latch" if self.risk_off else "close_only_window"
            log_entry["reasons"].append(f"close_only_enabled_{reason}")
            LOGGER.warning("Close-only mode active; new buy quotes are disabled")

        orders: list[TargetOrder] = []
        for token in tokens:
            quote = quotes[token.token_id]
            position = positions[token.token_id]
            log_entry[f"{token.outcome}_gamma"] = self.config.risk_aversion
            
            imbalance = self._calculate_orderbook_imbalance(snapshots[token.token_id].bids, snapshots[token.token_id].asks)
            log_entry[f"{token.outcome}_imbalance"] = imbalance
            dynamic_gamma = self.config.risk_aversion
            if abs(imbalance) > self.config.toxicity_imbalance_threshold:
                dynamic_gamma *= self.config.toxicity_gamma_multiplier
                log_entry["reasons"].append(f"toxicity_filter_active_{token.outcome}")
                log_entry[f"{token.outcome}_gamma"] = dynamic_gamma
                LOGGER.warning(
                    "[TOXICITY_FILTER] Market for %s is toxic (imbalance=%.2f). Increasing gamma to %.2f for this tick.",
                    token.outcome, imbalance, dynamic_gamma
                )
            
            emergency_signal = self.emergency_exit_engine.evaluate(
                token=token, quote=quote, position=position, previous_mid=self.last_midpoints.get(token.token_id), imbalance=imbalance
            )
            if emergency_signal.triggered:
                log_entry["reasons"].append(f"emergency_exit_{token.outcome}_{emergency_signal.reason}")
                emergency_order = self.emergency_exit_engine.build_maker_exit_order(token, quote, position, emergency_signal)
                if emergency_order:
                    orders.append(emergency_order)
                    if self.config.latch_on_stop_loss:
                        self.risk_off = self.config.risk_off_after_stop
                    self.last_midpoints[token.token_id] = quote.mid
                    continue

            allow_new_buy = not close_only and not self._position_risk_off(token, quote, position)

            # --- A-S Armor: Midpoint Jump Detection & Momentum Rider ---
            if self._midpoint_jump_detected(token, quote):
                log_entry['reasons'].append(f"midpoint_jump_detected_{token.outcome}")
                LOGGER.warning("[STRATEGY] %s midpoint jumped too fast.", token.outcome)

                # --- Momentum Rider Mode ---
                if (self.config.enable_momentum_rider_mode and 
                    (quote.mid >= self.config.momentum_entry_threshold or quote.mid <= (Decimal('1') - self.config.momentum_entry_threshold))):
                    
                    log_entry['reasons'].append(f"momentum_rider_triggered_{token.outcome}")
                    LOGGER.info("[MOMENTUM] Triggered for %s at price %s", token.outcome, quote.mid)
                    
                    # Determine side based on which way the price moved
                    side = "BUY" if quote.mid > self.last_midpoints.get(token.token_id, quote.mid) else "SELL"
                    
                    if side == "BUY":
                        # Place an aggressive buy order slightly below the new mid
                        momentum_price = self.adapter.round_price(quote.mid - token.tick_size * 2, token.tick_size, ROUND_DOWN)
                        shares = self._shares_for_notional(self.config.momentum_order_notional, momentum_price, token)
                        if shares:
                            orders.append(TargetOrder(token=token, side="BUY", price=momentum_price, shares=shares, reason="momentum_buy"))
                            LOGGER.info("[MOMENTUM] Placing aggressive BUY order for %s at %s", token.outcome, momentum_price)
                    else: # side == "SELL"
                        # Place an aggressive sell order slightly above the new mid
                        momentum_price = self.adapter.round_price(quote.mid + token.tick_size * 2, token.tick_size, ROUND_UP)
                        shares = self._shares_for_notional(self.config.momentum_order_notional, momentum_price, token)
                        if shares:
                            # We can only sell what we have
                            shares_to_sell = min(position.size, shares)
                            if shares_to_sell >= token.min_order_size:
                                orders.append(TargetOrder(token=token, side="SELL", price=momentum_price, shares=shares_to_sell, reason="momentum_sell"))
                                LOGGER.info("[MOMENTUM] Placing aggressive SELL order for %s at %s", token.outcome, momentum_price)
                
                # In either case (momentum or not), we disallow normal A-S buy orders after a jump
                allow_new_buy = False

            if allow_new_buy:
                maybe_buy_order = self._build_buy_order(token, quote, position, time_to_expiry, dynamic_gamma)
                if maybe_buy_order:
                    orders.append(maybe_buy_order)
                else:
                    log_entry["reasons"].append(f"buy_order_not_built_{token.outcome}")
            else:
                log_entry["reasons"].append(f"new_buys_disallowed_{token.outcome}")

            maybe_sell_order = self._build_sell_order(token, quote, position, time_to_expiry, dynamic_gamma, force_exit=close_only)
            if maybe_sell_order:
                orders.append(maybe_sell_order)
            else:
                log_entry["reasons"].append(f"sell_order_not_built_{token.outcome}")

            self.last_midpoints[token.token_id] = quote.mid

        self._sync_target_orders(orders, tokens, log_entry)
        
        if self.data_recorder:
            self.data_recorder.record(log_entry)

    def _build_buy_order(self, token: TokenConfig, quote: Quote, position: Position, time_to_expiry: Decimal, gamma: Decimal) -> TargetOrder | None:
        market_exposure = sum(p.notional for p in self.adapter.get_positions((self.yes, self.no), quotes={self.yes.token_id: quote} if token.token_id == self.yes.token_id else {self.no.token_id: quote}).values())
        if position.notional >= self.config.max_token_exposure_usdc or market_exposure >= self.config.max_market_exposure_usdc:
            return None

        limit_price, diagnostics = self.quote_engine.calculate_limit_price(token, "BUY", quote, position, time_to_expiry, gamma)
        self.quote_engine.log_diagnostics(diagnostics)

        if limit_price <= 0 or limit_price < self.config.min_price:
            return None

        remaining_market = self.config.max_market_exposure_usdc - market_exposure
        remaining_token = self.config.max_token_exposure_usdc - position.notional
        order_notional = min(self.dynamic_order_notional, remaining_market, remaining_token)
        
        shares = self._shares_for_notional(order_notional, limit_price, token)
        if shares is None:
            return None

        return TargetOrder(token=token, side="BUY", price=limit_price, shares=shares, reason="as_model_buy")

    def _build_sell_order(self, token: TokenConfig, quote: Quote, position: Position, time_to_expiry: Decimal, gamma: Decimal, force_exit: bool = False) -> TargetOrder | None:
        if position.size <= 0:
            return None

        # Determine if this is a risk-reducing sell (stop-loss, take-profit, or forced exit)
        is_risk_exit = force_exit or self._position_risk_off(token, quote, position)
        reason = "risk_exit" if is_risk_exit else "as_model_sell"

        limit_price, diagnostics = self.quote_engine.calculate_limit_price(token, "SELL", quote, position, time_to_expiry, gamma)
        self.quote_engine.log_diagnostics(diagnostics)

        if limit_price <= 0 or limit_price > self.config.max_price:
            return None

        shares_to_sell = position.size if is_risk_exit else (self.dynamic_order_notional / limit_price)
        shares = min(position.size, shares_to_sell)

        if shares < token.min_order_size:
            return None

        return TargetOrder(token=token, side="SELL", price=limit_price, shares=shares, reason=reason)

    # ... (Keep all helper methods like _position_return, _position_risk_off, etc.) ...
    def _position_return(self, position: Position) -> Decimal:
        if position.size <= 0 or position.avg_price <= 0:
            return Decimal("0")
        return (position.current_price - position.avg_price) / position.avg_price

    def _position_risk_off(self, token: TokenConfig, quote: Quote, position: Position) -> bool:
        ret = self._position_return(position)
        if position.size <= 0:
            return False
        if ret <= -self.config.stop_loss_fraction:
            LOGGER.warning(
                "%s stop loss hit: return=%.2f%% avg=%s mid=%s",
                token.outcome, float(ret * Decimal("100")), position.avg_price, quote.mid)
            if self.config.latch_on_stop_loss:
                self.risk_off = self.config.risk_off_after_stop
                LOGGER.warning("[DEBUG_RISK_OFF] Global risk_off latch has been ACTIVATED due to stop-loss.")
            return True
        if ret >= self.config.take_profit_fraction:
            LOGGER.info(
                "%s take profit hit: return=%.2f%% avg=%s mid=%s",
                token.outcome, float(ret * Decimal("100")), position.avg_price, quote.mid)
            return True
        return False

    def _get_seconds_to_expiry(self) -> float:
        end_date_iso = self.yes.end_date_iso or self.no.end_date_iso
        if not end_date_iso:
            return 0.0
        try:
            normalized = end_date_iso.replace("Z", "+00:00")
            end_dt = datetime.fromisoformat(normalized)
            if end_dt.tzinfo is None:
                end_dt = end_dt.replace(tzinfo=timezone.utc)
            return (end_dt - datetime.now(timezone.utc)).total_seconds()
        except ValueError:
            LOGGER.warning("Could not parse market end_date_iso: %s", end_date_iso)
            return 0.0

    def _is_risk_limit_breached(self, positions: dict[str, Position]) -> bool:
        account_exposure = self.adapter.get_global_exposure()
        market_exposure = sum(p.notional for p in positions.values())
        bot_exposure = max(Decimal("0"), market_exposure - self.initial_market_exposure)
        risk_exposure = account_exposure if self.config.count_existing_positions_in_global_limit else bot_exposure

        if risk_exposure >= self.config.max_global_exposure_usdc:
            LOGGER.warning(
                "Global exposure limit hit; risk_exposure=$%.2f max=$%.2f count_existing=%s",
                float(risk_exposure), float(self.config.max_global_exposure_usdc), self.config.count_existing_positions_in_global_limit)
            return True
        return False

    def _log_status_and_positions(self, tokens: tuple[TokenConfig, ...], positions: dict[str, Position]) -> None:
        account_exposure = self.adapter.get_global_exposure()
        market_exposure = sum(p.notional for p in positions.values())
        unrealized_pnl = sum(p.unrealized_pnl for p in positions.values())
        rewards = self.adapter.today_rewards_estimate(self.condition_id)

        LOGGER.info(
            "[STATUS] Market=%s | AccountExposure=$%.2f | MarketExposure=$%.2f | uPnL=$%.2f | Rewards~$%.4f",
            self.condition_id or "N/A", float(account_exposure), float(market_exposure), float(unrealized_pnl), float(rewards))

        for token in tokens:
            pos = positions[token.token_id]
            if pos.size > 0:
                LOGGER.info(
                    "[POSITION] %s: Size=%s | AvgPrice=%s | CurPrice=%s | uPnL=$%.2f",
                    token.outcome, pos.size, pos.avg_price, pos.current_price, float(pos.unrealized_pnl))

    # ... (Keep remaining helper methods like _calculate_orderbook_imbalance, _shares_for_notional, etc.) ...
    def _calculate_orderbook_imbalance(self, bids: list[tuple[Decimal, Decimal]], asks: list[tuple[Decimal, Decimal]]) -> Decimal:
        if not bids or not asks:
            return Decimal("0")
        levels = self.config.orderbook_imbalance_levels
        bid_depth = sum(size for _price, size in sorted(bids, reverse=True)[:levels])
        ask_depth = sum(size for _price, size in sorted(asks)[:levels])
        total_depth = bid_depth + ask_depth
        if total_depth <= 0:
            return Decimal("0")
        return (bid_depth - ask_depth) / total_depth

    def _midpoint_jump_detected(self, token: TokenConfig, quote: Quote) -> bool:
        last_mid = self.last_midpoints.get(token.token_id)
        if not last_mid or last_mid <= 0:
            return False
        move = abs(quote.mid - last_mid) / last_mid
        return move >= self.config.max_midpoint_move_fraction

    def _install_signal_handlers(self) -> None:
        def stop(_signum, _frame) -> None:
            LOGGER.info("Stop signal received")
            self.should_stop = True
        signal.signal(signal.SIGINT, stop)
        signal.signal(signal.SIGTERM, stop)

    def _current_market_exposure(self) -> Decimal:
        positions = self.adapter.get_positions((self.yes, self.no))
        return sum(position.notional for position in positions.values())

    def _cancel_working_orders(self) -> None:
        if self.condition_id:
            self.adapter.cancel_market_orders(self.condition_id)
        else:
            self.adapter.cancel_token_orders(self.yes.token_id)
            self.adapter.cancel_token_orders(self.no.token_id)
        self.last_placed_orders = []

    def _sync_target_orders(self, orders: list[TargetOrder], tokens: tuple[TokenConfig, TokenConfig], log_entry: dict) -> None:
        plan = self.order_reconciler.plan(
            orders, self.last_placed_orders, tokens, condition_id=self.condition_id, use_live_orders=not self.config.dry_run)
        
        log_entry["reconcile_reason"] = plan.reason
        log_entry["reconcile_keep_n"] = len(plan.keep)
        log_entry["reconcile_place_n"] = len(plan.place)
        log_entry["reconcile_cancel_n"] = len(plan.cancel_order_ids)
        log_entry["reconcile_cancel_all"] = plan.cancel_all
        
        LOGGER.info(
            "[RECONCILE] %s | keep=%d place=%d cancel=%d cancel_all=%s",
            plan.reason, len(plan.keep), len(plan.place), len(plan.cancel_order_ids), plan.cancel_all)
        
        if not plan.has_work:
            if orders:
                 log_entry["reasons"].append("quotes_unchanged")
                 LOGGER.info("[STRATEGY] Quotes unchanged; keeping existing GTC PostOnly orders")
            self.last_placed_orders = orders
            return

        self._cancel_reconciled_orders(plan)
        placed_orders, placement_failures = self._place_orders(plan.place)
        
        log_entry["placed_orders"] = [o.to_dict() for o in placed_orders]
        log_entry["placement_failures"] = placement_failures

        self.last_placed_orders = plan.keep + placed_orders

    def _cancel_reconciled_orders(self, plan: ReconciliationPlan) -> None:
        if plan.cancel_all:
            self._cancel_working_orders()
            return
        if plan.cancel_order_ids:
            self.adapter.cancel_orders(plan.cancel_order_ids)

    def _place_orders(self, orders: list[TargetOrder]) -> tuple[list[TargetOrder], list[dict]]:
        LOGGER.info("[STRATEGY] Placing %d %s targets", len(orders), self.config.order_type)
        placed_orders: list[TargetOrder] = []
        placement_failures: list[dict] = []

        for order in orders:
            LOGGER.info(
                "[TRADE] Placing %s %s: Price=%s | Shares=%s | Notional=$%.2f | Reason=%s",
                order.side, order.token.outcome, order.price, order.shares, float(order.notional), order.reason)
            
            success, reason = self._execute_order_with_retry(order)
            if success:
                placed_orders.append(order)
            else:
                placement_failures.append({
                    "outcome": order.token.outcome,
                    "side": order.side,
                    "price": order.price,
                    "reason": reason,
                })

        return placed_orders, placement_failures
    
    def _execute_order_with_retry(self, order: TargetOrder, max_retries: int = 3) -> tuple[bool, str]:
        current_notional = self.dynamic_order_notional
        current_shares = order.shares
        current_price = order.price
        for attempt in range(max_retries):
            try:
                LOGGER.info(
                    "[TRADE] Attempt %d: Placing %s %s: Price=%s | Shares=%s | Notional=$%.2f",
                    attempt + 1, order.side, order.token.outcome, current_price, current_shares, float(current_price * current_shares))
                response = self.adapter.place_limit_order(
                    token=order.token, side=order.side, price=current_price, size=current_shares, dry_run=self.config.dry_run)
                if isinstance(response, dict) and response.get("success") is False:
                    error_msg = str(response.get("errorMsg") or response.get("error") or "").lower()
                    if "post only" in error_msg or "cross" in error_msg:
                        LOGGER.warning(
                            "[POST_ONLY] %s %s @ %s rejected: %s", order.side, order.token.outcome, current_price, error_msg)
                        return False, "post_only_rejected"
                    raise PolyApiException(error_msg)

                LOGGER.info("Successfully placed order.")
                self.last_success_order_time = time.time()
                return True, "placed"
            except PolyApiException as exc:
                msg = str(exc.error_msg).lower()
                if "not enough balance" in msg:
                    LOGGER.warning(
                        "[RETRY] Attempt %d failed: Not enough balance. Reducing notional and retrying.", attempt + 1)
                    current_notional *= Decimal("0.9")
                    new_shares = self._shares_for_notional(current_notional, current_price, order.token)
                    if new_shares is None:
                        LOGGER.error("Reduced notional is too small to place order. Aborting retry.")
                        self.dynamic_order_notional = max(current_notional, Decimal("1.0"))
                        return False, "not_enough_balance_final"
                    current_shares = new_shares
                    time.sleep(0.5)
                    continue
                elif "post only" in msg or "order crosses book" in msg:
                    LOGGER.warning(
                        "[POST_ONLY] %s %s @ %s rejected (would cross book)", order.side, order.token.outcome, current_price)
                    return False, "post_only_rejected"
                else:
                    LOGGER.error("[API_ERROR] Unhandled API error on order placement: %s", msg)
                    return False, f"unhandled_api_error:{msg}"
        
        final_reason = f"failed_after_{max_retries}_retries"
        LOGGER.error(
            "Failed to place order for %s %s after %d retries.", order.side, order.token.outcome, max_retries)
        self.dynamic_order_notional = max(current_notional, Decimal("1.0"))
        return False, final_reason

    def _shares_for_notional(self, order_notional: Decimal, price: Decimal, token: TokenConfig) -> Decimal | None:
        if price <= 0:
            return None
        desired_shares = (order_notional / price).quantize(Decimal("0.000001"), ROUND_DOWN)
        if desired_shares >= token.min_order_size:
            return desired_shares
        min_notional = token.min_order_size * price
        remaining_market = self.config.max_market_exposure_usdc - self._current_market_exposure()
        if min_notional <= self.config.order_notional_usdc and min_notional <= remaining_market:
            return token.min_order_size
        return None