engine/matcher/

stochastic.rs

use crate::matcher::Matcher;
use crate::types::{FillEvent, MarketState, Order, OrderState, Side};
use rand::prelude::*;
use rand::rng;
use rand::rngs::StdRng;

pub struct StochasticMatcher {
    orders: Vec<Order>,
    dt: f64,
    k: f64,
    a: f64,
    rng: StdRng,
    // Hawkes self-excitation of fill arrivals (optional; alpha=0 disables).
    // a_eff(t) = a + excitation(t), where excitation decays at rate beta
    // and jumps by alpha on each fill.
    hawkes_alpha: f64,
    hawkes_beta: f64,
    hawkes_excitation: f64,
    // Bilateral (2D) Hawkes: separate buy/sell intensities.
    // lambda_buy(t) = a + buy_excitation(t)
    // lambda_sell(t) = a + sell_excitation(t)
    // Each side only self-excites on fills of its own side.
    bilateral_hawkes: bool,
    hawkes_buy_excitation: f64,
    hawkes_sell_excitation: f64,
}

impl StochasticMatcher {
    pub fn new(dt: f64, k: f64, a: f64) -> Self {
        Self {
            orders: Vec::new(),
            dt,
            k,
            a,
            rng: StdRng::from_rng(&mut rng()),
            hawkes_alpha: 0.0,
            hawkes_beta: 0.0,
            hawkes_excitation: 0.0,
            bilateral_hawkes: false,
            hawkes_buy_excitation: 0.0,
            hawkes_sell_excitation: 0.0,
        }
    }

    /// Enable Hawkes self-excitation for fill arrivals.
    ///
    /// Each fill boosts the effective arrival rate by `alpha`.
    /// The boost decays exponentially at rate `beta` (per second).
    /// The branching ratio alpha/beta must be < 1 for stationarity.
    pub fn with_hawkes(mut self, alpha: f64, beta: f64) -> Self {
        self.hawkes_alpha = alpha;
        self.hawkes_beta = beta;
        self
    }

    /// Enable bilateral (2D) Hawkes: separate intensities for buy and sell arrivals.
    ///
    /// dλ+ = β(μ − λ+)dt + α dN+   (buy side, excited only by buy fills)
    /// dλ- = β(μ − λ-)dt + α dN-   (sell side, excited only by sell fills)
    ///
    /// where μ = `a` (the baseline rate), α = `alpha`, β = `beta`.
    /// Branching ratio alpha/beta must be < 1 for stationarity.
    pub fn with_bilateral_hawkes(mut self, alpha: f64, beta: f64) -> Self {
        self.hawkes_alpha = alpha;
        self.hawkes_beta = beta;
        self.bilateral_hawkes = true;
        self
    }

    /// Initialize unilateral Hawkes so effective intensity starts at `lambda0`.
    ///
    /// Keeps baseline `a` unchanged and adjusts only the excitation state.
    pub fn with_initial_hawkes_intensity(mut self, lambda0: f64) -> Self {
        self.hawkes_excitation = (lambda0 - self.a).max(0.0);
        self
    }

    /// Initialize bilateral Hawkes so buy/sell effective intensities start at
    /// (`lambda_buy0`, `lambda_sell0`).
    ///
    /// Keeps baseline `a` unchanged and adjusts only side-specific excitations.
    pub fn with_initial_bilateral_hawkes_intensities(
        mut self,
        lambda_buy0: f64,
        lambda_sell0: f64,
    ) -> Self {
        self.hawkes_buy_excitation = (lambda_buy0 - self.a).max(0.0);
        self.hawkes_sell_excitation = (lambda_sell0 - self.a).max(0.0);
        self
    }

    /// Seed this matcher with an explicit RNG (for reproducible, Send-safe copies).
    pub fn with_rng(mut self, rng: StdRng) -> Self {
        self.rng = rng;
        self
    }

    /// Current effective arrival rate (base + Hawkes excitation).
    #[inline]
    pub fn effective_a(&self) -> f64 {
        self.a + self.hawkes_excitation
    }

    /// Current Hawkes excitation level (0 when disabled or quiescent).
    #[inline]
    pub fn hawkes_excitation(&self) -> f64 {
        self.hawkes_excitation
    }

    /// Current effective buy-side arrival rate in bilateral mode.
    #[inline]
    pub fn effective_a_buy(&self) -> f64 {
        self.a + self.hawkes_buy_excitation
    }

    /// Current effective sell-side arrival rate in bilateral mode.
    #[inline]
    pub fn effective_a_sell(&self) -> f64 {
        self.a + self.hawkes_sell_excitation
    }
}

impl Matcher for StochasticMatcher {
    fn add_order(&mut self, order: Order) {
        self.orders.push(order);
    }

    fn cancel_order(&mut self, order_id: u64) {
        if let Some(order) = self.orders.iter_mut().find(|o| o.id == order_id) {
            order.state = OrderState::Cancelled;
        }
        self.orders.retain(|o| o.state == OrderState::Open);
    }

    fn cancel_all(&mut self) {
        self.orders.clear();
    }

    fn process_quote(&mut self, state: &MarketState) -> Vec<FillEvent> {
        // Hawkes: decay excitation over dt, compute effective arrival rate
        if self.bilateral_hawkes {
            let decay = (-self.hawkes_beta * self.dt).exp();
            self.hawkes_buy_excitation *= decay;
            self.hawkes_sell_excitation *= decay;
        } else if self.hawkes_alpha > 0.0 {
            self.hawkes_excitation *= (-self.hawkes_beta * self.dt).exp();
        }
        let a_eff = self.a + self.hawkes_excitation;

        let mut fills = Vec::new();
        // `state` is S_{t+1}: the market state AFTER the price process has stepped.
        // Orders were submitted by strategies that already observed this state,
        // so `mid_price` here is the post-move mid (the t+1 price).
        let mid_price = (state.best_bid + state.best_ask) / 2.0;
        let timestamp = state.timestamp;
        let best_ask = state.best_ask;
        let best_bid = state.best_bid;

        for order in self.orders.iter_mut() {
            if order.state != OrderState::Open {
                continue;
            }

            let mut filled = false;
            let mut fill_price = 0.0;

            // ----------------------------------------------------------------
            // Two-pronged evaluation
            // ----------------------------------------------------------------
            // A. AGGRESSIVE CROSSING
            //    The limit price is on the wrong side of the current BBO.
            //    A new buy order priced >= best_ask would immediately consume
            //    the resting offer; a new sell order priced <= best_bid would
            //    immediately consume the resting bid.
            //    Fill at the BBO price (price improvement over limit price).
            //
            // B. DETERMINISTIC SWEEP
            //    The limit was placed passively but the mid-price has since
            //    moved THROUGH it.  Buy @ L: mid <= L (market fell to the bid).
            //    Sell @ L: mid >= L (market rose to the ask).
            //    Because the price path crossed L, every market order that
            //    occurred during the interval would have filled this order.
            //    Result: guaranteed fill at the limit price.
            //
            // C. PROBABILISTIC POISSON ARRIVAL
            //    The mid has NOT yet crossed the limit.  Randomly arriving
            //    market orders may still hit the resting order during [t, t+1].
            //    We use the distance from the post-move mid (S_{t+1}) so the
            //    intensity reflects the current book depth, not a stale pre-move
            //    value.
            //        delta  = distance(limit, mid_{t+1})   [always >= 0 here]
            //        lambda = a * exp(-k * delta)           [arrivals / sec]
            //        p_fill = 1 - exp(-lambda * dt)
            // ----------------------------------------------------------------

            // A. Aggressive crossing -- fill at BBO with price improvement
            let aggressive = match order.side {
                Side::Buy => order.limit_price >= best_ask,
                Side::Sell => order.limit_price <= best_bid,
            };

            if aggressive {
                filled = true;
                fill_price = match order.side {
                    Side::Buy => best_ask,
                    Side::Sell => best_bid,
                };
            } else {
                // B. Price sweep -- mid moved through the limit
                let swept = match order.side {
                    Side::Buy => mid_price <= order.limit_price,
                    Side::Sell => mid_price >= order.limit_price,
                };

                if swept {
                    filled = true;
                    fill_price = order.limit_price;
                } else {
                    // C. Poisson arrival -- passive order not yet swept.
                    // Delta is positive: order is away from the post-move mid.
                    let delta = match order.side {
                        Side::Buy => mid_price - order.limit_price, // mid > limit (passive)
                        Side::Sell => order.limit_price - mid_price, // limit > mid (passive)
                    };

                    let side_a_eff = if self.bilateral_hawkes {
                        match order.side {
                            Side::Buy => self.a + self.hawkes_buy_excitation,
                            Side::Sell => self.a + self.hawkes_sell_excitation,
                        }
                    } else {
                        a_eff
                    };

                    let lambda = side_a_eff * (-self.k * delta).exp();
                    let prob = 1.0 - (-lambda * self.dt).exp();

                    if self.rng.random_bool(prob.clamp(0.0, 1.0)) {
                        filled = true;
                        fill_price = order.limit_price;
                    }
                }
            }

            if filled {
                order.state = OrderState::Filled;
                fills.push(FillEvent {
                    timestamp,
                    order_id: order.id,
                    side: order.side,
                    price: fill_price,
                    quantity: order.quantity,
                });
            }
        }

        // Hawkes: each fill excites future arrivals
        if self.bilateral_hawkes && !fills.is_empty() {
            for fill in &fills {
                match fill.side {
                    Side::Buy => self.hawkes_buy_excitation += self.hawkes_alpha,
                    Side::Sell => self.hawkes_sell_excitation += self.hawkes_alpha,
                }
            }
        } else if self.hawkes_alpha > 0.0 && !fills.is_empty() {
            self.hawkes_excitation += self.hawkes_alpha * fills.len() as f64;
        }

        // Clean up filled orders
        self.orders.retain(|o| o.state == OrderState::Open);

        fills
    }

    fn get_orders(&self) -> &Vec<Order> {
        &self.orders
    }

    fn augment_parameters(&self) -> std::collections::HashMap<String, f64> {
        let mut params = std::collections::HashMap::with_capacity(2);
        if self.bilateral_hawkes {
            // Convention (shared by all strategies and tests):
            //   hawkes_buy_intensity  = lambda+ = buy MO arrival rate
            //                        = rate of fills on MM's sell (ask) orders
            //   hawkes_sell_intensity = lambda- = sell MO arrival rate
            //                        = rate of fills on MM's buy (bid) orders
            //
            // hawkes_sell_excitation is excited by Side::Sell fills (ask fills = buy MOs hitting ask).
            // hawkes_buy_excitation  is excited by Side::Buy  fills (bid fills = sell MOs hitting bid).
            // Therefore: buy intensity = effective_a_sell(), sell intensity = effective_a_buy().
            params.insert("hawkes_buy_intensity".to_string(), self.effective_a_sell());
            params.insert("hawkes_sell_intensity".to_string(), self.effective_a_buy());
        } else if self.hawkes_alpha > 0.0 {
            params.insert("hawkes_intensity".to_string(), self.effective_a());
        }
        params
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::Side;

    fn state(bid: f64, ask: f64) -> MarketState {
        MarketState {
            timestamp: 0.0,
            best_bid: bid,
            best_ask: ask,
            last_price: None,
            true_volatility: None,
            true_drift: None,
            parameters: None,
        }
    }

    // --- Prong A: aggressive crossing ---

    #[test]
    fn test_aggressive_buy_fills_at_ask_not_at_limit() {
        // Buy limit above best_ask crosses the book (prong A).
        // Fill price must be best_ask, not the (higher) limit price.
        let mut m = StochasticMatcher::new(1.0, 1.5, 140.0);
        m.add_order(Order::new(1, Side::Buy, 101.0, 1.0));
        let fills = m.process_quote(&state(99.0, 100.0));
        assert_eq!(fills.len(), 1);
        assert_eq!(fills[0].price, 100.0); // best_ask, not 101.0
    }

    #[test]
    fn test_aggressive_sell_fills_at_bid_not_at_limit() {
        // Sell limit below best_bid crosses the book (prong A).
        // Fill price must be best_bid, not the (lower) limit price.
        let mut m = StochasticMatcher::new(1.0, 1.5, 140.0);
        m.add_order(Order::new(1, Side::Sell, 98.0, 1.0));
        let fills = m.process_quote(&state(99.0, 100.0));
        assert_eq!(fills.len(), 1);
        assert_eq!(fills[0].price, 99.0); // best_bid, not 98.0
    }

    // --- Prong B: deterministic sweep ---

    #[test]
    fn test_sweep_buy_fills_when_mid_drops_through_limit() {
        // Passive buy at 99.8; mid (99.0) < limit -> the market swept
        // through the order (prong B). Fill is guaranteed at limit price.
        let mut m = StochasticMatcher::new(1.0, 1.5, 140.0);
        m.add_order(Order::new(1, Side::Buy, 99.8, 1.0));
        // bid=98.5, ask=100.0 -> mid=99.25 < 99.8 -> swept
        let fills = m.process_quote(&state(98.5, 100.0));
        assert_eq!(fills.len(), 1);
        assert_eq!(fills[0].price, 99.8);
    }

    #[test]
    fn test_sweep_sell_fills_when_mid_rises_through_limit() {
        // Passive sell at 99.5; mid (100.0) >= limit -> swept (prong B).
        // Fill is guaranteed at the limit price.
        let mut m = StochasticMatcher::new(1.0, 1.5, 140.0);
        m.add_order(Order::new(1, Side::Sell, 99.5, 1.0));
        // bid=99.0, ask=101.0 -> mid=100.0 >= 99.5 -> swept
        let fills = m.process_quote(&state(99.0, 101.0));
        assert_eq!(fills.len(), 1);
        assert_eq!(fills[0].price, 99.5);
    }

    // --- Prong C: Poisson arrival ---

    #[test]
    fn test_passive_far_from_mid_has_near_zero_fill_probability() {
        // delta=10, a=1 -> lambda=1*exp(-15)~3e-7 -> p_fill~3e-7 per tick.
        // Over 1000 independent draws the order should fill at most a handful of times.
        let a = 1.0;
        let k = 1.5;
        let mut fill_count = 0u32;
        let trials = 1000;
        for i in 0..trials {
            let mut m = StochasticMatcher::new(1.0, k, a);
            m.add_order(Order::new(i as u64 + 1, Side::Buy, 80.0, 1.0));
            // mid=90.0, limit=80.0, delta=10
            if !m.process_quote(&state(89.5, 90.5)).is_empty() {
                fill_count += 1;
            }
        }
        assert!(
            fill_count <= 10,
            "Expected near-zero fills, got {fill_count}/{trials}"
        );
    }

    #[test]
    fn test_passive_at_zero_delta_has_near_certain_fill() {
        // delta=0 -> lambda=a=500 -> p = 1 - exp(-500) ~= 1.0.
        // Place the buy limit exactly at mid (passive: limit < ask).
        let a = 500.0;
        let k = 1.5;
        let mut miss_count = 0u32;
        let trials = 100;
        for i in 0..trials {
            let mut m = StochasticMatcher::new(1.0, k, a);
            // bid=99.99, ask=100.01, mid=100.0, limit=100.0
            m.add_order(Order::new(i as u64 + 1, Side::Buy, 100.0, 1.0));
            if m.process_quote(&state(99.99, 100.01)).is_empty() {
                miss_count += 1;
            }
        }
        assert!(
            miss_count == 0,
            "Expected all fills, missed {miss_count}/{trials}"
        );
    }

    // --- Multiple orders from a single agent ---

    #[test]
    fn test_multiple_passive_bids_only_swept_one_fills() {
        // Three bids at different depths; only the one the mid sweeps through
        // (prong B) fills deterministically. The others are so far away (with
        // a=0.001) that Poisson fills are negligible.
        //
        // bid=98.9, ask=99.1 -> mid=99.0
        // order 1: limit=99.05 -> mid(99.0) <= 99.05 -> SWEPT -> guaranteed fill
        // order 2: limit=98.0  -> delta=1.0, lambda=0.001*exp(-1.5)~0.00022 -> p~0
        // order 3: limit=95.0  -> delta=4.0, lambda~3e-10 -> p~0
        let mut m = StochasticMatcher::new(1.0, 1.5, 0.001);
        m.add_order(Order::new(1, Side::Buy, 99.05, 1.0));
        m.add_order(Order::new(2, Side::Buy, 98.0, 1.0));
        m.add_order(Order::new(3, Side::Buy, 95.0, 1.0));
        let fills = m.process_quote(&state(98.9, 99.1));
        let ids: Vec<u64> = fills.iter().map(|f| f.order_id).collect();
        assert!(ids.contains(&1), "Swept order must fill");
        assert!(
            !ids.contains(&2),
            "Passive order at delta=1 must not fill with a=0.001"
        );
        assert!(
            !ids.contains(&3),
            "Passive order at delta=4 must not fill with a=0.001"
        );
    }

    #[test]
    fn test_two_aggressive_bids_both_fill_at_bbo() {
        // Two bids both above best_ask must fill independently, both at best_ask.
        let mut m = StochasticMatcher::new(1.0, 1.5, 140.0);
        m.add_order(Order::new(1, Side::Buy, 101.0, 1.0));
        m.add_order(Order::new(2, Side::Buy, 102.0, 2.0));
        let fills = m.process_quote(&state(99.0, 100.0));
        assert_eq!(fills.len(), 2);
        for f in &fills {
            assert_eq!(f.price, 100.0);
        }
    }

    // --- Hawkes self-excitation ---

    #[test]
    fn test_hawkes_disabled_by_default() {
        let m = StochasticMatcher::new(1.0, 1.5, 10.0);
        assert_eq!(m.hawkes_excitation(), 0.0);
        assert_eq!(m.effective_a(), 10.0);
    }

    #[test]
    fn test_hawkes_builder_sets_params() {
        let m = StochasticMatcher::new(1.0, 1.5, 10.0).with_hawkes(0.5, 2.0);
        assert_eq!(m.hawkes_excitation(), 0.0);
        assert_eq!(m.effective_a(), 10.0); // no fills yet
    }

    #[test]
    fn test_hawkes_excitation_after_fills() {
        // After a fill, excitation should increase by alpha.
        let mut m = StochasticMatcher::new(1.0, 1.5, 500.0).with_hawkes(3.0, 2.0);

        // Place a bid at the ask to guarantee fill (aggressive crossing).
        m.add_order(Order::new(1, Side::Buy, 101.0, 1.0));
        let fills = m.process_quote(&state(99.0, 100.0));
        assert_eq!(fills.len(), 1);
        // Excitation should be alpha * 1 fill = 3.0
        assert!((m.hawkes_excitation() - 3.0).abs() < 1e-10);
        assert!((m.effective_a() - 503.0).abs() < 1e-10);
    }

    #[test]
    fn test_hawkes_excitation_decays() {
        let mut m = StochasticMatcher::new(1.0, 1.5, 500.0).with_hawkes(3.0, 2.0);

        // Trigger a fill to build excitation.
        m.add_order(Order::new(1, Side::Buy, 101.0, 1.0));
        m.process_quote(&state(99.0, 100.0));
        assert!((m.hawkes_excitation() - 3.0).abs() < 1e-10);

        // Next tick (no new fills): excitation should decay by exp(-beta*dt).
        m.process_quote(&state(99.0, 100.0)); // no orders -> no fills
        let expected = 3.0 * (-2.0_f64).exp(); // exp(-2.0 * 1.0)
        assert!((m.hawkes_excitation() - expected).abs() < 1e-10);
    }

    #[test]
    fn test_hawkes_increases_fill_rate_statistically() {
        // Compare fill counts: Hawkes (with excitation primed) should fill
        // more often than base matcher at the same spread distance.
        //
        // We use a moderate base rate where fills are not saturated.
        // alpha=10, beta=5 gives meaningful excitation after each fill.
        //
        // Retry up to 3 times — the test is inherently stochastic and
        // 2000 trials may occasionally not show a clear signal.
        const MAX_ATTEMPTS: usize = 3;
        let mut best_margin: i32 = i32::MIN;
        let mut best_base = 0u32;
        let mut best_hawkes = 0u32;

        for attempt in 0..MAX_ATTEMPTS {
            let trials = 2000;
            let mut base_fills = 0u32;
            let mut hawkes_fills = 0u32;

            // First run base paths
            for i in 0..trials {
                let mut m = StochasticMatcher::new(1.0, 1.5, 1.0);
                m.add_order(Order::new(i as u64 + 1, Side::Buy, 99.5, 1.0));
                if !m.process_quote(&state(99.9, 100.1)).is_empty() {
                    base_fills += 1;
                }
            }

            // Then run Hawkes paths: prime excitation to simulate post-burst state
            for i in 0..trials {
                let mut m = StochasticMatcher::new(1.0, 1.5, 1.0).with_hawkes(10.0, 5.0);
                // Prime: inject excitation as if 2 fills just happened
                m.add_order(Order::new(1000000 + i as u64, Side::Buy, 101.0, 1.0));
                m.process_quote(&state(99.0, 100.0)); // aggressive fill -> excitation += 10
                m.add_order(Order::new(2000000 + i as u64, Side::Buy, 101.0, 1.0));
                m.process_quote(&state(99.0, 100.0)); // another fill -> excitation += 10

                // Now place a passive order: higher effective_a should give more fills
                m.add_order(Order::new(3000000 + i as u64, Side::Buy, 99.5, 1.0));
                if !m.process_quote(&state(99.9, 100.1)).is_empty() {
                    hawkes_fills += 1;
                }
            }

            let margin = hawkes_fills as i32 - base_fills as i32;
            if margin > best_margin {
                best_margin = margin;
                best_base = base_fills;
                best_hawkes = hawkes_fills;
            }

            if hawkes_fills > base_fills {
                return; // success
            }

            eprintln!(
                "attempt {} failed: base={}, hawkes={}; retrying…",
                attempt + 1,
                base_fills,
                hawkes_fills,
            );
        }

        panic!(
            "Hawkes should produce more fills after {} attempts. \
             Best result: base={}, hawkes={} (margin={})",
            MAX_ATTEMPTS, best_base, best_hawkes, best_margin,
        );
    }
}