// docs: processes#geometric-brownian-motion
//! Backs the GBM example on the processes catalog page.

use stochastic_rs::prelude::*;
use stochastic_rs::simd_rng::Deterministic;
use stochastic_rs::stochastic::diffusion::gbm::Gbm;

#[test]
fn gbm_sample_and_sample_par() {
  let p = Gbm::<f64, _>::new(
    0.05,
    0.2,
    1_000,
    Some(100.0),
    Some(1.0),
    Deterministic::new(42),
  );
  let path = p.sample();
  assert_eq!(path.len(), 1_000);

  let paths = p.sample_par(10_000);
  assert_eq!(paths.len(), 10_000);
  assert_eq!(paths[0].len(), 1_000);
}

import stochastic_rs as srs

p = srs.PyGbm(mu=0.05, sigma=0.2, n=1000, x0=100.0, t=1.0)
path = p.sample()                   # numpy.ndarray, shape (1000,)
paths = p.sample_par(10_000)        # shape (10_000, 1000)

// docs: processes#heston-stochastic-volatility
//! Backs the Heston example on the processes catalog page.

use stochastic_rs::simd_rng::Unseeded;
use stochastic_rs::stochastic::volatility::HestonPow;
use stochastic_rs::stochastic::volatility::heston::Heston;
use stochastic_rs::traits::ProcessExt;

#[test]
fn heston_two_factor_sample() {
  let p = Heston::<f64, _>::new(
    Some(100.0),
    Some(0.04),
    2.0,
    0.04,
    0.3,
    -0.7,
    0.03,
    1_000,
    Some(1.0),
    HestonPow::Sqrt,
    Some(true),
    Unseeded,
  );
  let [s_path, v_path] = p.sample();
  assert_eq!(s_path.len(), 1_000);
  assert_eq!(v_path.len(), 1_000);
  assert!(v_path.iter().all(|&v| v >= 0.0));
}

import stochastic_rs as srs

p = srs.PyHeston(kappa=2.0, theta=0.04, sigma=0.3, rho=-0.7, mu=0.03,
               n=1000, s0=100.0, v0=0.04, t=1.0)
s, v = p.sample()                   # both numpy arrays

// docs: processes#heston-stochastic-local-volatility
//! Backs the Heston SLV example on the processes catalog page.

use ndarray::Array1;
use ndarray::Array2;
use stochastic_rs::simd_rng::Deterministic;
use stochastic_rs::stochastic::volatility::heston_slv::HestonSlv;
use stochastic_rs::traits::Expr;
use stochastic_rs::traits::Grid2D;
use stochastic_rs::traits::ProcessExt;

#[test]
fn heston_slv_two_factor_sample_under_an_expression_leverage() {
  // L(t, S) = 1.2 − 0.003 S, a leverage that falls with the spot, written as
  // an expression so the process can also run on a device; eta = 0.6 scales
  // the vol-of-vol.
  let leverage = Expr::lit(1.2) - Expr::x() * 0.003;
  let p = HestonSlv::<f64, _>::new(
    Some(100.0),
    Some(0.04),
    2.0,
    0.04,
    0.3,
    -0.7,
    0.03,
    0.6,
    leverage,
    1_000,
    Some(1.0),
    Deterministic::new(7),
  );
  assert!(p.device_ready());
  let [s_path, v_path] = p.sample();
  assert_eq!(s_path.len(), 1_000);
  assert!(s_path.iter().all(|&s| s > 0.0));
  assert!(v_path.iter().all(|&v| v >= 0.0));

  // A tabulated leverage — what a calibrated `LeverageSurface` converts
  // into — drives the same sampler on the host.
  let grid = Grid2D::new(
    Array1::from_vec(vec![0.0, 1.0]),
    Array1::from_vec(vec![50.0, 150.0]),
    Array2::from_elem((2, 2), 1.0),
  );
  let tabulated = p.with_leverage(grid);
  assert!(!tabulated.device_ready());
  let [s_grid, _] = tabulated.sample();
  assert_eq!(s_grid.len(), 1_000);
}

import stochastic_rs as srs

# leverage: a callable L(t, s), a (spots, times, values) triple of arrays, or
# the LeverageSurface a HestonSlvCalibrator returns
p = srs.PyHestonSlv(kappa=2.0, theta=0.04, sigma=0.3, rho=-0.7, mu=0.03, eta=0.6,
                    leverage=lambda t, s: 1.2 - 0.003 * s, n=1000,
                    s0=100.0, v0=0.04, t=1.0)
s, v = p.sample()                   # both numpy arrays

// docs: processes#fractional-brownian-motion
//! Backs the fGN example on the processes catalog page.

use stochastic_rs::simd_rng::Deterministic;
use stochastic_rs::stochastic::noise::fgn::Fgn;
use stochastic_rs::traits::ProcessExt;

#[test]
fn fgn_sample() {
  let fgn = Fgn::<f64, _>::new(
    /* hurst */ 0.3,
    /* n */ 4096,
    /* t */ Some(1.0),
    Deterministic::new(42),
  );
  let increments = fgn.sample();
  assert_eq!(increments.len(), 4096);
}

import stochastic_rs as srs
import numpy as np

fgn = srs.PyFgn(hurst=0.3, n=4096, t=1.0, seed=42)
increments = fgn.sample()
fbm = np.cumsum(increments)         # fBM = cumsum(fGN)

// docs: processes#quasi-monte-carlo-paths-sobol--brownian-bridge
//! Backs the quasi-Monte Carlo example on the processes page.

use ndarray::Array2;
use stochastic_rs::simd_rng::Deterministic;
use stochastic_rs::stochastic::mc::brownian_bridge_qmc::BrownianBridgeQmc;
use stochastic_rs::stochastic::mc::sobol::SobolSeq;

#[test]
fn qmc_brownian_paths_price_a_martingale() {
  // 2^12 scrambled Sobol paths of a one-year Brownian motion on 64 steps.
  let qmc = BrownianBridgeQmc::scrambled(64, 1.0, &Deterministic::new(7));
  let w: Array2<f64> = qmc.paths(4_096);
  assert_eq!(w.dim(), (4_096, 64));

  // exp(σW_T − σ²T/2) is a martingale, so its QMC average sits at 1 to a
  // few parts in a thousand — tighter than plain MC with the same budget.
  let sigma = 0.2;
  let mean = w
    .column(63)
    .iter()
    .map(|x| (sigma * x - 0.5 * sigma * sigma).exp())
    .sum::<f64>()
    / 4_096.0;
  assert!((mean - 1.0).abs() < 5e-3, "martingale mean {mean}");

  // The sequence itself reaches far beyond the old 21-dimension ceiling.
  let deep: Array2<f64> = SobolSeq::new(5_000).sample(16);
  assert!(deep.iter().all(|u| (0.0..1.0).contains(u)));
}

import numpy as np
import stochastic_rs as srs

qmc = srs.PyBrownianBridgeQmc(steps=64, horizon=1.0, seed=7)   # seed → scrambled
w = qmc.paths(4096)                                            # (4096, 64) Brownian levels
print(np.cov(w[:, -1], w[:, 31])[0, 1])                        # ≈ min(t, s) = 0.5

sobol = srs.PySobolSeq(1000)                                   # 1000 dimensions, unscrambled
print(sobol.sample(8)[:, 999])

// docs: processes#correlated-multi-asset-gbm
//! Backs the correlated multi-asset GBM example on the processes page.

use ndarray::array;
use stochastic_rs::simd_rng::Deterministic;
use stochastic_rs::stochastic::diffusion::multi_gbm::MultiGbm;
use stochastic_rs::traits::ProcessExt;

#[test]
fn multi_gbm_terminal_correlation_follows_rho() {
  // Two assets over one year on 252 steps with correlation −0.4.
  let model = MultiGbm::<f64, _>::new(
    array![0.05, 0.02],
    array![0.2, 0.3],
    array![[1.0, -0.4], [-0.4, 1.0]],
    253,
    array![100.0, 50.0],
    Some(1.0),
    Deterministic::new(7),
  );
  let paths = model.sample_par(8_000);
  assert_eq!(paths[0].dim(), (2, 253));

  // Correlation of the terminal log-returns tracks ρ.
  let logs: Vec<(f64, f64)> = paths
    .iter()
    .map(|p| ((p[(0, 252)] / 100.0).ln(), (p[(1, 252)] / 50.0).ln()))
    .collect();
  let n = logs.len() as f64;
  let (m0, m1) = (
    logs.iter().map(|l| l.0).sum::<f64>() / n,
    logs.iter().map(|l| l.1).sum::<f64>() / n,
  );
  let v0 = logs.iter().map(|l| (l.0 - m0).powi(2)).sum::<f64>() / n;
  let v1 = logs.iter().map(|l| (l.1 - m1).powi(2)).sum::<f64>() / n;
  let cov = logs.iter().map(|l| (l.0 - m0) * (l.1 - m1)).sum::<f64>() / n;
  let corr = cov / (v0 * v1).sqrt();
  assert!((corr + 0.4).abs() < 0.05, "terminal correlation {corr}");
}

import numpy as np
import stochastic_rs as srs

rho = np.array([[1.0, 0.6, -0.2], [0.6, 1.0, 0.1], [-0.2, 0.1, 1.0]])
model = srs.PyMultiGbm(mu=[0.05, 0.03, 0.01], sigma=[0.2, 0.3, 0.15], rho=rho, n=253, x0=[100.0, 50.0, 10.0], t=1.0, seed=7)
paths = model.sample_par(10_000)          # list of (3, 253) arrays
terminal = np.array([p[:, -1] for p in paths])
print(np.corrcoef(np.log(terminal).T))    # ≈ rho

// docs: processes#wishart-process-stochastic-covariance
//! Backs the Wishart process example on the processes page.

use ndarray::Array2;
use ndarray::array;
use stochastic_rs::simd_rng::Deterministic;
use stochastic_rs::stochastic::diffusion::wishart::Wishart;
use stochastic_rs::traits::ProcessExt;

#[test]
fn wishart_terminal_law_matches_the_closed_forms() {
  // 2 × 2 Wishart with a full drift, four exact steps over one year.
  let process = Wishart::<f64, _>::new(
    2.5,
    array![[-0.5, 0.1], [0.05, -0.3]],
    array![[0.3, 0.1], [0.0, 0.2]],
    array![[1.0, 0.2], [0.2, 0.5]],
    5,
    Some(1.0),
    Deterministic::new(7),
  );
  let paths = process.sample_par(8_000);
  assert_eq!(paths[0].dim(), (5, 2, 2));

  // Terminal mean against E[X_1] = m x₀ mᵀ + α q (affine moment formula).
  let mut mean = Array2::<f64>::zeros((2, 2));
  for p in &paths {
    mean += &p.index_axis(ndarray::Axis(0), 4);
  }
  mean /= paths.len() as f64;
  let want = process.mean(1.0);
  assert!(
    (&mean - &want).iter().all(|e| e.abs() < 0.05),
    "mean {mean:?} vs {want:?}"
  );

  // Laplace transform E[exp(Tr(v X_1))] for a negative definite v, eq. (10).
  let v = array![[-0.4, -0.12], [-0.12, -0.4]];
  let mc = paths
    .iter()
    .map(|p| {
      let x = p.index_axis(ndarray::Axis(0), 4);
      (v.dot(&x)).diag().sum().exp()
    })
    .sum::<f64>()
    / paths.len() as f64;
  let exact = process.laplace_transform(&v, 1.0);
  assert!((mc - exact).abs() < 0.03, "laplace {mc} vs {exact}");
}

import numpy as np
import stochastic_rs as srs

b = np.array([[-0.5, 0.1], [0.05, -0.3]])
a = np.array([[0.3, 0.1], [0.0, 0.2]])
x0 = np.array([[1.0, 0.2], [0.2, 0.5]])
process = srs.PyWishart(alpha=2.5, b=b, a=a, x0=x0, n=5, t=1.0, seed=7)
paths = process.sample_par(10_000)                    # list of (5, 2, 2) arrays
terminal = np.mean([p[-1] for p in paths], axis=0)
print(np.abs(terminal - process.mean(1.0)).max())     # ≈ 0: the scheme is exact
v = -0.4 * np.array([[1.0, 0.3], [0.3, 1.0]])
print(process.laplace_transform(v, 1.0))              # E[exp(Tr(v X_1))], eq. (10)

// docs: processes#cheyette-quasi-gaussian-short-rate
//! Backs the Cheyette example on the processes page.

use stochastic_rs::simd_rng::Deterministic;
use stochastic_rs::stochastic::interest::cheyette::Cheyette;
use stochastic_rs::traits::ProcessExt;

fn flat_forward(_t: f64) -> f64 {
  0.03
}

fn displaced_vol(_t: f64, x: f64) -> f64 {
  0.01 + 0.3 * x
}

#[test]
fn cheyette_state_and_bond_reconstruction() {
  // Flat 3 % initial curve, κ = 0.5, displaced local volatility; one year on 100 steps.
  let model = Cheyette::<f64, _>::new(
    flat_forward as fn(f64) -> f64,
    0.5,
    displaced_vol as fn(f64, f64) -> f64,
    101,
    Some(1.0),
    Deterministic::new(7),
  );
  let [x, y] = model.sample();
  assert_eq!(x.len(), 101);
  assert_eq!((x[0], y[0]), (0.0, 0.0));
  // y accumulates the variance of x, so it stays positive along the path.
  assert!(y.iter().skip(1).all(|v| *v > 0.0));

  // Rates and bonds rebuild from the terminal state.
  let r_1 = model.short_rate(1.0, x[100]);
  let p_1_3 = model.zero_bond(1.0, 3.0, x[100], y[100]);
  assert!((r_1 - 0.03).abs() < 0.05 && p_1_3 > 0.8 && p_1_3 < 1.0);

  // At the origin the reconstruction is the initial curve.
  assert!((model.zero_bond(0.0, 2.0, 0.0, 0.0) - (-0.06_f64).exp()).abs() < 1e-12);
}

import numpy as np
import stochastic_rs as srs

model = srs.PyCheyette(f0=lambda t: 0.03, kappa=0.5, sigma=lambda t, x: 0.01 + 0.3 * x, n=101, t=1.0, seed=7)
x, y = model.sample()                          # one state path (x_t, y_t), both length 101
xs, ys = model.sample_par(2000)                # (2000, 101) each
print(model.short_rate(1.0, xs[:, -1].mean()))  # ≈ 0.03: x is centred
print(model.zero_bond(1.0, 3.0, 0.0, 0.0))      # P_0(3) / P_0(1) at the zero state

// docs: processes#euler-engine
//! Backs the Euler-engine example on the processes page.

use stochastic_rs::stochastic::device::Cpu;
use stochastic_rs::stochastic::diffusion::gbm::Gbm;
use stochastic_rs::traits::ProcessExt;
use stochastic_rs_core::simd_rng::Deterministic;

#[test]
fn euler_engine_prices_a_forward() {
  let gbm = Gbm::new(
    0.05,
    0.2,
    253,
    Some(100.0),
    Some(1.0),
    Deterministic::new(7),
  );
  // The backend is a type parameter of the process: `Cpu` (the default) is the
  // process's own SIMD sampler; `.on::<Metal>()` or `.on::<Cuda>()` (with the
  // matching feature) run the Euler kernel on the device through the very
  // same `sample_par` call.
  let paths = gbm.clone().on::<Cpu>().sample_par(20_000); // Vec<Array1<f64>>, 20_000 × 253
  assert_eq!(paths.len(), 20_000);
  let terminal_mean = paths.iter().map(|p| p[252]).sum::<f64>() / 20_000.0;
  assert!((terminal_mean / (100.0 * 0.05_f64.exp()) - 1.0).abs() < 0.01);

  // `on::<Cpu>()` changes nothing: it is exactly what the process returns.
  assert_eq!(paths[0].to_vec(), gbm.sample_par(20_000)[0].to_vec());
}

import numpy as np
import stochastic_rs as srs

# PyGbm(mu, sigma, n, x0=None, t=None, seed=None, dtype=None, device=None)
paths = srs.PyGbm(0.05, 0.2, 253, x0=100.0, t=1.0, seed=7).sample_par(20_000)
print(paths.shape, paths[:, -1].mean())   # (20000, 253), ≈ 100·e^{0.05}
ou = srs.PyOu(2.0, 1.0, 0.5, 501, x0=1.0, t=2.0).sample_par(10_000)
# device= names one backend: "cuda", "metal" or "accelerate", each needing
# the matching cargo feature of the build;
# the single-precision devices want dtype="f32" and return float32 arrays
gpu = srs.PyGbm(0.05, 0.2, 253, x0=100.0, t=1.0, seed=7, dtype="f32", device="metal").sample_par(20_000)