solver/models/

bilateral_hawkes_order_flow_imbalance.rs

use super::control::{ControlProblem, StateDerivatives};
use super::market_making::MarketMakingControl;
use super::normal_cdf;
use crate::numeric::finite_difference::discretization::{DimensionKind, Transport};
use crate::numeric::finite_difference::pde::PdeProblem;

// Re-export so existing paths compile.
pub use super::TerminalCondition;

/// Bilateral Hawkes model with price-impact-driven adverse selection.
///
/// # State Space: `[q, lambda_plus, lambda_minus]`
///
/// | dim | symbol       | description                                      |
/// |-----|--------------|--------------------------------------------------|
/// | 0   | q            | Inventory level (integer, jump-controlled)       |
/// | 1   | lambda_plus  | Buy market-order intensity (ODE drift, no noise) |
/// | 2   | lambda_minus | Sell market-order intensity (ODE drift, no noise)|
///
/// # Dynamics
///
/// `d(lambda+) = beta*(mu - lambda+) dt + alpha dN+`
/// `d(lambda-) = beta*(mu - lambda-) dt + alpha dN-`
/// `dS = sigma dW + eta_ofi (dN+ - dN-)`
///
/// Each buy market order (dN+) moves the mid price UP by `eta_ofi`; each sell market
/// order (dN-) moves it DOWN by `eta_ofi`. The market maker accounts for this
/// price impact when setting optimal spreads.
///
/// # Optimal Spreads Under Price Impact
///
/// Deriving the FOC from the CARA HJB with per-fill wealth effect `(q±1)*eta_ofi` gives:
///
/// $$ \delta_{bid}^* = \delta_{base} - \frac{\partial V}{\partial q} + (q+1)\eta_{OFI} $$
/// $$ \delta_{ask}^* = \delta_{base} + \frac{\partial V}{\partial q} - (q-1)\eta_{OFI} $$
///
/// Economic interpretation:
/// - Long inventory (q > 0): ask narrows (selling into a rising price is favourable),
///   bid widens (avoid accumulating more when sell MOs push price down).
/// - Short inventory (q < 0): ask widens (avoid shorting further into a rising price),
///   bid narrows (cover cheaply when sell MOs push price down).
///
/// Setting `eta_ofi = 0` recovers the plain `BilateralHawkes` (BHK) model exactly.
///
/// In the simulation pass `impact_factor = -eta_ofi` to `SimulatedDataSource` so the
/// price process matches the assumed dynamics `dS = sigma dW + eta_ofi (dN+ - dN-)`.
///
/// # Hawkes Approximation
///
/// Same fluid (mean-field) approximation as `BilateralHawkes`:
///
/// $$ d(\lambda_+)/dt \approx \beta(\mu - \lambda_+) + \alpha \lambda_+ $$
/// $$ d(\lambda_-)/dt \approx \beta(\mu - \lambda_-) + \alpha \lambda_- $$
///
#[derive(Clone, Debug)]
pub struct BilateralHawkesOrderFlowImbalance {
    /// Risk-aversion parameter.
    pub gamma: f64,
    /// Price volatility.
    pub sigma: f64,
    /// Fill-rate decay (Avellaneda-Stoikov kappa).
    pub kappa: f64,
    // --- Hawkes parameters (shared for both sides) ---
    /// Jump size added to the excited side's intensity.
    pub alpha: f64,
    /// Mean-reversion speed.
    pub beta: f64,
    /// Baseline intensity (mean-reversion target).
    pub mu: f64,
    // --- Price-impact parameter ---
    /// Price impact per fill event (eta_ofi in dS = sigma*dW + eta_ofi*(dN+ - dN-)).
    /// Higher eta_ofi = stronger q-dependent spread correction.
    /// Setting eta_ofi = 0 recovers the plain BilateralHawkes (BHK) model.
    pub eta_ofi: f64,
    /// Grid step for both lambda+ and lambda- dimensions.
    pub lambda_step: f64,
    /// Grid step for the inventory dimension.
    pub dq: f64,
    /// Minimum inventory (lower hard boundary). Defaults to -infinity.
    pub q_min: f64,
    /// Maximum inventory (upper hard boundary). Defaults to +infinity.
    pub q_max: f64,
    /// Terminal condition for the value function at T. Defaults to Zero.
    pub terminal_condition: TerminalCondition,
    /// Optional terminal liquidation half-spread used when terminal_condition is
    /// LiquidationCost. If None, uses the AS base spread.
    pub terminal_liquidation_half_spread: Option<f64>,
}

impl BilateralHawkesOrderFlowImbalance {
    pub fn new(
        gamma: f64,
        sigma: f64,
        kappa: f64,
        alpha: f64,
        beta: f64,
        mu: f64,
        eta_ofi: f64,
    ) -> Self {
        Self {
            gamma,
            sigma,
            kappa,
            alpha,
            beta,
            mu,
            eta_ofi,
            lambda_step: 1.0,
            dq: 1.0,
            q_min: f64::NEG_INFINITY,
            q_max: f64::INFINITY,
            terminal_condition: TerminalCondition::Zero,
            terminal_liquidation_half_spread: None,
        }
    }

    pub fn with_terminal_condition(mut self, terminal_condition: TerminalCondition) -> Self {
        self.terminal_condition = terminal_condition;
        self
    }

    pub fn with_terminal_liquidation_half_spread(mut self, half_spread: f64) -> Self {
        self.terminal_liquidation_half_spread = Some(half_spread.max(0.0));
        self
    }

    pub fn with_inventory_bounds(mut self, q_min: f64, q_max: f64) -> Self {
        self.q_min = q_min;
        self.q_max = q_max;
        self
    }

    pub fn with_lambda_step(mut self, lambda_step: f64) -> Self {
        self.lambda_step = lambda_step.abs().max(1e-8);
        self
    }

    pub fn with_dq(mut self, dq: f64) -> Self {
        self.dq = dq.abs().max(1e-8);
        self
    }

    /// Computes optimal bid/ask half-spreads with price-impact-driven inventory correction.
    ///
    /// When eta_ofi > 0 each fill carries an adverse-selection wealth effect `(q±1)*eta_ofi`.
    /// The FOC on the CARA HJB gives:
    ///   delta_bid = base_spread - dV/dq_fwd + (q+1)*eta_ofi
    ///   delta_ask = base_spread + dV/dq_bwd - (q-1)*eta_ofi
    ///
    /// Long inventory (q > 0): ask narrows (sell into rising price), bid widens.
    /// Short inventory (q < 0): ask widens (avoid adding short), bid narrows.
    /// eta_ofi = 0 recovers the plain BHK spread formula identically.
    pub fn get_spreads(&self, state: &[f64; 3], derivs: &StateDerivatives<3>) -> (f64, f64) {
        let q = state[0];
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).ln();
        let delta_bid = (base_spread - derivs.fwd[0] + (q + 1.0) * self.eta_ofi).clamp(-10.0, 10.0);
        let delta_ask = (base_spread + derivs.bwd[0] - (q - 1.0) * self.eta_ofi).clamp(-10.0, 10.0);
        (delta_bid, delta_ask)
    }

    /// The base arrival intensity is state-dependent: returns the average of
    /// the buy/sell intensity states for intensity-to-spread conversion.
    pub fn fill_rate_base(&self, state: &[f64; 3]) -> f64 {
        ((state[1] + state[2]) * 0.5).max(1e-10)
    }

    /// The fill-rate decay parameter used for intensity-to-spread conversion.
    pub fn fill_rate_decay(&self) -> f64 {
        self.kappa
    }
}

impl ControlProblem<3> for BilateralHawkesOrderFlowImbalance {
    type Control = MarketMakingControl;

    fn optimize(&self, _t: f64, state: &[f64; 3], derivs: &StateDerivatives<3>) -> Self::Control {
        let q = state[0];
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);

        let (d_bid, d_ask) = self.get_spreads(state, derivs);

        let max_mult = 20.0_f64;
        let lambda_bid_fill = lambda_minus * (-self.kappa * d_bid).exp().min(max_mult);
        let lambda_ask_fill = lambda_plus * (-self.kappa * d_ask).exp().min(max_mult);

        let lambda_bid_fill = if q >= self.q_max {
            0.0
        } else {
            lambda_bid_fill
        };
        let lambda_ask_fill = if q <= self.q_min {
            0.0
        } else {
            lambda_ask_fill
        };

        MarketMakingControl::new(lambda_bid_fill, lambda_ask_fill)
    }

    fn running_reward(&self, _t: f64, _state: &[f64; 3], control: &Self::Control) -> f64 {
        (control.bid_intensity + control.ask_intensity) / (self.gamma + self.kappa)
    }

    fn bsde_driver(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        _derivs: &StateDerivatives<3>,
        dt: f64,
    ) -> f64 {
        // Full-value problem: the two Hawkes intensity drifts are already
        // simulated forward, so the backward driver adds only the running
        // reward plus the local inventory-risk source. The reward is bounded at
        // `1/dt` per fill side to match the forward fill-probability clamp.
        let q = state[0];
        let local = -0.5 * self.gamma * self.sigma.powi(2) * q.powi(2);
        let rate_cap = 1.0 / dt.max(1e-12);
        let reward = (control.bid_intensity.min(rate_cap) + control.ask_intensity.min(rate_cap))
            / (self.gamma + self.kappa);
        reward + local
    }

    fn generator(
        &self,
        _t: f64,
        state: &[f64; 3],
        _control: &Self::Control,
        derivs: &StateDerivatives<3>,
    ) -> f64 {
        let q = state[0];
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let net_drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let net_drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;

        -0.5 * self.gamma * self.sigma.powi(2) * q.powi(2)
            + net_drift_plus * derivs.grad[1]
            + net_drift_minus * derivs.grad[2]
    }

    fn terminal(&self, state: &[f64; 3]) -> f64 {
        match self.terminal_condition {
            TerminalCondition::Zero => 0.0,
            TerminalCondition::LiquidationCost => {
                let q = state[0];
                let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).ln();
                let half_spread = self.terminal_liquidation_half_spread.unwrap_or(base_spread);
                -q.abs() * half_spread
            }
        }
    }

    fn discount_rate(&self, _state: &[f64; 3]) -> f64 {
        0.0
    }

    fn constant_discount_rate(&self) -> Option<f64> {
        Some(0.0)
    }

    fn next_step(&self, _t: f64, state: &[f64; 3], dt: f64, noise: &[f64; 3]) -> [f64; 3] {
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;
        let mut next = [
            state[0],
            (lambda_plus + drift_plus * dt).max(0.0),
            (lambda_minus + drift_minus * dt).max(0.0),
        ];

        let u = normal_cdf(noise[0]);
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).ln();
        let lambda_bid_fill = lambda_minus * (-self.kappa * base_spread).exp();
        let lambda_ask_fill = lambda_plus * (-self.kappa * base_spread).exp();
        let p_bid = (lambda_bid_fill * dt).clamp(0.0, 1.0);
        let p_ask = (lambda_ask_fill * dt).clamp(0.0, 1.0);
        if u < p_bid {
            next[0] += 1.0;
            next[2] += self.alpha;
        } else if u > 1.0 - p_ask {
            next[0] -= 1.0;
            next[1] += self.alpha;
        }

        next[1] = next[1].max(0.0);
        next[2] = next[2].max(0.0);
        next
    }

    fn next_step_controlled(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        dt: f64,
        noise: &[f64; 3],
    ) -> [f64; 3] {
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;
        let mut next = [
            state[0],
            (lambda_plus + drift_plus * dt).max(0.0),
            (lambda_minus + drift_minus * dt).max(0.0),
        ];

        // The optimal control already carries the fill intensities, so the
        // Bernoulli fill events use them instead of a frozen proxy quote.
        let u = normal_cdf(noise[0]);
        let p_bid = (control.bid_intensity * dt).clamp(0.0, 1.0);
        let p_ask = (control.ask_intensity * dt).clamp(0.0, 1.0);
        if u < p_bid {
            next[0] += 1.0;
            next[2] += self.alpha;
        } else if u > 1.0 - p_ask {
            next[0] -= 1.0;
            next[1] += self.alpha;
        }

        next[1] = next[1].max(0.0);
        next[2] = next[2].max(0.0);
        next
    }

    fn is_diffusion_dimension(&self, _dim: usize) -> bool {
        false
    }

    fn gradient_step(&self, dim: usize) -> f64 {
        match dim {
            1 | 2 => self.lambda_step.abs().max(1e-8),
            _ => 1.0,
        }
    }
}

impl PdeProblem<3> for BilateralHawkesOrderFlowImbalance {
    fn dimension_kind(&self, dim: usize) -> DimensionKind {
        match dim {
            0 => DimensionKind::DiscreteJump,
            _ => DimensionKind::DeterministicDrift,
        }
    }

    fn transport(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        derivs: &StateDerivatives<3>,
    ) -> Transport<3> {
        let q = state[0];
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let net_drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let net_drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;
        let step = self.lambda_step.abs().max(1e-12);

        let rate_plus = net_drift_plus.abs() / step;
        let rate_minus = net_drift_minus.abs() / step;
        let (lp_plus, lp_minus) = if net_drift_plus >= 0.0 {
            (rate_plus, 0.0)
        } else {
            (0.0, rate_plus)
        };
        let (lm_plus, lm_minus) = if net_drift_minus >= 0.0 {
            (rate_minus, 0.0)
        } else {
            (0.0, rate_minus)
        };

        let hamiltonian =
            (control.bid_intensity + control.ask_intensity) / (self.gamma + self.kappa);
        let jump_transport =
            control.bid_intensity * derivs.fwd[0] - control.ask_intensity * derivs.bwd[0];
        let local = -0.5 * self.gamma * self.sigma.powi(2) * q.powi(2);

        Transport::new(
            [control.bid_intensity, lp_plus, lm_plus],
            [control.ask_intensity, lp_minus, lm_minus],
            hamiltonian + local - jump_transport,
        )
    }
}

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

    fn zero_derivs() -> StateDerivatives<3> {
        StateDerivatives::with_directional([0.0; 3], [0.0; 3], [0.0; 3], [0.0; 3])
    }

    fn derivs_with(fwd0: f64, bwd0: f64) -> StateDerivatives<3> {
        let mut d = zero_derivs();
        d.fwd[0] = fwd0;
        d.bwd[0] = bwd0;
        d
    }

    fn zero_state() -> [f64; 3] {
        [0.0, 1.0, 1.0]
    }

    fn default_model() -> BilateralHawkesOrderFlowImbalance {
        BilateralHawkesOrderFlowImbalance::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0, 0.5)
            .with_lambda_step(0.1)
            .with_dq(1.0)
    }

    #[test]
    fn model_creation() {
        let m = default_model();
        assert_eq!(m.gamma, 0.5);
        assert_eq!(m.sigma, 0.2);
        assert_eq!(m.kappa, 1.5);
        assert_eq!(m.alpha, 0.5);
        assert_eq!(m.beta, 2.0);
        assert_eq!(m.mu, 1.0);
        assert_eq!(m.eta_ofi, 0.5);
    }

    #[test]
    fn bsde_driver_excludes_intensity_transport() {
        let m = default_model();
        let state = [2.0, 1.0, 1.0];
        let control = ControlProblem::optimize(&m, 0.0, &state, &zero_derivs());

        let mut derivs = zero_derivs();
        derivs.grad[1] = 1.0;
        derivs.grad[2] = 1.0;

        let dt = 0.01;
        let q = state[0];
        let local = -0.5 * m.gamma * m.sigma.powi(2) * q.powi(2);
        let rate_cap = 1.0 / dt;
        let reward = (control.bid_intensity.min(rate_cap) + control.ask_intensity.min(rate_cap))
            / (m.gamma + m.kappa);
        let expected = reward + local;

        assert_eq!(
            ControlProblem::bsde_driver(&m, 0.0, &state, &control, &derivs, dt),
            expected
        );
        assert!(
            ControlProblem::bsde_driver(&m, 0.0, &state, &control, &derivs, dt)
                != ControlProblem::driver(&m, 0.0, &state, &control, &derivs),
            "full driver adds intensity transport, bsde driver must not"
        );
    }

    #[test]
    fn next_step_controlled_uses_optimal_intensities() {
        let m = default_model();
        let state = [1.0, 1.0, 1.0];

        let ctrl = MarketMakingControl::new(0.0, 0.0);
        let next =
            ControlProblem::next_step_controlled(&m, 0.0, &state, &ctrl, 0.01, &[0.0, 0.0, 0.0]);
        assert_eq!(next[0], 1.0);

        let ctrl = MarketMakingControl::new(1000.0, 0.0);
        let next =
            ControlProblem::next_step_controlled(&m, 0.0, &state, &ctrl, 0.01, &[0.0, 0.0, 0.0]);
        assert_eq!(next[0], 2.0);
    }

    #[test]
    fn dimension_kinds_match_state_space() {
        let m = default_model();
        assert_eq!(
            PdeProblem::dimension_kind(&m, 0),
            DimensionKind::DiscreteJump
        );
        assert_eq!(
            PdeProblem::dimension_kind(&m, 1),
            DimensionKind::DeterministicDrift
        );
        assert_eq!(
            PdeProblem::dimension_kind(&m, 2),
            DimensionKind::DeterministicDrift
        );
    }

    #[test]
    fn zero_inventory_gives_symmetric_spreads() {
        let m = default_model();
        let state = [0.0, 1.0, 1.0];
        let (bid, ask) = m.get_spreads(&state, &zero_derivs());
        let base = (1.0 / m.gamma) * (1.0 + m.gamma / m.kappa).ln();
        assert!(
            (bid - ask).abs() < 1e-12,
            "bid and ask should be equal at q=0"
        );
        assert!((bid - (base + m.eta_ofi)).abs() < 1e-12);
    }

    #[test]
    fn long_inventory_widens_bid_narrows_ask() {
        let m = default_model(); // xi = 0.5
        let q = 5.0;
        let state_long = [q, 1.0, 1.0];
        let state_zero = [0.0, 1.0, 1.0];
        let (bid_long, ask_long) = m.get_spreads(&state_long, &zero_derivs());
        let (bid_zero, ask_zero) = m.get_spreads(&state_zero, &zero_derivs());
        assert!(bid_long > bid_zero, "long position widens bid");
        assert!(ask_long < ask_zero, "long position narrows ask");
    }

    #[test]
    fn short_inventory_narrows_bid_widens_ask() {
        let m = default_model(); // xi = 0.5
        let q = -5.0;
        let state_short = [q, 1.0, 1.0];
        let state_zero = [0.0, 1.0, 1.0];
        let (bid_short, ask_short) = m.get_spreads(&state_short, &zero_derivs());
        let (bid_zero, ask_zero) = m.get_spreads(&state_zero, &zero_derivs());
        assert!(bid_short < bid_zero, "short position narrows bid");
        assert!(ask_short > ask_zero, "short position widens ask");
    }

    #[test]
    fn larger_xi_gives_larger_spread_corrections() {
        let m1 = BilateralHawkesOrderFlowImbalance::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0, 0.5)
            .with_lambda_step(0.1);
        let m2 = BilateralHawkesOrderFlowImbalance::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0, 1.0)
            .with_lambda_step(0.1);
        let m0 = BilateralHawkesOrderFlowImbalance::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0, 0.0)
            .with_lambda_step(0.1);

        let state = [5.0, 1.0, 1.0]; // q=5 (long)
        let derivs = zero_derivs();

        let (bid0, ask0) = m0.get_spreads(&state, &derivs);
        let (bid1, ask1) = m1.get_spreads(&state, &derivs);
        let (bid2, ask2) = m2.get_spreads(&state, &derivs);

        assert!(bid1 > bid0);
        assert!(bid2 > bid1);
        assert!(ask1 < ask0);
        assert!(ask2 < ask1);
    }

    #[test]
    fn terminal_condition_zero() {
        let m = default_model().with_terminal_condition(TerminalCondition::Zero);
        let state = [5.0, 1.0, 1.0];
        assert_eq!(ControlProblem::terminal(&m, &state), 0.0);
    }

    #[test]
    fn terminal_condition_liquidation_cost() {
        let m = default_model()
            .with_terminal_condition(TerminalCondition::LiquidationCost)
            .with_terminal_liquidation_half_spread(0.1);

        let state_q5 = [5.0, 1.0, 1.0];
        let state_q0 = [0.0, 1.0, 1.0];
        let state_qm3 = [-3.0, 1.0, 1.0];

        assert_eq!(ControlProblem::terminal(&m, &state_q5), -5.0 * 0.1);
        assert_eq!(ControlProblem::terminal(&m, &state_q0), 0.0);
        assert_eq!(ControlProblem::terminal(&m, &state_qm3), -3.0 * 0.1);
    }

    #[test]
    fn optimize_output_nonnegative() {
        let m = default_model();
        let state = zero_state();
        let derivs = zero_derivs();
        let output = ControlProblem::optimize(&m, 0.0, &state, &derivs);

        assert!(output.bid_intensity >= 0.0);
        assert!(output.ask_intensity >= 0.0);
    }

    #[test]
    fn inventory_bounds_enforcement() {
        let m = default_model().with_inventory_bounds(-2.0, 2.0);
        let state = [2.0, 2.0, 1.0]; // at upper bound
        let derivs = derivs_with(0.0, -1.0); // negative gradient encourages sells

        let output = ControlProblem::optimize(&m, 0.0, &state, &derivs);
        // At upper bound, can't increase inventory further: no bid fill.
        assert_eq!(output.bid_intensity, 0.0);
    }

    #[test]
    fn hawkes_dynamics_stability() {
        let m = default_model();

        let rho = m.alpha / m.beta;
        let stationary_intensity = m.mu / (1.0 - rho);

        let mut state = [0.0, stationary_intensity, stationary_intensity];

        for _ in 0..1000 {
            state = ControlProblem::next_step(&m, 0.0, &state, 0.01, &[0.0, 0.0, 0.0]);
        }

        assert!(state[1] > 0.0 && state[1] < 10.0 * stationary_intensity);
        assert!(state[2] > 0.0 && state[2] < 10.0 * stationary_intensity);
    }

    #[test]
    fn next_step_consistency() {
        let m = default_model();
        let state = [0.0, 1.5, 1.0];
        let dt = 0.01;

        let next = ControlProblem::next_step(&m, 0.0, &state, dt, &[0.0; 3]);

        assert_eq!(next[0], state[0]);

        let drift_plus = m.beta * (m.mu - state[1]) + m.alpha * state[1];
        let expected_lambda_plus = state[1] + drift_plus * dt;
        assert!((next[1] - expected_lambda_plus).abs() < 1e-10);
    }

    #[test]
    fn next_step_can_fill_inventory() {
        let m = default_model();
        let state = [0.0, 2.0, 2.0];
        let dt = 0.01;

        // Very negative noise -> uniform(0,1) near 0 -> bid fill -> q increases.
        let up = ControlProblem::next_step(&m, 0.0, &state, dt, &[-10.0, 0.0, 0.0]);
        // Very positive noise -> uniform near 1 -> ask fill -> q decreases.
        let down = ControlProblem::next_step(&m, 0.0, &state, dt, &[10.0, 0.0, 0.0]);

        assert!(up[0] > state[0], "bid fill should increase inventory");
        assert!(down[0] < state[0], "ask fill should decrease inventory");
    }

    #[test]
    fn gradient_step_returns_correct_values() {
        let m = default_model();

        assert_eq!(ControlProblem::gradient_step(&m, 0), 1.0);
        assert_eq!(ControlProblem::gradient_step(&m, 1), 0.1);
        assert_eq!(ControlProblem::gradient_step(&m, 2), 0.1);
    }

    #[test]
    fn is_diffusion_dimension() {
        let m = default_model();

        assert!(!ControlProblem::is_diffusion_dimension(&m, 0));
        assert!(!ControlProblem::is_diffusion_dimension(&m, 1));
        assert!(!ControlProblem::is_diffusion_dimension(&m, 2));
    }

    #[test]
    fn constant_discount_rate() {
        let m = default_model();
        assert_eq!(ControlProblem::constant_discount_rate(&m), Some(0.0));
    }

    #[test]
    fn spread_formula_matches_analytical_expression() {
        let m = default_model(); // xi = 0.5
        let base = (1.0 / m.gamma) * (1.0 + m.gamma / m.kappa).ln();
        for &q in &[-5.0_f64, -1.0, 0.0, 1.0, 5.0] {
            let state = [q, 1.0, 1.0];
            let (bid, ask) = m.get_spreads(&state, &zero_derivs());
            let expected_bid = base + (q + 1.0) * m.eta_ofi;
            let expected_ask = base - (q - 1.0) * m.eta_ofi;
            assert!((bid - expected_bid).abs() < 1e-12, "bid mismatch at q={q}");
            assert!((ask - expected_ask).abs() < 1e-12, "ask mismatch at q={q}");
        }
    }

    #[test]
    fn compare_with_balanced_no_imbalance() {
        let m = default_model();
        let state = [0.0, 1.0, 1.0];

        let (bid, ask) = m.get_spreads(&state, &zero_derivs());
        assert!((bid - ask).abs() < 1e-12);
    }

    #[test]
    fn extreme_imbalance_clamping() {
        let m = BilateralHawkesOrderFlowImbalance::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0, 100.0)
            .with_lambda_step(0.1);

        let state = [0.0, 10.0, 0.0]; // extreme imbalance
        let derivs = derivs_with(100.0, 100.0);

        let (bid, ask) = m.get_spreads(&state, &derivs);

        assert!((-10.0..=10.0).contains(&bid));
        assert!((-10.0..=10.0).contains(&ask));
    }

    #[test]
    fn inventory_control_respects_bounds() {
        let m = default_model().with_inventory_bounds(-5.0, 5.0);
        let state_at_max = [5.0, 2.0, 1.0];
        let state_at_min = [-5.0, 1.0, 2.0];
        let derivs = zero_derivs();

        let output_max = ControlProblem::optimize(&m, 0.0, &state_at_max, &derivs);
        let output_min = ControlProblem::optimize(&m, 0.0, &state_at_min, &derivs);

        // At max inventory, cannot increase: bid fill must be zero.
        assert_eq!(output_max.bid_intensity, 0.0);

        // At min inventory, cannot decrease: ask fill must be zero.
        assert_eq!(output_min.ask_intensity, 0.0);
    }

    #[test]
    fn xi_zero_recovers_baseline() {
        let m = BilateralHawkesOrderFlowImbalance::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0, 0.0)
            .with_lambda_step(0.1);
        let state = [0.0, 2.0, 1.0];
        let (bid, ask) = m.get_spreads(&state, &zero_derivs());
        let base = (1.0 / m.gamma) * (1.0 + m.gamma / m.kappa).ln();
        assert!((bid - base).abs() < 1e-12);
        assert!((ask - base).abs() < 1e-12);
    }

    #[test]
    fn model_builder_methods() {
        let m = BilateralHawkesOrderFlowImbalance::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0, 0.5)
            .with_lambda_step(0.2)
            .with_dq(0.5)
            .with_inventory_bounds(-10.0, 10.0)
            .with_terminal_condition(TerminalCondition::LiquidationCost)
            .with_terminal_liquidation_half_spread(0.05);

        assert_eq!(m.lambda_step, 0.2);
        assert_eq!(m.dq, 0.5);
        assert_eq!(m.q_min, -10.0);
        assert_eq!(m.q_max, 10.0);
        assert_eq!(m.terminal_liquidation_half_spread, Some(0.05));
    }
}