optimiz-rs/src/optimal_control/mrsjd.rs

//! Markov Regime Switching Jump Diffusion (MRSJD)
//! ==============================================
//!
//! Complete implementation following the paper:
//! "Markov Regime Switching Jump Diffusion Model and the Control Problem"
//!
//! # Mathematical Framework
//!
//! ## Full Dynamics
//!
//! State: (X_t, α_t) where α_t ∈ {1,...,K} is regime
//!
//! dX_t = μ^{α_t}(X_t)dt + σ^{α_t}(X_t)dW_t + dJ_t^{α_t}
//!
//! - Regime transitions: q_{ij}dt probability
//! - Jump intensity and distribution depend on regime: λ^i, F^i
//!
//! ## Coupled HJB System with Jumps
//!
//! ρV^i(x) = sup_u [μ^i·∇V^i + (σ^i)²/2·∇²V^i + L^i(x,u)
//!            + λ^i∫[V^i(x+y) - V^i(x)]F^i(dy)
//!            + Σ_{j≠i} q_{ij}[V^j(x) - V^i(x)]]
//!
//! This is the most general formulation combining:
//! 1. Diffusion processes
//! 2. Jump processes  
//! 3. Regime switching
//! 4. Optimal control

use crate::optimal_control::{
    jump_diffusion::JumpDistribution,
    OptimalControlError, Result,
};
use ndarray::{Array1, Array2};

/// Regime-specific jump parameters
pub struct RegimeJumpParameters {
    /// Drift μ^i(x)
    pub drift: Box<dyn Fn(f64) -> f64 + Send + Sync>,
    /// Diffusion σ^i(x)
    pub diffusion: Box<dyn Fn(f64) -> f64 + Send + Sync>,
    /// Running cost L^i(x, u)
    pub cost: Box<dyn Fn(f64, f64) -> f64 + Send + Sync>,
    /// Jump intensity λ^i
    pub jump_intensity: f64,
    /// Jump distribution F^i
    pub jump_distribution: JumpDistribution,
}

/// MRSJD configuration
#[derive(Debug, Clone)]
pub struct MRSJDConfig {
    /// Number of regimes
    pub n_regimes: usize,
    /// Transition rate matrix Q
    pub transition_rates: Array2<f64>,
    /// Discount rate
    pub rho: f64,
    /// Transaction cost
    pub transaction_cost: f64,
    /// State space bounds
    pub state_bounds: (f64, f64),
    /// Number of grid points
    pub n_points: usize,
    /// Maximum iterations
    pub max_iter: usize,
    /// Convergence tolerance
    pub tolerance: f64,
}

impl Default for MRSJDConfig {
    fn default() -> Self {
        let mut q = Array2::<f64>::zeros((2, 2));
        q[[0, 1]] = 0.5;
        q[[1, 0]] = 0.3;

        Self {
            n_regimes: 2,
            transition_rates: q,
            rho: 0.04,
            transaction_cost: 0.001,
            state_bounds: (-4.0, 4.0),
            n_points: 400, // More points needed for jumps
            max_iter: 3000,
            tolerance: 1e-6,
        }
    }
}

/// MRSJD result
#[derive(Debug, Clone)]
pub struct MRSJDResult {
    /// State space grid
    pub x: Array1<f64>,
    /// Value functions V^i(x)
    pub values: Array2<f64>,
    /// Optimal controls u^i(x)
    pub controls: Array2<f64>,
    /// Gradients
    pub gradients: Array2<f64>,
    /// Jump integral contributions
    pub jump_integrals: Array2<f64>,
    /// Stationary distribution
    pub stationary_distribution: Array1<f64>,
    /// Iterations
    pub iterations: usize,
    /// Residual
    pub residual: f64,
}

/// Markov Regime Switching Jump Diffusion Solver
pub struct MRSJDSolver {
    config: MRSJDConfig,
    regime_params: Vec<RegimeJumpParameters>,
}

impl MRSJDSolver {
    /// Create new MRSJD solver
    pub fn new(config: MRSJDConfig, regime_params: Vec<RegimeJumpParameters>) -> Result<Self> {
        // Validation
        if config.n_regimes != regime_params.len() {
            return Err(OptimalControlError::InvalidParameters(format!(
                "Need {} regime parameters",
                config.n_regimes
            )));
        }

        if config.n_points < 200 {
            return Err(OptimalControlError::InvalidParameters(
                "Need at least 200 grid points for MRSJD".to_string(),
            ));
        }

        // Validate transition rates
        for i in 0..config.n_regimes {
            for j in 0..config.n_regimes {
                if i != j && config.transition_rates[[i, j]] < 0.0 {
                    return Err(OptimalControlError::InvalidParameters(
                        "Transition rates must be non-negative".to_string(),
                    ));
                }
            }
        }

        Ok(Self {
            config,
            regime_params,
        })
    }

    /// Solve the coupled MRSJD system
    pub fn solve(&self) -> Result<MRSJDResult> {
        let cfg = &self.config;

        // Create grid
        let (x_min, x_max) = cfg.state_bounds;
        let dx = (x_max - x_min) / (cfg.n_points - 1) as f64;
        let x = Array1::from_iter((0..cfg.n_points).map(|i| x_min + i as f64 * dx));

        // Precompute jump kernels for each regime
        let jump_kernels = self.compute_all_jump_kernels(&x, dx)?;

        // Setup transition matrix
        let mut q = cfg.transition_rates.clone();
        for i in 0..cfg.n_regimes {
            let row_sum: f64 = (0..cfg.n_regimes)
                .filter(|&j| j != i)
                .map(|j| q[[i, j]])
                .sum();
            q[[i, i]] = -row_sum;
        }

        // Initialize
        let mut v = Array2::<f64>::zeros((cfg.n_regimes, cfg.n_points));
        let mut v_old = Array2::<f64>::zeros((cfg.n_regimes, cfg.n_points));
        let mut u = Array2::<f64>::zeros((cfg.n_regimes, cfg.n_points));
        let mut jump_integrals = Array2::<f64>::zeros((cfg.n_regimes, cfg.n_points));

        // Main iteration loop
        let mut iterations = 0;
        let mut residual = f64::INFINITY;

        for iter in 0..cfg.max_iter {
            v_old.assign(&v);

            // Solve for each regime
            for regime in 0..cfg.n_regimes {
                self.solve_regime_mrsjd(
                    regime,
                    &x,
                    &mut v,
                    &v_old,
                    &mut u,
                    &mut jump_integrals,
                    &q,
                    &jump_kernels[regime],
                    dx,
                )?;
            }

            // Check convergence
            residual =
                (&v - &v_old).mapv(|x| x.abs()).sum() / (cfg.n_regimes * cfg.n_points) as f64;
            iterations = iter + 1;

            if residual < cfg.tolerance {
                break;
            }
            // No under-relaxation: the implicit-in-space solve is
            // unconditionally stable, damping only slows convergence.
        }

        // NOTE: `!(residual < tolerance)` (rather than `residual >= tolerance`)
        // also catches NaN residuals, which otherwise slip through both
        // comparisons and produce a silently-invalid Ok result.
        if !(residual < cfg.tolerance) {
            return Err(OptimalControlError::ConvergenceError(format!(
                "Failed to converge after {} iterations, residual = {:.2e}",
                iterations, residual
            )));
        }

        // Compute gradients
        let gradients = self.compute_all_gradients(&v, dx);

        // Stationary distribution
        let stationary_dist = self.compute_stationary_distribution(&q)?;

        Ok(MRSJDResult {
            x,
            values: v,
            controls: u,
            gradients,
            jump_integrals,
            stationary_distribution: stationary_dist,
            iterations,
            residual,
        })
    }

    /// Solve HJB for one regime with jumps
    fn solve_regime_mrsjd(
        &self,
        regime: usize,
        x: &Array1<f64>,
        v: &mut Array2<f64>,
        v_old: &Array2<f64>,
        u: &mut Array2<f64>,
        jump_int: &mut Array2<f64>,
        q: &Array2<f64>,
        jump_kernel: &Array2<f64>,
        dx: f64,
    ) -> Result<()> {
        let cfg = &self.config;
        let params = &self.regime_params[regime];
        let n = cfg.n_points;

        // Jump inflow λ·Σ_j k_ij v_old_j and outflow mass λ·Σ_j k_ij per node.
        // The kernel row sum can be < 1 (jumps leaving the grid are dropped),
        // so track it explicitly to keep the scheme conservative.
        let lambda = params.jump_intensity;
        let mut jump_mass = vec![0.0_f64; n];
        for i in 0..n {
            let mut inflow = 0.0;
            let mut mass = 0.0;
            for j in 0..n {
                inflow += jump_kernel[[i, j]] * v_old[[regime, j]];
                mass += jump_kernel[[i, j]];
            }
            jump_int[[regime, i]] = lambda * (inflow - mass * v_old[[regime, i]]);
            jump_mass[i] = lambda * mass;
        }

        // Implicit-in-space solve (Kushner–Dupuis upwind discretisation).
        // The stationary HJB  ρv = μ v' + ½σ² v'' + jump + switching + cost
        // is rearranged into a diagonally dominant tridiagonal system per
        // regime (jumps and regime coupling explicit via v_old), which is
        // unconditionally stable — a pointwise Jacobi update diverges here
        // because σ²/dx² ≫ ρ.
        let mut sub = vec![0.0_f64; n]; // a_i · v_{i-1}
        let mut diag = vec![0.0_f64; n]; // b_i · v_i
        let mut sup = vec![0.0_f64; n]; // c_i · v_{i+1}
        let mut rhs = vec![0.0_f64; n];

        for i in 1..n - 1 {
            let xi = x[i];
            let mu = (params.drift)(xi);
            let sigma = (params.diffusion)(xi);
            let sig2 = sigma * sigma;
            let mu_p = mu.max(0.0);
            let mu_m = mu.min(0.0);

            // Control from the current value gradient (policy-iteration style)
            let dv_forward = (v_old[[regime, i + 1]] - v_old[[regime, i]]) / dx;
            let dv_backward = (v_old[[regime, i]] - v_old[[regime, i - 1]]) / dx;
            let optimal_control = self.optimize_control_mrsjd(xi, dv_forward, dv_backward, params);
            u[[regime, i]] = optimal_control;
            let cost = (params.cost)(xi, optimal_control);

            // Total outflow rate to other regimes
            let q_out: f64 = (0..cfg.n_regimes)
                .filter(|&j| j != regime)
                .map(|j| q[[regime, j]])
                .sum();

            sub[i] = -(mu_p / dx + 0.5 * sig2 / (dx * dx));
            sup[i] = mu_m / dx - 0.5 * sig2 / (dx * dx);
            diag[i] = cfg.rho + mu_p / dx - mu_m / dx + sig2 / (dx * dx) + jump_mass[i] + q_out;

            // Explicit couplings: jump inflow + other-regime values
            let switching_in: f64 = (0..cfg.n_regimes)
                .filter(|&j| j != regime)
                .map(|j| q[[regime, j]] * v_old[[j, i]])
                .sum();
            let jump_inflow = jump_int[[regime, i]] + jump_mass[i] * v_old[[regime, i]];
            rhs[i] = cost + jump_inflow + switching_in;
        }

        // Neumann boundaries: v_0 = v_1, v_{n-1} = v_{n-2}
        diag[0] = 1.0;
        sup[0] = -1.0;
        rhs[0] = 0.0;
        sub[n - 1] = -1.0;
        diag[n - 1] = 1.0;
        rhs[n - 1] = 0.0;

        // Thomas algorithm (forward sweep + back substitution)
        for i in 1..n {
            let w = sub[i] / diag[i - 1];
            diag[i] -= w * sup[i - 1];
            rhs[i] -= w * rhs[i - 1];
        }
        v[[regime, n - 1]] = rhs[n - 1] / diag[n - 1];
        for i in (0..n - 1).rev() {
            v[[regime, i]] = (rhs[i] - sup[i] * v[[regime, i + 1]]) / diag[i];
        }

        u[[regime, 0]] = u[[regime, 1]];
        u[[regime, n - 1]] = u[[regime, n - 2]];

        Ok(())
    }

    /// Compute jump kernels for all regimes
    #[allow(unused_variables)]  // dx parameter reserved for future extensions
    fn compute_all_jump_kernels(&self, x: &Array1<f64>, dx: f64) -> Result<Vec<Array2<f64>>> {
        use rand::thread_rng;
        let mut rng = thread_rng();

        let n = x.len();
        let mut kernels = Vec::with_capacity(self.config.n_regimes);

        for regime in 0..self.config.n_regimes {
            let mut kernel = Array2::<f64>::zeros((n, n));
            let dist = &self.regime_params[regime].jump_distribution;

            // Monte Carlo discretization
            let n_samples = 20000;
            for i in 0..n {
                let mut jump_counts = vec![0; n];

                for _ in 0..n_samples {
                    let jump_size = dist.sample(&mut rng);
                    let target_x = x[i] + jump_size;

                    if let Some(j) = self.find_nearest_index(x, target_x) {
                        jump_counts[j] += 1;
                    }
                }

                // Normalize
                for j in 0..n {
                    kernel[[i, j]] = jump_counts[j] as f64 / n_samples as f64;
                }
            }

            kernels.push(kernel);
        }

        Ok(kernels)
    }

    /// Find nearest grid point
    fn find_nearest_index(&self, x: &Array1<f64>, target: f64) -> Option<usize> {
        let (x_min, x_max) = self.config.state_bounds;

        if target < x_min || target > x_max {
            return None;
        }

        let mut best_idx = 0;
        let mut best_dist = (x[0] - target).abs();

        for (i, &xi) in x.iter().enumerate() {
            let dist = (xi - target).abs();
            if dist < best_dist {
                best_dist = dist;
                best_idx = i;
            }
        }

        Some(best_idx)
    }

    /// Optimize control (problem-specific)
    fn optimize_control_mrsjd(
        &self,
        _x: f64,
        _dv_forward: f64,
        _dv_backward: f64,
        _params: &RegimeJumpParameters,
    ) -> f64 {
        // Placeholder - implement specific optimization
        0.0
    }

    /// Compute gradients for all regimes
    fn compute_all_gradients(&self, v: &Array2<f64>, dx: f64) -> Array2<f64> {
        let (n_regimes, n_points) = v.dim();
        let mut grad = Array2::<f64>::zeros((n_regimes, n_points));

        for i in 0..n_regimes {
            for j in 1..n_points - 1 {
                grad[[i, j]] = (v[[i, j + 1]] - v[[i, j - 1]]) / (2.0 * dx);
            }
            grad[[i, 0]] = (v[[i, 1]] - v[[i, 0]]) / dx;
            grad[[i, n_points - 1]] = (v[[i, n_points - 1]] - v[[i, n_points - 2]]) / dx;
        }

        grad
    }

    /// Compute stationary distribution
    fn compute_stationary_distribution(&self, q: &Array2<f64>) -> Result<Array1<f64>> {
        use ndarray_linalg::Solve;

        let n = q.nrows();
        let mut a = q.t().to_owned();

        for i in 0..n {
            a[[i, i]] -= q[[i, i]];
        }

        // Replace last equation with Σπ_i = 1
        for j in 0..n {
            a[[n - 1, j]] = 1.0;
        }

        let mut b = Array1::<f64>::zeros(n);
        b[n - 1] = 1.0;

        match a.solve(&b) {
            Ok(pi) => Ok(pi),
            Err(_) => Err(OptimalControlError::MatrixError(
                "Failed to compute stationary distribution".to_string(),
            )),
        }
    }
}

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

    #[test]
    fn test_mrsjd_solver_generic() {
        // Generic test: State evolution with regime switching and jumps
        // Using abstract state variable (not portfolio-specific)

        let mut q = Array2::<f64>::zeros((2, 2));
        q[[0, 1]] = 0.5; // Regime 0 → 1
        q[[1, 0]] = 0.3; // Regime 1 → 0

        let config = MRSJDConfig {
            n_regimes: 2,
            transition_rates: q,
            state_bounds: (-1.0, 3.0),
            n_points: 200,
            rho: 0.05,
            transaction_cost: 0.0,
            max_iter: 200,
            tolerance: 1e-4,
        };

        // Regime 0: Higher drift, lower volatility, fewer jumps
        let params_0 = RegimeJumpParameters {
            drift: Box::new(|x| 0.5 - 0.1 * x), // Mean-reverting to 0.5
            diffusion: Box::new(|_x| 0.2),
            cost: Box::new(|x, _u| x.powi(2)), // Quadratic cost
            jump_intensity: 0.1,
            jump_distribution: JumpDistribution::Normal {
                mean: -0.05,
                std: 0.1,
            },
        };

        // Regime 1: Lower drift, higher volatility, more frequent jumps
        let params_1 = RegimeJumpParameters {
            drift: Box::new(|x| 0.2 - 0.05 * x),
            diffusion: Box::new(|_x| 0.4),
            cost: Box::new(|x, _u| x.powi(2)),
            jump_intensity: 0.3,
            jump_distribution: JumpDistribution::Normal {
                mean: -0.1,
                std: 0.15,
            },
        };

        let solver = MRSJDSolver::new(config, vec![params_0, params_1]).unwrap();
        let result = solver.solve().unwrap();

        assert_eq!(result.values.nrows(), 2);
        assert!(result.iterations > 0);
        assert!(result.residual < 1e-4);

        // Stationary distribution should sum to 1
        let sum: f64 = result.stationary_distribution.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6);
    }
}