//! # Svcgmy
//!
//! $$
//! dS_t=S_t\sqrt{v_t}dW_t + S_{t^-}dJ_t^{\mathrm{CGMY}},\quad dv_t=\kappa(\theta-v_t)dt+\xi\sqrt{v_t}dB_t
//! $$
//!
use ndarray::Array1;
use ndarray_rand::RandomExt;
use rand::Rng;
use scilib::math::basic::gamma;

use crate::distributions::exp::SimdExp;
use crate::distributions::uniform::SimdUniform;
use crate::stats::non_central_chi_squared;
use crate::stochastic::process::poisson::Poisson;
use crate::traits::FloatExt;
use crate::traits::ProcessExt;

/// CGMY Stochastic Volatility process
///
/// <https://www.econstor.eu/bitstream/10419/239493/1/175133161X.pdf>
pub struct SVCGMY<T: FloatExt> {
  /// Positive jump rate lambda_plus (corresponds to G)
  pub lambda_plus: T, // G
  /// Negative jump rate lambda_minus (corresponds to M)
  pub lambda_minus: T, // M
  /// Jump activity parameter alpha (corresponds to Y), with 0 < alpha < 2
  pub alpha: T,
  /// Mean reversion rate
  pub kappa: T,
  /// Long-term volatility
  pub eta: T,
  /// Volatility of volatility
  pub zeta: T,
  /// Instantaneous correlation parameter.
  pub rho: T,
  /// Number of time steps
  pub n: usize,
  /// Jumps
  pub j: usize,
  /// Initial value
  pub x0: Option<T>,
  /// Initial value
  pub v0: Option<T>,
  /// Total time horizon
  pub t: Option<T>,
}

impl<T: FloatExt> SVCGMY<T> {
  pub fn new(
    lambda_plus: T,
    lambda_minus: T,
    alpha: T,
    kappa: T,
    eta: T,
    zeta: T,
    rho: T,
    n: usize,
    j: usize,
    x0: Option<T>,
    v0: Option<T>,
    t: Option<T>,
  ) -> Self {
    assert!(lambda_plus > T::zero(), "lambda_plus must be positive");
    assert!(lambda_minus > T::zero(), "lambda_minus must be positive");
    assert!(
      alpha > T::zero() && alpha < T::from_usize_(2),
      "alpha must be in (0, 2)"
    );
    assert!(kappa > T::zero(), "kappa must be positive");
    assert!(eta >= T::zero(), "eta must be non-negative");
    assert!(zeta > T::zero(), "zeta must be positive");
    assert!(
      (-T::one()..=T::one()).contains(&rho),
      "rho must be in [-1, 1]"
    );
    if let Some(v0) = v0 {
      assert!(v0 >= T::zero(), "v0 must be non-negative");
    }

    Self {
      lambda_plus,
      lambda_minus,
      alpha,
      kappa,
      eta,
      zeta,
      rho,
      n,
      j,
      x0,
      v0,
      t,
    }
  }
}

impl<T: FloatExt> ProcessExt<T> for SVCGMY<T> {
  type Output = Array1<T>;

  fn sample(&self) -> Self::Output {
    let mut rng = rand::rng();

    let t_max = self.t.unwrap_or(T::one());
    let dt = t_max / T::from_usize_(self.n - 1);

    let mut x = Array1::<T>::zeros(self.n);
    let mut v = Array1::<T>::zeros(self.n);
    let mut y = Array1::<T>::zeros(self.n);

    x[0] = self.x0.unwrap_or(T::zero());
    v[0] = self.v0.unwrap_or(T::zero());
    y[0] = x[0] - self.rho * v[0];

    let f2 = T::from_usize_(2);

    let g = gamma(2.0 - self.alpha.to_f64().unwrap());
    let C = T::one()
      / (T::from_f64_fast(g)
        * (self.lambda_plus.powf(self.alpha - f2) + self.lambda_minus.powf(self.alpha - f2)));
    let c = (f2 * self.kappa) / ((T::one() - (-self.kappa * dt).exp()) * self.zeta.powi(2));
    let df = T::from_usize_(4) * self.kappa * self.eta / self.zeta.powi(2);

    for i in 1..self.n {
      let ncp = f2 * c * v[i - 1] * (-self.kappa * dt).exp();
      let xi = non_central_chi_squared::sample(df, ncp, &mut rng);
      v[i] = xi / (f2 * c);
    }

    let uniform = SimdUniform::new(T::zero(), T::one());
    let exp = SimdExp::new(T::one());

    let U = Array1::<T>::random(self.j, &uniform);
    let E = Array1::<T>::random(self.j, exp);
    let P = Poisson::new(T::one(), Some(self.j), None);
    let P = P.sample();
    let tau = Array1::<T>::random(self.j, &uniform) * t_max;

    let mut c_tau = Array1::<T>::zeros(self.j);
    for (idx, tau_j) in tau.iter().enumerate() {
      let k = ((*tau_j / dt).ceil()).min(T::from_usize_(self.n - 1));
      let v_k = if k == T::zero() {
        v[0]
      } else {
        v[k.to_usize().unwrap() - 1]
      };
      c_tau[idx] = C * v_k;
    }

    for i in 1..self.n {
      let numerator = v[i - 1]
        * (self.lambda_plus.powf(self.alpha - T::one())
          - self.lambda_minus.powf(self.alpha - T::one()));
      let denominator = (T::one() - self.alpha)
        * (self.lambda_plus.powf(self.alpha - f2) + self.lambda_minus.powf(self.alpha - f2));
      let b = -numerator / denominator;

      let mut jump_component = T::zero();

      let t_1 = T::from_usize_(i - 1) * dt;
      let t = T::from_usize_(i) * dt;

      for j in 0..self.j {
        if tau[j] > t_1 && tau[j] <= t {
          let v_j = if rng.random_bool(0.5) {
            self.lambda_plus
          } else {
            -self.lambda_minus
          };

          let numerator = self.alpha * P[j];
          let denominator = f2 * c_tau[j] * t_max;
          let term1 = (numerator / denominator).powf(-T::one() / self.alpha);
          let term2 = E[j] * U[j].powf(T::one() / self.alpha) / v_j.abs();
          let min_term = term1.min(term2);
          let jump_size = min_term * (v_j / v_j.abs());
          jump_component += jump_size;
        }
      }

      y[i] = y[i - 1] + jump_component + b * dt;
    }

    for i in 1..self.n {
      x[i] = y[i] + self.rho * v[i];
    }

    x
  }
}

py_process_1d!(PySVCGMY, SVCGMY,
  sig: (lambda_plus, lambda_minus, alpha, kappa, eta, zeta, rho, n, j, x0=None, v0=None, t=None, dtype=None),
  params: (lambda_plus: f64, lambda_minus: f64, alpha: f64, kappa: f64, eta: f64, zeta: f64, rho: f64, n: usize, j: usize, x0: Option<f64>, v0: Option<f64>, t: Option<f64>)
);

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

  #[test]
  fn x0_shift_is_preserved_when_no_jumps_and_zero_loading() {
    let model = SVCGMY::new(
      2.0_f64, // lambda_plus
      2.0_f64, // lambda_minus
      0.5,     // alpha
      1.0,     // kappa
      0.04,    // eta
      0.2,     // zeta
      0.0,     // rho
      8,       // n
      0,       // j
      Some(5.0),
      Some(0.0),
      Some(1.0),
    );

    let x = model.sample();
    assert!(x.iter().all(|v| (*v - 5.0).abs() < 1e-12));
  }

  #[test]
  #[should_panic(expected = "alpha must be in (0, 2)")]
  fn invalid_alpha_panics() {
    let _ = SVCGMY::new(
      2.0_f64,
      2.0,
      2.5,
      1.0,
      0.04,
      0.2,
      0.0,
      8,
      0,
      Some(0.0),
      Some(0.01),
      Some(1.0),
    );
  }
}
