//! SDE expression virtual machine.
//!
//! A register-based scalar VM for evaluating stochastic differential equation
//! components (drift, diffusion, jump, state updates). Each tape is a flat
//! `Vec<SdeOp>` that gets executed once per time-step per path.

use rand::Rng;
use rand_chacha::ChaCha8Rng;
use rand_distr::{Distribution, Normal};

// ---------------------------------------------------------------------------
// Instruction set
// ---------------------------------------------------------------------------

/// A single instruction in an SDE expression tape.
#[derive(Clone, Debug)]
pub enum SdeOp {
    // ── Leaf ──────────────────────────────────────────────────────────────
    /// Push a constant.
    Literal(f64),
    /// Load a state variable by index (0 = price S).
    LoadState(usize),
    /// Load a model parameter by index.
    LoadParam(usize),
    /// Load a built-in variable (dt, t, step index, last log-return).
    LoadBuiltin(Builtin),

    // ── Arithmetic ───────────────────────────────────────────────────────
    Add(usize, usize),
    Sub(usize, usize),
    Mul(usize, usize),
    Div(usize, usize),
    Pow(usize, usize),
    Neg(usize),

    // ── Math functions ───────────────────────────────────────────────────
    Abs(usize),
    Sqrt(usize),
    Log(usize),
    Exp(usize),
    Floor(usize),
    Max(usize, usize),
    Min(usize, usize),

    // ── Comparison / logic ───────────────────────────────────────────────
    Gt(usize, usize),
    Lt(usize, usize),
    Gte(usize, usize),
    Lte(usize, usize),
    Eq(usize, usize),
    And(usize, usize),
    Or(usize, usize),
    Not(usize),
    IfElse {
        cond: usize,
        then_val: usize,
        else_val: usize,
    },

    // ── Random (consume the per-path RNG) ────────────────────────────────
    /// Sample from N(0, 1).
    RandNormal,
    /// Sample from U(0, 1).
    RandUniform,
}

/// Built-in variables injected by the runner at each time-step.
#[derive(Clone, Copy, Debug)]
pub enum Builtin {
    /// Time-step size in years.
    Dt,
    /// Current simulation time (step_idx * dt).
    T,
    /// Current step index (0-based).
    StepIdx,
    /// Log-return from the previous step (0.0 on first step).
    Ret,
}

// ---------------------------------------------------------------------------
// Runtime context
// ---------------------------------------------------------------------------

/// Immutable per-step context supplied by the runner.
#[derive(Clone, Copy, Debug)]
pub struct SdeBuiltins {
    pub dt: f64,
    pub t: f64,
    pub step_idx: usize,
    pub ret: f64,
}

// ---------------------------------------------------------------------------
// Tape evaluator
// ---------------------------------------------------------------------------

/// Evaluate an SDE op tape and return the final register value.
///
/// `registers` is caller-owned scratch space (avoids allocation per call).
/// It must be at least as long as `ops`.
#[inline]
pub fn eval_tape(
    ops: &[SdeOp],
    registers: &mut [f64],
    states: &[f64],
    params: &[f64],
    builtins: &SdeBuiltins,
    rng: &mut ChaCha8Rng,
) -> f64 {
    debug_assert!(registers.len() >= ops.len());

    for (i, op) in ops.iter().enumerate() {
        registers[i] = match *op {
            // Leaf
            SdeOp::Literal(v) => v,
            SdeOp::LoadState(idx) => states[idx],
            SdeOp::LoadParam(idx) => params[idx],
            SdeOp::LoadBuiltin(b) => match b {
                Builtin::Dt => builtins.dt,
                Builtin::T => builtins.t,
                Builtin::StepIdx => builtins.step_idx as f64,
                Builtin::Ret => builtins.ret,
            },

            // Arithmetic
            SdeOp::Add(a, b) => registers[a] + registers[b],
            SdeOp::Sub(a, b) => registers[a] - registers[b],
            SdeOp::Mul(a, b) => registers[a] * registers[b],
            SdeOp::Div(a, b) => {
                let denom = registers[b];
                if denom.abs() < 1e-300 {
                    0.0
                } else {
                    registers[a] / denom
                }
            }
            SdeOp::Pow(a, b) => registers[a].powf(registers[b]),
            SdeOp::Neg(a) => -registers[a],

            // Math
            SdeOp::Abs(a) => registers[a].abs(),
            SdeOp::Sqrt(a) => registers[a].max(0.0).sqrt(),
            SdeOp::Log(a) => registers[a].ln(),
            SdeOp::Exp(a) => registers[a].exp(),
            SdeOp::Floor(a) => registers[a].floor(),
            SdeOp::Max(a, b) => registers[a].max(registers[b]),
            SdeOp::Min(a, b) => registers[a].min(registers[b]),

            // Comparison (1.0 = true, 0.0 = false)
            SdeOp::Gt(a, b) => bool_f64(registers[a] > registers[b]),
            SdeOp::Lt(a, b) => bool_f64(registers[a] < registers[b]),
            SdeOp::Gte(a, b) => bool_f64(registers[a] >= registers[b]),
            SdeOp::Lte(a, b) => bool_f64(registers[a] <= registers[b]),
            SdeOp::Eq(a, b) => bool_f64((registers[a] - registers[b]).abs() < 1e-12),
            SdeOp::And(a, b) => bool_f64(is_truthy(registers[a]) && is_truthy(registers[b])),
            SdeOp::Or(a, b) => bool_f64(is_truthy(registers[a]) || is_truthy(registers[b])),
            SdeOp::Not(a) => bool_f64(!is_truthy(registers[a])),
            SdeOp::IfElse {
                cond,
                then_val,
                else_val,
            } => {
                if is_truthy(registers[cond]) {
                    registers[then_val]
                } else {
                    registers[else_val]
                }
            }

            // Random
            SdeOp::RandNormal => {
                let dist = Normal::new(0.0, 1.0).unwrap();
                dist.sample(rng)
            }
            SdeOp::RandUniform => rng.random::<f64>(),
        };
    }

    // Return last register value
    if ops.is_empty() {
        0.0
    } else {
        registers[ops.len() - 1]
    }
}

#[inline(always)]
fn bool_f64(b: bool) -> f64 {
    if b { 1.0 } else { 0.0 }
}

#[inline(always)]
fn is_truthy(v: f64) -> bool {
    v.abs() > 1e-12
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn test_arithmetic_tape() {
        // Compute: param[0] - 0.5 * state[1]  (i.e. mu - 0.5*h)
        let ops = vec![
            SdeOp::LoadParam(0),  // r0 = mu
            SdeOp::Literal(0.5),  // r1 = 0.5
            SdeOp::LoadState(1),  // r2 = h
            SdeOp::Mul(1, 2),     // r3 = 0.5 * h
            SdeOp::Sub(0, 3),     // r4 = mu - 0.5*h
        ];
        let mut regs = vec![0.0; ops.len()];
        let states = [100.0, 0.04]; // S=100, h=0.04
        let params = [0.08];        // mu=0.08
        let builtins = SdeBuiltins {
            dt: 1.0 / 252.0,
            t: 0.0,
            step_idx: 0,
            ret: 0.0,
        };
        let mut rng = ChaCha8Rng::seed_from_u64(42);

        let result = eval_tape(&ops, &mut regs, &states, &params, &builtins, &mut rng);
        assert!((result - 0.06).abs() < 1e-12); // 0.08 - 0.5*0.04 = 0.06
    }

    #[test]
    fn test_sqrt_tape() {
        // sqrt(state[1])  => sqrt(0.04) = 0.2
        let ops = vec![
            SdeOp::LoadState(1), // r0 = h
            SdeOp::Sqrt(0),      // r1 = sqrt(h)
        ];
        let mut regs = vec![0.0; ops.len()];
        let states = [100.0, 0.04];
        let builtins = SdeBuiltins { dt: 0.0, t: 0.0, step_idx: 0, ret: 0.0 };
        let mut rng = ChaCha8Rng::seed_from_u64(42);

        let result = eval_tape(&ops, &mut regs, &states, &[], &builtins, &mut rng);
        assert!((result - 0.2).abs() < 1e-12);
    }

    #[test]
    fn test_if_else() {
        // if(state[0] > 100, 0.1, 0.2)
        let ops = vec![
            SdeOp::LoadState(0),  // r0 = S
            SdeOp::Literal(100.0),// r1 = 100
            SdeOp::Gt(0, 1),     // r2 = S > 100
            SdeOp::Literal(0.1), // r3
            SdeOp::Literal(0.2), // r4
            SdeOp::IfElse { cond: 2, then_val: 3, else_val: 4 }, // r5
        ];
        let mut regs = vec![0.0; ops.len()];
        let states = [105.0];
        let builtins = SdeBuiltins { dt: 0.0, t: 0.0, step_idx: 0, ret: 0.0 };
        let mut rng = ChaCha8Rng::seed_from_u64(42);

        let result = eval_tape(&ops, &mut regs, &states, &[], &builtins, &mut rng);
        assert!((result - 0.1).abs() < 1e-12);
    }

    #[test]
    fn test_rand_normal_deterministic() {
        let ops = vec![SdeOp::RandNormal];
        let mut regs = vec![0.0; 1];
        let builtins = SdeBuiltins { dt: 0.0, t: 0.0, step_idx: 0, ret: 0.0 };

        let mut rng1 = ChaCha8Rng::seed_from_u64(42);
        let v1 = eval_tape(&ops, &mut regs, &[], &[], &builtins, &mut rng1);

        let mut rng2 = ChaCha8Rng::seed_from_u64(42);
        let v2 = eval_tape(&ops, &mut regs, &[], &[], &builtins, &mut rng2);

        assert_eq!(v1, v2); // deterministic
    }
}
