oxictl 0.1.1

//! Scenario-based chance-constrained stochastic MPC.
//!
//! Instead of optimising expected cost analytically, scenario MPC samples
//! a finite number of disturbance realisations ("scenarios") and converts the
//! stochastic problem into a deterministic one with scenario-specific constraints.
//!
//! The key idea (scenario approach, Campi & Garatti 2008): if S scenarios are
//! sampled i.i.d. from the disturbance distribution and the resulting finite
//! optimisation is feasible, then the solution satisfies chance constraints
//! with high confidence.
//!
//! Implementation notes:
//! - Scenarios are generated by a deterministic LCG (no `rand` crate).
//! - Disturbance model: additive Gaussian-like noise approximated via Box-Muller
//!   with LCG uniform samples (pure Rust / libm only).
//! - Constraint: Pr{||x_k||_∞ ≤ x_max} ≥ 1 - ε, enforced via scenario count.
//! - Sample average approximation (SAA) optimises the average scenario cost.
#![allow(
    unused,
    clippy::needless_range_loop,
    clippy::too_many_arguments,
    clippy::type_complexity
)]

use crate::core::matrix::{matmul, Matrix};
use crate::core::scalar::ControlScalar;

/// Error type for stochastic MPC operations.
#[derive(Debug)]
pub enum StochasticMpcError {
    /// No scenarios have been generated yet.
    NoScenarios,
    /// The scenario count exceeds the compile-time maximum.
    TooManyScenarios,
}

impl core::fmt::Display for StochasticMpcError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            StochasticMpcError::NoScenarios => write!(f, "Stochastic MPC: no scenarios generated"),
            StochasticMpcError::TooManyScenarios => {
                write!(f, "Stochastic MPC: scenario count exceeds capacity")
            }
        }
    }
}

/// Linear Congruential Generator (LCG) for deterministic pseudo-random numbers.
///
/// Parameters follow the Numerical Recipes convention:
///   x_{n+1} = (a * x_n + c) mod m
///
/// Produces uniform samples in [0, 1) when divided by m.
#[derive(Clone, Copy, Debug)]
pub struct Lcg {
    state: u64,
    a: u64,
    c: u64,
    m: u64,
}

impl Lcg {
    /// Create a new LCG with given seed and standard NR parameters.
    pub fn new(seed: u64) -> Self {
        Self {
            state: seed,
            a: 1_664_525,
            c: 1_013_904_223,
            m: 1u64 << 32,
        }
    }

    /// Advance the LCG and return the next raw state value.
    pub fn next_u64(&mut self) -> u64 {
        self.state = self.a.wrapping_mul(self.state).wrapping_add(self.c) & (self.m - 1);
        self.state
    }

    /// Return the next uniform sample in [0, 1).
    pub fn next_f64(&mut self) -> f64 {
        self.next_u64() as f64 / self.m as f64
    }

    /// Return the next uniform sample in [0, 1) as a generic scalar.
    pub fn next_scalar<S: ControlScalar>(&mut self) -> S {
        S::from_f64(self.next_f64())
    }

    /// Box-Muller transform: convert two uniform samples to one standard normal sample.
    ///
    /// Uses `libm::log` and `libm::cos` for no_std / pure-Rust compatibility.
    pub fn next_normal_f64(&mut self) -> f64 {
        let u1 = self.next_f64().max(1e-15); // avoid log(0)
        let u2 = self.next_f64();
        let two_pi = 2.0 * core::f64::consts::PI;
        libm::sqrt(-2.0 * libm::log(u1)) * libm::cos(two_pi * u2)
    }
}

/// A single disturbance scenario: one realisation of process noise over H steps.
///
/// - N: state dimension
/// - H: prediction horizon
#[derive(Clone, Copy)]
pub struct Scenario<S: ControlScalar, const N: usize, const H: usize> {
    /// Process noise w_k (N×1) at each step.
    pub noise: [Matrix<S, N, 1>; H],
    /// Weight / probability of this scenario (used in SAA).
    pub weight: S,
}

impl<S: ControlScalar, const N: usize, const H: usize> Scenario<S, N, H> {
    /// Create a zero-noise scenario with given weight.
    pub fn zeros(weight: S) -> Self {
        Self {
            noise: [Matrix::zeros(); H],
            weight,
        }
    }
}

/// Scenario-based stochastic MPC.
///
/// Generates `C` scenarios using the deterministic LCG, then optimises the
/// sample-average cost subject to chance constraints enforced by scenario pruning.
///
/// Type parameters:
/// - N: state dimension
/// - I: input dimension
/// - H: prediction horizon
/// - C: maximum number of scenarios (compile-time constant)
pub struct StochasticMpc<
    S: ControlScalar,
    const N: usize,
    const I: usize,
    const H: usize,
    const C: usize,
> {
    /// Nominal state transition matrix A (N×N).
    pub a: Matrix<S, N, N>,
    /// Input matrix B (N×I).
    pub b: Matrix<S, N, I>,
    /// State cost weight Q (N×N).
    pub q: Matrix<S, N, N>,
    /// Input cost weight R (I×I).
    pub r: Matrix<S, I, I>,
    /// Generated scenario set.
    pub scenarios: [Scenario<S, N, H>; C],
    /// Number of active scenarios (≤ C).
    pub n_scenarios: usize,
    /// Current state.
    pub x: Matrix<S, N, 1>,
    /// Process noise standard deviation (isotropic).
    pub noise_std: S,
    /// State constraint: ||x_k||_inf ≤ x_max.
    pub x_max: S,
    /// Input constraint: ||u_k||_inf ≤ u_max.
    pub u_max: S,
    /// Chance constraint violation tolerance ε (0 < ε < 1).
    pub epsilon: S,
    /// Number of gradient iterations for SAA optimisation.
    pub iterations: usize,
    /// LCG seed for scenario generation.
    lcg_seed: u64,
}

impl<S: ControlScalar, const N: usize, const I: usize, const H: usize, const C: usize>
    StochasticMpc<S, N, I, H, C>
{
    /// Create a new StochasticMpc controller.
    pub fn new(
        a: Matrix<S, N, N>,
        b: Matrix<S, N, I>,
        q: Matrix<S, N, N>,
        r: Matrix<S, I, I>,
        noise_std: S,
        x_max: S,
        u_max: S,
        epsilon: S,
        iterations: usize,
        lcg_seed: u64,
    ) -> Self {
        Self {
            a,
            b,
            q,
            r,
            scenarios: [Scenario::zeros(S::ZERO); C],
            n_scenarios: 0,
            x: Matrix::zeros(),
            noise_std,
            x_max,
            u_max,
            epsilon,
            iterations,
            lcg_seed,
        }
    }

    /// Generate `n` scenarios using the deterministic LCG.
    ///
    /// Each scenario consists of H additive noise vectors w_k ~ N(0, noise_std^2 I).
    /// Returns an error if n > C.
    pub fn generate_scenarios(&mut self, n: usize) -> Result<(), StochasticMpcError> {
        if n > C {
            return Err(StochasticMpcError::TooManyScenarios);
        }
        let mut lcg = Lcg::new(self.lcg_seed);
        let weight = S::ONE / S::from_f64(n as f64);
        let std = self.noise_std;

        for s in 0..n {
            let mut scenario = Scenario::zeros(weight);
            for k in 0..H {
                for i in 0..N {
                    let z = lcg.next_normal_f64();
                    scenario.noise[k].data[i][0] = std * S::from_f64(z);
                }
            }
            self.scenarios[s] = scenario;
        }
        self.n_scenarios = n;
        Ok(())
    }

    /// Compute the stage cost: x^T Q x + u^T R u.
    fn stage_cost(&self, x: &Matrix<S, N, 1>, u: &Matrix<S, I, 1>) -> S {
        let qx = matmul(&self.q, x);
        let xt = x.transpose();
        let cx = matmul(&xt, &qx).data[0][0];

        let ru = matmul(&self.r, u);
        let ut = u.transpose();
        let cu = matmul(&ut, &ru).data[0][0];

        cx + cu
    }

    /// Propagate state under nominal dynamics + noise: x_{k+1} = A x_k + B u_k + w_k.
    fn propagate_noisy(
        &self,
        x: &Matrix<S, N, 1>,
        u: &Matrix<S, I, 1>,
        w: &Matrix<S, N, 1>,
    ) -> Matrix<S, N, 1> {
        let ax = matmul(&self.a, x);
        let bu = matmul(&self.b, u);
        ax.add_mat(&bu).add_mat(w)
    }

    /// Compute the sample-average approximation (SAA) cost for a given input sequence.
    pub fn saa_cost(&self, u_seq: &[Matrix<S, I, 1>; H]) -> S {
        if self.n_scenarios == 0 {
            return S::ZERO;
        }
        let mut total = S::ZERO;
        for s in 0..self.n_scenarios {
            let sc = &self.scenarios[s];
            let mut x = self.x;
            let mut sc_cost = S::ZERO;
            for k in 0..H {
                sc_cost += self.stage_cost(&x, &u_seq[k]);
                x = self.propagate_noisy(&x, &u_seq[k], &sc.noise[k]);
            }
            total += sc.weight * sc_cost;
        }
        total
    }

    /// Count the number of scenarios that violate the state constraint ||x||_inf ≤ x_max
    /// at any step along the horizon.
    pub fn constraint_violations(&self, u_seq: &[Matrix<S, I, 1>; H]) -> usize {
        let mut violations = 0_usize;
        for s in 0..self.n_scenarios {
            let sc = &self.scenarios[s];
            let mut x = self.x;
            let mut violated = false;
            for k in 0..H {
                x = self.propagate_noisy(&x, &u_seq[k], &sc.noise[k]);
                for i in 0..N {
                    if x.data[i][0] > self.x_max || x.data[i][0] < -self.x_max {
                        violated = true;
                        break;
                    }
                }
                if violated {
                    break;
                }
            }
            if violated {
                violations += 1;
            }
        }
        violations
    }

    /// Empirical chance constraint satisfaction ratio.
    ///
    /// Returns the fraction of scenarios satisfying the state constraint.
    pub fn constraint_satisfaction_ratio(&self, u_seq: &[Matrix<S, I, 1>; H]) -> S {
        if self.n_scenarios == 0 {
            return S::ONE;
        }
        let violations = self.constraint_violations(u_seq);
        let satisfied = self.n_scenarios.saturating_sub(violations);
        S::from_f64(satisfied as f64 / self.n_scenarios as f64)
    }

    /// Numerical gradient of SAA cost w.r.t. u_k (central differences).
    fn gradient_uk(&self, k: usize, u_seq: &[Matrix<S, I, 1>; H], eps: S) -> Matrix<S, I, 1> {
        let two_eps = S::TWO * eps;
        let mut grad = Matrix::<S, I, 1>::zeros();
        for i in 0..I {
            let mut u_p = *u_seq;
            let mut u_m = *u_seq;
            u_p[k].data[i][0] += eps;
            u_m[k].data[i][0] -= eps;
            let cp = self.saa_cost(&u_p);
            let cm = self.saa_cost(&u_m);
            grad.data[i][0] = (cp - cm) / two_eps;
        }
        grad
    }

    /// Project input onto box constraint [-u_max, u_max].
    fn project_input(&self, u: Matrix<S, I, 1>) -> Matrix<S, I, 1> {
        let mut out = u;
        for i in 0..I {
            out.data[i][0] = out.data[i][0].clamp_val(-self.u_max, self.u_max);
        }
        out
    }

    /// Solve the stochastic MPC problem via SAA gradient descent.
    ///
    /// Requires that `generate_scenarios` has been called first.
    /// Returns the first optimal input action, or an error if no scenarios exist.
    pub fn solve(&mut self) -> Result<Matrix<S, I, 1>, StochasticMpcError> {
        if self.n_scenarios == 0 {
            return Err(StochasticMpcError::NoScenarios);
        }

        let eps = S::from_f64(1e-5);
        let step = S::from_f64(1e-3);

        let mut u_seq: [Matrix<S, I, 1>; H] = [Matrix::zeros(); H];

        for _iter in 0..self.iterations {
            for k in 0..H {
                let g = self.gradient_uk(k, &u_seq, eps);
                for i in 0..I {
                    u_seq[k].data[i][0] -= step * g.data[i][0];
                }
                u_seq[k] = self.project_input(u_seq[k]);
            }
        }

        Ok(u_seq[0])
    }

    /// Set the current state.
    pub fn set_state(&mut self, x: Matrix<S, N, 1>) {
        self.x = x;
    }

    /// Reseed the LCG for a new scenario draw at the next call to `generate_scenarios`.
    pub fn reseed(&mut self, seed: u64) {
        self.lcg_seed = seed;
    }

    /// Prune scenarios whose constraint violations exceed the allowed fraction ε.
    ///
    /// Removes individual scenarios that violate the state constraint, keeping
    /// only the scenarios satisfying ||x_k||_inf ≤ x_max for all k.
    /// Renormalises weights after pruning.
    pub fn prune_violating_scenarios(&mut self, u_seq: &[Matrix<S, I, 1>; H]) {
        let mut kept = 0_usize;
        let mut kept_indices = [0usize; C];

        for s in 0..self.n_scenarios {
            let sc = &self.scenarios[s];
            let mut x = self.x;
            let mut ok = true;
            for k in 0..H {
                x = self.propagate_noisy(&x, &u_seq[k], &sc.noise[k]);
                for i in 0..N {
                    if x.data[i][0] > self.x_max || x.data[i][0] < -self.x_max {
                        ok = false;
                        break;
                    }
                }
                if !ok {
                    break;
                }
            }
            if ok {
                kept_indices[kept] = s;
                kept += 1;
            }
        }

        // Re-pack scenarios and renormalise
        if kept > 0 {
            let new_weight = S::ONE / S::from_f64(kept as f64);
            let old_scenarios = self.scenarios;
            for j in 0..kept {
                self.scenarios[j] = old_scenarios[kept_indices[j]];
                self.scenarios[j].weight = new_weight;
            }
            self.n_scenarios = kept;
        }
    }
}

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

    fn make_stochastic_mpc() -> StochasticMpc<f64, 2, 1, 4, 8> {
        let mut a = Matrix::<f64, 2, 2>::identity();
        a.data[0][1] = 0.1;

        let mut b = Matrix::<f64, 2, 1>::zeros();
        b.data[0][0] = 0.005;
        b.data[1][0] = 0.1;

        let q = Matrix::<f64, 2, 2>::identity();
        let mut r = Matrix::<f64, 1, 1>::zeros();
        r.data[0][0] = 0.1;

        StochasticMpc::new(a, b, q, r, 0.01_f64, 5.0_f64, 1.0_f64, 0.05_f64, 30, 42)
    }

    #[test]
    fn lcg_produces_deterministic_sequence() {
        let mut lcg1 = Lcg::new(42);
        let mut lcg2 = Lcg::new(42);
        for _ in 0..100 {
            assert_eq!(lcg1.next_u64(), lcg2.next_u64());
        }
    }

    #[test]
    fn lcg_uniform_in_range() {
        let mut lcg = Lcg::new(1337);
        for _ in 0..200 {
            let u = lcg.next_f64();
            assert!((0.0..1.0).contains(&u), "LCG out of range: {}", u);
        }
    }

    #[test]
    fn lcg_normal_finite() {
        let mut lcg = Lcg::new(999);
        for _ in 0..50 {
            let z = lcg.next_normal_f64();
            assert!(z.is_finite(), "Box-Muller produced non-finite value: {}", z);
        }
    }

    #[test]
    fn generate_scenarios_fills_correctly() {
        let mut mpc = make_stochastic_mpc();
        let result = mpc.generate_scenarios(6);
        assert!(result.is_ok());
        assert_eq!(mpc.n_scenarios, 6);
    }

    #[test]
    fn too_many_scenarios_returns_error() {
        let mut mpc = make_stochastic_mpc();
        let result = mpc.generate_scenarios(100); // C = 8
        assert!(matches!(result, Err(StochasticMpcError::TooManyScenarios)));
    }

    #[test]
    fn saa_cost_zero_with_zero_scenarios_and_zero_state() {
        let mpc = make_stochastic_mpc();
        let u_seq = [Matrix::<f64, 1, 1>::zeros(); 4];
        let cost = mpc.saa_cost(&u_seq);
        assert!(
            cost < 1e-12,
            "SAA cost should be 0 with no scenarios: {}",
            cost
        );
    }

    #[test]
    fn solve_without_scenarios_returns_error() {
        let mut mpc = make_stochastic_mpc();
        let result = mpc.solve();
        assert!(matches!(result, Err(StochasticMpcError::NoScenarios)));
    }

    #[test]
    fn solve_with_scenarios_succeeds() {
        let mut mpc = make_stochastic_mpc();
        mpc.generate_scenarios(4).expect("scenario generation");
        let mut x0 = Matrix::<f64, 2, 1>::zeros();
        x0.data[0][0] = 0.5;
        mpc.set_state(x0);
        let result = mpc.solve();
        assert!(result.is_ok(), "solve should succeed: {:?}", result);
    }

    #[test]
    fn constraint_satisfaction_is_in_range() {
        let mut mpc = make_stochastic_mpc();
        mpc.generate_scenarios(8).expect("scenario generation");
        let u_seq = [Matrix::<f64, 1, 1>::zeros(); 4];
        let ratio = mpc.constraint_satisfaction_ratio(&u_seq);
        assert!(
            (0.0..=1.0).contains(&ratio),
            "Ratio out of [0,1]: {}",
            ratio
        );
    }
}