"""
Spread Dynamics Model — model how spread changes based on supply/demand.

Improves quote placement by predicting spread movements.
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

import numpy as np
from numba import njit


@njit(cache=True)
def compute_spread_tendency(
    current_spread_bps: float,
    bid_depth: float,
    ask_depth: float,
    recent_trade_imbalance: float,
    toxicity: float,
    volatility: float,
) -> float:
    """
    Compute spread tendency (positive = tightening, negative = widening).

    Factors:
      - Depth imbalance: more depth on one side → spread tends to tighten
      - Trade imbalance: buying pressure → ask side thins → spread widens
      - Toxicity: toxic flow widens spread
      - Volatility: high volatility widens spread

    Returns tendency in bps per second.
    """
    # Depth factor: balanced depth → tightening
    depth_balance = (bid_depth - ask_depth) / max(bid_depth + ask_depth, 1e-12)
    depth_factor = -depth_balance * 0.5  # negative = tightening when balanced

    # Trade imbalance factor: buying pressure widens spread
    trade_factor = recent_trade_imbalance * 0.3

    # Toxicity factor: toxic flow widens spread
    toxicity_factor = toxicity * 0.5

    # Volatility factor: high volatility widens spread
    volatility_factor = volatility * 0.02

    return depth_factor + trade_factor + toxicity_factor + volatility_factor


@njit(cache=True)
def predict_spread(
    current_spread_bps: float,
    spread_tendency: float,
    time_horizon_s: float,
    min_spread_bps: float = 0.1,
    max_spread_bps: float = 100.0,
) -> float:
    """
    Predict spread after time_horizon_s.

    Model: spread adjusts toward equilibrium with mean reversion.
    """
    # Mean reversion toward current level
    reversion_rate = 0.1  # 10% reversion per second
    target = current_spread_bps + spread_tendency * time_horizon_s
    target = max(min_spread_bps, min(max_spread_bps, target))

    # Apply mean reversion
    predicted = current_spread_bps + (target - current_spread_bps) * (1 - math.exp(-reversion_rate * time_horizon_s))
    return max(min_spread_bps, min(max_spread_bps, predicted))


class SpreadDynamicsModel:
    """
    Spread dynamics model for the CWM.

    Predicts spread movements to improve quote placement.
    """

    def __init__(self) -> None:
        self._spread_history: list[float] = []
        self._last_spread_bps: float = 0.0

    def update(self, spread_bps: float) -> None:
        """Update with current spread."""
        self._spread_history.append(spread_bps)
        self._last_spread_bps = spread_bps
        # Keep only recent history
        if len(self._spread_history) > 1000:
            self._spread_history = self._spread_history[-500:]

    def predict(self, time_horizon_s: float = 5.0) -> float:
        """Predict spread after time_horizon_s."""
        if not self._spread_history:
            return self._last_spread_bps

        # Simple trend-based prediction
        if len(self._spread_history) < 10:
            return self._last_spread_bps

        recent = self._spread_history[-10:]
        trend = (recent[-1] - recent[0]) / len(recent)
        predicted = self._last_spread_bps + trend * time_horizon_s
        return max(0.1, predicted)

    @property
    def current_spread(self) -> float:
        return self._last_spread_bps

    @property
    def spread_volatility(self) -> float:
        if len(self._spread_history) < 10:
            return 0.0
        recent = self._spread_history[-50:]
        mean = sum(recent) / len(recent)
        variance = sum((x - mean) ** 2 for x in recent) / len(recent)
        return math.sqrt(variance)