//! GPU-accelerated stochastic simulation via CUDA.
//!
//! Transpiles the [`SdeOp`] tape into a CUDA kernel, compiles with NVRTC,
//! and launches on the GPU. Each CUDA thread simulates one path.

use bt_kernel::{BtError, Result};
use cudarc::driver::{CudaDevice, CudaSlice, LaunchAsync, LaunchConfig};
use cudarc::nvrtc::compile_ptx;
use std::collections::HashSet;
use std::sync::{Arc, Mutex};

use super::compiler::SdeProgram;
use super::runner::{StochasticResult, StochasticSimConfig};
use super::sde_ops::{Builtin, SdeOp};
use crate::monte_carlo::DistributionStats;

// ---------------------------------------------------------------------------
// Cached CUDA device + compiled modules
// ---------------------------------------------------------------------------

struct CudaState {
    dev: Arc<CudaDevice>,
    loaded_modules: HashSet<String>,
}

static CUDA_STATE: Mutex<Option<CudaState>> = Mutex::new(None);

/// Get or create the CUDA device, compile PTX if needed, load module if needed.
/// Returns (device, module_name) ready to get_func().
fn prepare_cuda(cuda_src: &str) -> Result<(Arc<CudaDevice>, String)> {
    use std::hash::{Hash, Hasher};
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    cuda_src.hash(&mut hasher);
    let hash = hasher.finish();
    let module_name = format!("sde_{hash:016x}");

    let mut guard = CUDA_STATE.lock().unwrap();

    // Init device on first call
    let state = if let Some(ref mut s) = *guard {
        s
    } else {
        let dev = CudaDevice::new(0)
            .map_err(|e| BtError::Strategy(format!("CUDA device init failed: {e}")))?;
        *guard = Some(CudaState {
            dev,
            loaded_modules: HashSet::new(),
        });
        guard.as_mut().unwrap()
    };

    // Compile + load module if not cached
    if !state.loaded_modules.contains(&module_name) {
        let ptx = compile_ptx(cuda_src).map_err(|e| {
            BtError::Strategy(format!(
                "CUDA kernel compilation failed: {e}\n--- kernel source ---\n{cuda_src}"
            ))
        })?;

        state
            .dev
            .load_ptx(ptx, &module_name, &["simulate"])
            .map_err(|e| BtError::Strategy(format!("CUDA module load failed: {e}")))?;

        state.loaded_modules.insert(module_name.clone());
    }

    Ok((Arc::clone(&state.dev), module_name))
}

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// Run stochastic simulation on GPU.
///
/// Falls back to CPU if `store_paths` is enabled (too much VRAM for large N).
pub fn run_stochastic_gpu(
    program: &SdeProgram,
    config: &StochasticSimConfig,
) -> Result<StochasticResult> {
    if config.store_paths {
        return Err(BtError::Strategy(
            "store_paths=True is not supported on GPU (VRAM limit). \
             Use device='cpu' for path storage or set store_paths=False."
                .into(),
        ));
    }
    if config.n_paths == 0 {
        return Err(BtError::Strategy("stochastic: n_paths must be > 0".into()));
    }
    if config.n_steps == 0 {
        return Err(BtError::Strategy("stochastic: n_steps must be > 0".into()));
    }
    if config.dt <= 0.0 {
        return Err(BtError::Strategy("stochastic: dt must be > 0".into()));
    }
    if config.s0 <= 0.0 {
        return Err(BtError::Strategy("stochastic: s0 must be > 0".into()));
    }

    // License gate
    let max_sims = bt_license::max_monte_carlo_sims();
    if config.n_paths > max_sims {
        return Err(BtError::Strategy(format!(
            "stochastic: {} paths requested but Community license allows max {}. \
             Upgrade to Pro for unlimited simulations.",
            config.n_paths, max_sims
        )));
    }

    let use_f32 = config.precision == "f32";

    // 1. Transpile + compile + load (all cached after first call)
    let cuda_src = generate_cuda_kernel(program, use_f32);
    let (dev, module_name) = prepare_cuda(&cuda_src)?;

    let n = config.n_paths;
    let base_seed = config.rng_seed.unwrap_or_else(rand::random::<u64>);
    let n_states = program.state_names.len() as i32;

    let block_size = 256u32;
    let grid_size = (n as u32 + block_size - 1) / block_size;
    let launch_cfg = LaunchConfig {
        grid_dim: (grid_size, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: 0,
    };

    let f = dev
        .get_func(&module_name, "simulate")
        .ok_or_else(|| BtError::Strategy("CUDA kernel 'simulate' not found".into()))?;

    // Branch on precision for typed GPU buffers
    if use_f32 {
        let params_f32: Vec<f32> = program.param_values.iter().map(|&v| v as f32).collect();
        let state_f32: Vec<f32> = program.state_inits.iter().map(|&v| v as f32).collect();

        let d_params: CudaSlice<f32> = dev.htod_copy(params_f32).map_err(map_cuda)?;
        let d_state: CudaSlice<f32> = dev.htod_copy(state_f32).map_err(map_cuda)?;
        let d_out: CudaSlice<f32> = dev.alloc_zeros(n * 5).map_err(map_cuda)?;

        unsafe {
            f.launch(
                launch_cfg,
                (
                    n as i32,
                    config.n_steps as i32,
                    config.dt as f32,
                    config.s0 as f32,
                    &d_params,
                    &d_state,
                    n_states,
                    &d_out,
                    base_seed,
                ),
            )
            .map_err(map_cuda)?;
        }

        let raw: Vec<f32> = dev.dtoh_sync_copy(&d_out).map_err(map_cuda)?;

        // Stats directly on f32 — avoid 50M f32→f64 widen
        let cl = &config.confidence_levels;
        let s = |offset: usize| build_stats_f32(&raw[offset * n..(offset + 1) * n], cl);

        let ((final_price, final_return), (max_drawdown, (annualized_return, annualized_vol))) =
            rayon::join(
                || rayon::join(|| s(0), || s(1)),
                || rayon::join(|| s(2), || rayon::join(|| s(3), || s(4))),
            );

        return Ok(StochasticResult {
            n_paths: n,
            n_steps: config.n_steps,
            model_name: String::new(),
            final_price,
            final_return,
            max_drawdown,
            annualized_return,
            annualized_vol,
            paths: None,
        });
    } else {
        let d_params: CudaSlice<f64> = dev
            .htod_copy(program.param_values.clone())
            .map_err(map_cuda)?;
        let d_state: CudaSlice<f64> = dev
            .htod_copy(program.state_inits.clone())
            .map_err(map_cuda)?;
        let d_out: CudaSlice<f64> = dev.alloc_zeros(n * 5).map_err(map_cuda)?;

        unsafe {
            f.launch(
                launch_cfg,
                (
                    n as i32,
                    config.n_steps as i32,
                    config.dt,
                    config.s0,
                    &d_params,
                    &d_state,
                    n_states,
                    &d_out,
                    base_seed,
                ),
            )
            .map_err(map_cuda)?;
        }

        // Stats in-place on mutable slices — zero extra copies
        let mut raw = dev.dtoh_sync_copy(&d_out).map_err(map_cuda)?;
        let cl = &config.confidence_levels;

        let (fp, rest) = raw.split_at_mut(n);
        let (fr, rest) = rest.split_at_mut(n);
        let (dd, rest) = rest.split_at_mut(n);
        let (ar, av) = rest.split_at_mut(n);

        let ((final_price, final_return), (max_drawdown, (annualized_return, annualized_vol))) =
            rayon::join(
                || rayon::join(|| build_stats_mut(fp, cl), || build_stats_mut(fr, cl)),
                || rayon::join(
                    || build_stats_mut(dd, cl),
                    || rayon::join(|| build_stats_mut(ar, cl), || build_stats_mut(av, cl)),
                ),
            );

        return Ok(StochasticResult {
            n_paths: n,
            n_steps: config.n_steps,
            model_name: String::new(),
            final_price,
            final_return,
            max_drawdown,
            annualized_return,
            annualized_vol,
            paths: None,
        });
    };
}

fn map_cuda(e: cudarc::driver::DriverError) -> BtError {
    BtError::Strategy(format!("CUDA error: {e}"))
}

// ---------------------------------------------------------------------------
// Stats (same as runner.rs)
// ---------------------------------------------------------------------------

fn build_stats(values: &[f64], confidence_levels: &[f64]) -> DistributionStats {
    use rayon::prelude::*;

    if values.is_empty() {
        return DistributionStats {
            percentiles: confidence_levels.iter().map(|&q| (q, 0.0)).collect(),
            mean: 0.0,
            std: 0.0,
            min: 0.0,
            max: 0.0,
        };
    }

    let n = values.len();
    let cmp = |a: &f64, b: &f64| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal);

    // Single parallel pass: sum + min + max
    let (sum, min, max) = values
        .par_iter()
        .fold(
            || (0.0f64, f64::INFINITY, f64::NEG_INFINITY),
            |(s, mn, mx), &v| (s + v, mn.min(v), mx.max(v)),
        )
        .reduce(
            || (0.0, f64::INFINITY, f64::NEG_INFINITY),
            |(s1, mn1, mx1), (s2, mn2, mx2)| (s1 + s2, mn1.min(mn2), mx1.max(mx2)),
        );
    let mean = sum / n as f64;
    let variance: f64 =
        values.par_iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n as f64;
    let std = variance.sqrt();

    // Percentiles via select_nth_unstable: O(n) per percentile instead of O(n log n) sort.
    // Process in ascending index order so each selection shrinks the working partition.
    let mut data = values.to_vec();

    let mut targets: Vec<(usize, f64, usize)> = confidence_levels
        .iter()
        .enumerate()
        .map(|(orig, &q)| {
            let idx = ((q * (n - 1) as f64).round() as usize).min(n - 1);
            (orig, q, idx)
        })
        .collect();
    targets.sort_by_key(|&(_, _, idx)| idx);

    let mut percentiles = vec![(0.0, 0.0); confidence_levels.len()];
    let mut lo = 0usize;
    for &(orig, q, idx) in &targets {
        if idx >= lo {
            data[lo..].select_nth_unstable_by(idx - lo, cmp);
        }
        percentiles[orig] = (q, data[idx]);
        lo = idx + 1;
    }

    DistributionStats {
        percentiles,
        mean,
        std,
        min,
        max,
    }
}

/// Stats in-place on mutable f64 slice — zero copies.
fn build_stats_mut(values: &mut [f64], confidence_levels: &[f64]) -> DistributionStats {
    use rayon::prelude::*;

    if values.is_empty() {
        return DistributionStats {
            percentiles: confidence_levels.iter().map(|&q| (q, 0.0)).collect(),
            mean: 0.0, std: 0.0, min: 0.0, max: 0.0,
        };
    }

    let n = values.len();
    let cmp = |a: &f64, b: &f64| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal);

    let (sum, min, max) = values
        .par_iter()
        .fold(
            || (0.0f64, f64::INFINITY, f64::NEG_INFINITY),
            |(s, mn, mx), &v| (s + v, mn.min(v), mx.max(v)),
        )
        .reduce(
            || (0.0, f64::INFINITY, f64::NEG_INFINITY),
            |(s1, mn1, mx1), (s2, mn2, mx2)| (s1 + s2, mn1.min(mn2), mx1.max(mx2)),
        );
    let mean = sum / n as f64;
    let variance: f64 =
        values.par_iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n as f64;
    let std = variance.sqrt();

    // select_nth directly on the mutable slice — no clone needed
    let mut targets: Vec<(usize, f64, usize)> = confidence_levels
        .iter()
        .enumerate()
        .map(|(orig, &q)| {
            let idx = ((q * (n - 1) as f64).round() as usize).min(n - 1);
            (orig, q, idx)
        })
        .collect();
    targets.sort_by_key(|&(_, _, idx)| idx);

    let mut percentiles = vec![(0.0, 0.0); confidence_levels.len()];
    let mut lo = 0usize;
    for &(orig, q, idx) in &targets {
        if idx >= lo {
            values[lo..].select_nth_unstable_by(idx - lo, cmp);
        }
        percentiles[orig] = (q, values[idx]);
        lo = idx + 1;
    }

    DistributionStats { percentiles, mean, std, min, max }
}

/// Stats computed directly on f32 data — avoids 50M f32→f64 widen.
/// select_nth on f32, results stored as f64.
fn build_stats_f32(values: &[f32], confidence_levels: &[f64]) -> DistributionStats {
    use rayon::prelude::*;

    if values.is_empty() {
        return DistributionStats {
            percentiles: confidence_levels.iter().map(|&q| (q, 0.0)).collect(),
            mean: 0.0,
            std: 0.0,
            min: 0.0,
            max: 0.0,
        };
    }

    let n = values.len();
    let cmp = |a: &f32, b: &f32| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal);

    let (sum, min, max) = values
        .par_iter()
        .fold(
            || (0.0f64, f32::INFINITY, f32::NEG_INFINITY),
            |(s, mn, mx), &v| (s + v as f64, mn.min(v), mx.max(v)),
        )
        .reduce(
            || (0.0, f32::INFINITY, f32::NEG_INFINITY),
            |(s1, mn1, mx1), (s2, mn2, mx2)| (s1 + s2, mn1.min(mn2), mx1.max(mx2)),
        );
    let mean = sum / n as f64;
    let variance: f64 =
        values.par_iter().map(|&v| (v as f64 - mean).powi(2)).sum::<f64>() / n as f64;
    let std = variance.sqrt();

    let mut data = values.to_vec();

    let mut targets: Vec<(usize, f64, usize)> = confidence_levels
        .iter()
        .enumerate()
        .map(|(orig, &q)| {
            let idx = ((q * (n - 1) as f64).round() as usize).min(n - 1);
            (orig, q, idx)
        })
        .collect();
    targets.sort_by_key(|&(_, _, idx)| idx);

    let mut percentiles = vec![(0.0, 0.0); confidence_levels.len()];
    let mut lo = 0usize;
    for &(orig, q, idx) in &targets {
        if idx >= lo {
            data[lo..].select_nth_unstable_by(idx - lo, cmp);
        }
        percentiles[orig] = (q, data[idx] as f64);
        lo = idx + 1;
    }

    DistributionStats {
        percentiles,
        mean,
        std,
        min: min as f64,
        max: max as f64,
    }
}

// ---------------------------------------------------------------------------
// CUDA kernel transpiler
// ---------------------------------------------------------------------------

/// Generate a complete CUDA kernel source from an SDE program.
fn generate_cuda_kernel(program: &SdeProgram, use_f32: bool) -> String {
    let n_states = program.state_names.len();
    let r = if use_f32 { "float" } else { "double" };
    let drift_code = transpile_tape(&program.drift, "drift", r);
    let diff_code = transpile_tape(&program.diffusion, "diff", r);

    let jump_code = if let (Some(ref ji), Some(ref js)) =
        (&program.jump_intensity, &program.jump_size)
    {
        let ji_code = transpile_tape(ji, "ji", r);
        let js_code = transpile_tape(js, "js", r);
        format!(
            r#"
        // -- Jump component --
        {ji_code}
        REAL lambda_val = ji_result;
        if (lambda_val > (REAL)0.0) {{
            REAL expected = lambda_val * dt;
            int n_jumps;
            if (expected < (REAL)0.1) {{
                n_jumps = (pcg_uniform(&rng) < expected) ? 1 : 0;
            }} else {{
                REAL L = EXP(-expected);
                int k = 0;
                REAL p = (REAL)1.0;
                do {{ k++; p *= pcg_uniform(&rng); }} while (p > L);
                n_jumps = k - 1;
            }}
            for (int jj = 0; jj < n_jumps; jj++) {{
                {js_code}
                log_ret += js_result;
            }}
        }}
"#
        )
    } else {
        String::new()
    };

    let state_update_code = program
        .state_updates
        .iter()
        .map(|(idx, tape)| {
            let label = format!("su{idx}");
            let code = transpile_tape(tape, &label, r);
            format!(
                r#"
        {{
            {code}
            state[{idx}] = {label}_result;
        }}"#
            )
        })
        .collect::<Vec<_>>()
        .join("\n");

    // f32: single pcg32 call for uniform (24-bit mantissa), faster Box-Muller
    // f64: two pcg32 calls combined (53-bit mantissa)
    let (rng_uniform, rng_normal, zero_guard) = if use_f32 {
        (
            r#"__device__ float pcg_uniform(unsigned long long* state) {
    return (float)(pcg32(state) >> 8) * (1.0f / 16777216.0f);
}"#,
            r#"__device__ float pcg_normal(unsigned long long* state) {
    float u1 = pcg_uniform(state) + 1e-30f;
    float u2 = pcg_uniform(state);
    return sqrtf(-2.0f * logf(u1)) * cosf(2.0f * 3.14159265f * u2);
}"#,
            "1e-30f",
        )
    } else {
        (
            r#"__device__ double pcg_uniform(unsigned long long* state) {
    unsigned int a = pcg32(state);
    unsigned int b = pcg32(state);
    unsigned long long combined = ((unsigned long long)a << 32) | (unsigned long long)b;
    return (double)(combined >> 11) * (1.0 / 9007199254740992.0);
}"#,
            r#"__device__ double pcg_normal(unsigned long long* state) {
    double u1 = pcg_uniform(state) + 1e-300;
    double u2 = pcg_uniform(state);
    return sqrt(-2.0 * log(u1)) * cos(2.0 * 3.14159265358979323846 * u2);
}"#,
            "1e-300",
        )
    };

    let eps_cmp = if use_f32 { "1e-6f" } else { "1e-12" };

    format!(
        r#"
// Auto-generated CUDA kernel — precision: {r}

typedef {r} REAL;
#define EXP(x) {exp_fn}(x)
#define LOG(x) {log_fn}(x)
#define SQRT(x) {sqrt_fn}(x)
#define FABS(x) {fabs_fn}(x)
#define FMAX(x,y) {fmax_fn}(x,y)
#define FMIN(x,y) {fmin_fn}(x,y)
#define FLOOR(x) {floor_fn}(x)
#define POW(x,y) {pow_fn}(x,y)
#define COS(x) {cos_fn}(x)
#define ZERO_GUARD {zero_guard}
#define EPS_CMP {eps_cmp}

__device__ unsigned int pcg32(unsigned long long* state) {{
    unsigned long long s = *state;
    *state = s * 6364136223846793005ULL + 1442695040888963407ULL;
    unsigned int xorshifted = (unsigned int)(((s >> 18u) ^ s) >> 27u);
    unsigned int rot = (unsigned int)(s >> 59u);
    return (xorshifted >> rot) | (xorshifted << ((~rot + 1u) & 31u));
}}

{rng_uniform}

{rng_normal}

extern "C" __global__ void simulate(
    int n_paths,
    int n_steps,
    REAL dt,
    REAL s0,
    const REAL* __restrict__ params,
    const REAL* __restrict__ state_inits,
    int n_states,
    REAL* __restrict__ out,
    unsigned long long base_seed
) {{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= n_paths) return;

    unsigned long long rng = (base_seed + (unsigned long long)tid) * 0x9E3779B97F4A7C15ULL;
    pcg32(&rng); pcg32(&rng); pcg32(&rng); pcg32(&rng);

    REAL state[{n_states}];
    state[0] = s0;
    for (int i = 1; i < {n_states}; i++) state[i] = state_inits[i];

    REAL sqrt_dt = SQRT(dt);
    REAL peak = s0;
    REAL max_dd = (REAL)0.0;
    REAL sum_lr_sq = (REAL)0.0;
    REAL last_ret = (REAL)0.0;

    for (int step = 0; step < n_steps; step++) {{
        REAL t = (REAL)step * dt;

        {drift_code}
        REAL drift_val = drift_result;

        {diff_code}
        REAL diff_val = diff_result;

        REAL z = pcg_normal(&rng);
        REAL log_ret = (drift_val - (REAL)0.5 * diff_val * diff_val) * dt
                     + diff_val * sqrt_dt * z;

        {jump_code}

        state[0] *= EXP(log_ret);
        last_ret = log_ret;
        sum_lr_sq += log_ret * log_ret;

        {state_update_code}

        if (state[0] > peak) peak = state[0];
        if (peak > (REAL)0.0) {{
            REAL dd = (peak - state[0]) / peak;
            if (dd > max_dd) max_dd = dd;
        }}
    }}

    REAL final_price = state[0];
    REAL ratio = final_price / s0;
    REAL total_time = (REAL)n_steps * dt;

    out[tid]               = final_price;
    out[tid + n_paths]     = ratio - (REAL)1.0;
    out[tid + 2 * n_paths] = max_dd;
    out[tid + 3 * n_paths] = (total_time > (REAL)0.0 && ratio > (REAL)0.0)
        ? POW(ratio, (REAL)1.0 / total_time) - (REAL)1.0 : (REAL)0.0;
    out[tid + 4 * n_paths] = (sum_lr_sq > (REAL)0.0 && n_steps > 1)
        ? SQRT(sum_lr_sq / (REAL)n_steps / dt) : (REAL)0.0;
}}
"#,
        r = r,
        exp_fn = if use_f32 { "expf" } else { "exp" },
        log_fn = if use_f32 { "logf" } else { "log" },
        sqrt_fn = if use_f32 { "sqrtf" } else { "sqrt" },
        fabs_fn = if use_f32 { "fabsf" } else { "fabs" },
        fmax_fn = if use_f32 { "fmaxf" } else { "fmax" },
        fmin_fn = if use_f32 { "fminf" } else { "fmin" },
        floor_fn = if use_f32 { "floorf" } else { "floor" },
        pow_fn = if use_f32 { "powf" } else { "pow" },
        cos_fn = if use_f32 { "cosf" } else { "cos" },
        zero_guard = zero_guard,
        eps_cmp = eps_cmp,
        rng_uniform = rng_uniform,
        rng_normal = rng_normal,
        n_states = n_states,
        drift_code = drift_code,
        diff_code = diff_code,
        jump_code = jump_code,
        state_update_code = state_update_code,
    )
}

/// Transpile a single SdeOp tape into CUDA C code.
///
/// Each op becomes a register assignment. The final register value is stored
/// in `{prefix}_result`.
fn transpile_tape(tape: &[SdeOp], prefix: &str, real_type: &str) -> String {
    if tape.is_empty() {
        return format!("{real_type} {prefix}_result = ({real_type})0.0;");
    }

    let mut lines = Vec::with_capacity(tape.len() + 1);

    for (i, op) in tape.iter().enumerate() {
        let reg = format!("{prefix}_r{i}");
        let expr = match op {
            // Leaf
            SdeOp::Literal(v) => format_f64(*v),
            SdeOp::LoadState(idx) => format!("state[{idx}]"),
            SdeOp::LoadParam(idx) => format!("params[{idx}]"),
            SdeOp::LoadBuiltin(b) => match b {
                Builtin::Dt => "dt".into(),
                Builtin::T => "t".into(),
                Builtin::StepIdx => format!("({real_type})step"),
                Builtin::Ret => "last_ret".into(),
            },

            // Arithmetic
            SdeOp::Add(a, b) => format!("{prefix}_r{a} + {prefix}_r{b}"),
            SdeOp::Sub(a, b) => format!("{prefix}_r{a} - {prefix}_r{b}"),
            SdeOp::Mul(a, b) => format!("{prefix}_r{a} * {prefix}_r{b}"),
            SdeOp::Div(a, b) => {
                format!("(FABS({prefix}_r{b}) < ZERO_GUARD ? ({real_type})0.0 : {prefix}_r{a} / {prefix}_r{b})")
            }
            SdeOp::Pow(a, b) => format!("POW({prefix}_r{a}, {prefix}_r{b})"),
            SdeOp::Neg(a) => format!("-{prefix}_r{a}"),

            // Math
            SdeOp::Abs(a) => format!("FABS({prefix}_r{a})"),
            SdeOp::Sqrt(a) => format!("SQRT(FMAX({prefix}_r{a}, ({real_type})0.0))"),
            SdeOp::Log(a) => format!("LOG({prefix}_r{a})"),
            SdeOp::Exp(a) => format!("EXP({prefix}_r{a})"),
            SdeOp::Floor(a) => format!("FLOOR({prefix}_r{a})"),
            SdeOp::Max(a, b) => format!("FMAX({prefix}_r{a}, {prefix}_r{b})"),
            SdeOp::Min(a, b) => format!("FMIN({prefix}_r{a}, {prefix}_r{b})"),

            // Comparison — use EPS_CMP macro
            SdeOp::Gt(a, b) => format!("({prefix}_r{a} > {prefix}_r{b} ? ({real_type})1.0 : ({real_type})0.0)"),
            SdeOp::Lt(a, b) => format!("({prefix}_r{a} < {prefix}_r{b} ? ({real_type})1.0 : ({real_type})0.0)"),
            SdeOp::Gte(a, b) => format!("({prefix}_r{a} >= {prefix}_r{b} ? ({real_type})1.0 : ({real_type})0.0)"),
            SdeOp::Lte(a, b) => format!("({prefix}_r{a} <= {prefix}_r{b} ? ({real_type})1.0 : ({real_type})0.0)"),
            SdeOp::Eq(a, b) => {
                format!("(FABS({prefix}_r{a} - {prefix}_r{b}) < EPS_CMP ? ({real_type})1.0 : ({real_type})0.0)")
            }
            SdeOp::And(a, b) => {
                format!("(FABS({prefix}_r{a}) > EPS_CMP && FABS({prefix}_r{b}) > EPS_CMP ? ({real_type})1.0 : ({real_type})0.0)")
            }
            SdeOp::Or(a, b) => {
                format!("(FABS({prefix}_r{a}) > EPS_CMP || FABS({prefix}_r{b}) > EPS_CMP ? ({real_type})1.0 : ({real_type})0.0)")
            }
            SdeOp::Not(a) => format!("(FABS({prefix}_r{a}) < EPS_CMP ? ({real_type})1.0 : ({real_type})0.0)"),
            SdeOp::IfElse { cond, then_val, else_val } => {
                format!("(FABS({prefix}_r{cond}) > EPS_CMP ? {prefix}_r{then_val} : {prefix}_r{else_val})")
            }

            // Random
            SdeOp::RandNormal => "pcg_normal(&rng)".into(),
            SdeOp::RandUniform => "pcg_uniform(&rng)".into(),
        };

        lines.push(format!("{real_type} {reg} = {expr};"));
    }

    let last = tape.len() - 1;
    lines.push(format!("{real_type} {prefix}_result = {prefix}_r{last};"));

    lines.join("\n        ")
}

/// Format an f64 as a CUDA double literal.
fn format_f64(v: f64) -> String {
    if v == f64::INFINITY {
        "INFINITY".into()
    } else if v == f64::NEG_INFINITY {
        "(-INFINITY)".into()
    } else if v.is_nan() {
        "nan(\"\")".into()
    } else {
        // Use enough precision to round-trip
        let s = format!("{v:.17e}");
        s
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stochastic::compiler::{compile_sde, SdeModelDef};
    use std::collections::HashMap;

    #[test]
    fn test_transpile_gbm_kernel() {
        let def = SdeModelDef {
            name: "gbm".into(),
            drift: "mu".into(),
            diffusion: "sigma".into(),
            jump_intensity: None,
            jump_size: None,
            state_vars: HashMap::new(),
            state_update: HashMap::new(),
            params: [("mu".into(), 0.05), ("sigma".into(), 0.2)]
                .into_iter()
                .collect(),
        };
        let program = compile_sde(&def).unwrap();
        let kernel = generate_cuda_kernel(&program, false);

        // Verify kernel contains expected components
        assert!(kernel.contains("extern \"C\" __global__ void simulate"));
        assert!(kernel.contains("pcg_normal"));
        assert!(kernel.contains("drift_result"));
        assert!(kernel.contains("diff_result"));
        assert!(kernel.contains("params["));

        // Also verify f32 kernel generates
        let kernel_f32 = generate_cuda_kernel(&program, true);
        assert!(kernel_f32.contains("typedef float REAL"));
    }

    #[test]
    fn test_transpile_garch_jd_kernel() {
        let def = SdeModelDef {
            name: "garch_jd".into(),
            drift: "mu".into(),
            diffusion: "sqrt(h)".into(),
            jump_intensity: Some("lambda".into()),
            jump_size: Some("normal(mu_j, sigma_j)".into()),
            state_vars: [("h".into(), 1e-4)].into_iter().collect(),
            state_update: [("h".into(), "omega + alpha * (ret - mu) ** 2 + beta * h".into())]
                .into_iter()
                .collect(),
            params: [
                ("mu".into(), 0.08),
                ("omega".into(), 1e-6),
                ("alpha".into(), 0.1),
                ("beta".into(), 0.85),
                ("lambda".into(), 5.0),
                ("mu_j".into(), -0.02),
                ("sigma_j".into(), 0.04),
            ]
            .into_iter()
            .collect(),
        };
        let program = compile_sde(&def).unwrap();
        let kernel = generate_cuda_kernel(&program, false);

        // Verify jump code is present
        assert!(kernel.contains("ji_result"));
        assert!(kernel.contains("js_result"));
        assert!(kernel.contains("n_jumps"));
        // Verify state update
        assert!(kernel.contains("su1_result"));
        // Verify 2 states (S, h)
        assert!(kernel.contains("double state[2]"));
    }
}
