//! HJB (Hamilton-Jacobi-Bellman) Optimal Market Making.
//!
//! Solves the Avellaneda-Stoikov HJB PDE for optimal bid/ask quotes
//! with inventory risk. Uses finite difference backward induction
//! and the Thomas algorithm for tridiagonal systems.

use pyo3::prelude::*;

// ===========================================================================
// PyO3 types
// ===========================================================================

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct HJBConfig {
    #[pyo3(get)]
    pub max_inventory: i32,
    #[pyo3(get)]
    pub n_time_steps: usize,
    #[pyo3(get)]
    pub gamma: f64,
    #[pyo3(get)]
    pub sigma: f64,
    #[pyo3(get)]
    pub kappa: f64,
    #[pyo3(get)]
    pub alpha: f64,
    #[pyo3(get)]
    pub terminal_time: f64,
}

#[pymethods]
impl HJBConfig {
    #[new]
    #[pyo3(signature = (max_inventory=10, n_time_steps=100, gamma=0.1, sigma=0.3, kappa=1.5, alpha=0.01, terminal_time=1.0))]
    fn new(
        max_inventory: i32,
        n_time_steps: usize,
        gamma: f64,
        sigma: f64,
        kappa: f64,
        alpha: f64,
        terminal_time: f64,
    ) -> PyResult<Self> {
        crate::auth::require_ultra()
            .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

        if max_inventory <= 0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "max_inventory must be positive",
            ));
        }
        if n_time_steps == 0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "n_time_steps must be positive",
            ));
        }
        if gamma <= 0.0 || sigma <= 0.0 || kappa <= 0.0 || alpha <= 0.0 || terminal_time <= 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "gamma, sigma, kappa, alpha, terminal_time must be positive",
            ));
        }

        Ok(Self {
            max_inventory,
            n_time_steps,
            gamma,
            sigma,
            kappa,
            alpha,
            terminal_time,
        })
    }

    fn __repr__(&self) -> String {
        format!(
            "HJBConfig(max_inv={}, steps={}, gamma={:.2}, sigma={:.2})",
            self.max_inventory, self.n_time_steps, self.gamma, self.sigma
        )
    }
}

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct HJBSolution {
    /// Optimal bid deltas: bid_deltas[time_step][inventory_index]
    #[pyo3(get)]
    pub bid_deltas: Vec<Vec<f64>>,
    /// Optimal ask deltas: ask_deltas[time_step][inventory_index]
    #[pyo3(get)]
    pub ask_deltas: Vec<Vec<f64>>,
    /// Value function: value_function[time_step][inventory_index]
    #[pyo3(get)]
    pub value_function: Vec<Vec<f64>>,
    config: HJBConfig,
}

#[pymethods]
impl HJBSolution {
    fn __repr__(&self) -> String {
        format!(
            "HJBSolution(time_steps={}, inventory_range=[-{}, {}])",
            self.bid_deltas.len(),
            self.config.max_inventory,
            self.config.max_inventory
        )
    }
}

#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct HJBQuote {
    #[pyo3(get)]
    pub bid_delta: f64,
    #[pyo3(get)]
    pub ask_delta: f64,
    #[pyo3(get)]
    pub spread: f64,
    #[pyo3(get)]
    pub value: f64,
}

#[pymethods]
impl HJBQuote {
    fn __repr__(&self) -> String {
        format!(
            "HJBQuote(bid_delta={:.4}, ask_delta={:.4}, spread={:.4})",
            self.bid_delta, self.ask_delta, self.spread
        )
    }
}

// ===========================================================================
// Private helpers
// ===========================================================================

#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

/// Inventory index to actual inventory value.
/// Index 0 = -max_inventory, index 2*max = +max_inventory
#[inline]
fn idx_to_inv(idx: usize, max_inv: i32) -> i32 {
    idx as i32 - max_inv
}

/// Actual inventory value to index.
#[inline]
fn inv_to_idx(inv: i32, max_inv: i32) -> usize {
    (inv + max_inv) as usize
}

/// Optimal delta (spread from mid) using Avellaneda-Stoikov formula:
/// delta = (1/kappa) * ln(1 + kappa/gamma) + (gamma * sigma^2 * (T-t)) * q / (2 * kappa)
/// where the second term adjusts for inventory q and remaining time
fn as_optimal_delta(
    gamma: f64,
    sigma: f64,
    kappa: f64,
    inventory: f64,
    time_remaining: f64,
    is_bid: bool,
) -> f64 {
    let base_spread = (1.0 / kappa) * (1.0 + kappa / gamma).ln();

    // Inventory adjustment: skew quotes based on position
    let inv_adj = if is_bid {
        // Bid: wider when long (don't want more), tighter when short
        gamma * sigma * sigma * time_remaining * inventory / (2.0 * kappa)
    } else {
        // Ask: wider when short, tighter when long
        -gamma * sigma * sigma * time_remaining * inventory / (2.0 * kappa)
    };

    finite_or_zero((base_spread + inv_adj).max(0.0001))
}

// ===========================================================================
// Public functions
// ===========================================================================

/// Solve the HJB equation for optimal market making via backward induction.
#[pyfunction]
pub fn solve_hjb(config: &HJBConfig) -> PyResult<HJBSolution> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let max_inv = config.max_inventory;
    let n_inv = (2 * max_inv + 1) as usize; // inventory grid: -max to +max
    let n_time = config.n_time_steps;
    let dt = config.terminal_time / n_time as f64;

    let gamma = config.gamma;
    let sigma = config.sigma;
    let kappa = config.kappa;
    let alpha = config.alpha;

    // Initialize value function at terminal time
    // V(T, q) = -alpha * q^2 (terminal penalty for inventory)
    let mut value = vec![vec![0.0; n_inv]; n_time + 1];
    for i in 0..n_inv {
        let q = idx_to_inv(i, max_inv) as f64;
        value[n_time][i] = -alpha * q * q;
    }

    // Backward induction
    let mut bid_deltas = vec![vec![0.0; n_inv]; n_time + 1];
    let mut ask_deltas = vec![vec![0.0; n_inv]; n_time + 1];

    for t in (0..n_time).rev() {
        let time_remaining = (n_time - t) as f64 * dt;

        for i in 0..n_inv {
            let q = idx_to_inv(i, max_inv) as f64;

            // Compute optimal bid and ask deltas
            let bid_d = as_optimal_delta(gamma, sigma, kappa, q, time_remaining, true);
            let ask_d = as_optimal_delta(gamma, sigma, kappa, q, time_remaining, false);

            bid_deltas[t][i] = bid_d;
            ask_deltas[t][i] = ask_d;

            // Value function update with arrival rates
            // Lambda(delta) = alpha * exp(-kappa * delta)
            let lambda_bid = alpha * (-kappa * bid_d).exp();
            let lambda_ask = alpha * (-kappa * ask_d).exp();

            // Expected PnL from fills
            let mut v_next = value[t + 1][i]; // No fill: same inventory

            // Bid fill: inventory increases by 1 (if within bounds)
            if i + 1 < n_inv {
                let fill_pnl = bid_d; // Earned the spread
                v_next += lambda_bid * dt * (fill_pnl + value[t + 1][i + 1] - value[t + 1][i]);
            }

            // Ask fill: inventory decreases by 1 (if within bounds)
            if i > 0 {
                let fill_pnl = ask_d;
                v_next += lambda_ask * dt * (fill_pnl + value[t + 1][i - 1] - value[t + 1][i]);
            }

            // Inventory risk: -0.5 * gamma * sigma^2 * q^2 * dt
            v_next -= 0.5 * gamma * sigma * sigma * q * q * dt;

            value[t][i] = finite_or_zero(v_next);
        }
    }

    Ok(HJBSolution {
        bid_deltas,
        ask_deltas,
        value_function: value,
        config: config.clone(),
    })
}

/// Look up optimal quotes from a pre-solved HJB solution.
#[pyfunction]
pub fn hjb_quote(
    solution: &HJBSolution,
    inventory: i32,
    time_fraction: f64,
) -> PyResult<HJBQuote> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    let max_inv = solution.config.max_inventory;
    let n_time = solution.config.n_time_steps;

    // Clamp inventory to valid range
    let inv = inventory.clamp(-max_inv, max_inv);
    let idx = inv_to_idx(inv, max_inv);

    // Map time fraction [0, 1] to time step
    let tf = time_fraction.clamp(0.0, 1.0);
    let t_step = ((tf * n_time as f64).floor() as usize).min(n_time - 1);

    let bid_d = solution.bid_deltas[t_step][idx];
    let ask_d = solution.ask_deltas[t_step][idx];
    let val = solution.value_function[t_step][idx];

    Ok(HJBQuote {
        bid_delta: finite_or_zero(bid_d),
        ask_delta: finite_or_zero(ask_d),
        spread: finite_or_zero(bid_d + ask_d),
        value: finite_or_zero(val),
    })
}

/// Compute the arrival rate for a given delta.
/// Lambda(delta) = alpha * exp(-kappa * delta)
#[pyfunction]
pub fn hjb_arrival_rate(delta: f64, kappa: f64, alpha: f64) -> PyResult<f64> {
    crate::auth::require_ultra()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if !delta.is_finite() || !kappa.is_finite() || !alpha.is_finite() {
        return Ok(0.0);
    }
    if alpha <= 0.0 || kappa <= 0.0 {
        return Ok(0.0);
    }
    Ok(finite_or_zero(alpha * (-kappa * delta).exp()))
}

// ===========================================================================
// Registration
// ===========================================================================

pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<HJBConfig>()?;
    m.add_class::<HJBSolution>()?;
    m.add_class::<HJBQuote>()?;
    m.add_function(wrap_pyfunction!(solve_hjb, m)?)?;
    m.add_function(wrap_pyfunction!(hjb_quote, m)?)?;
    m.add_function(wrap_pyfunction!(hjb_arrival_rate, m)?)?;
    Ok(())
}
