//! Built-in SDE model presets.
//!
//! Each preset returns an [`SdeModelDef`] with expression strings and default
//! parameters that the user can override.

use std::collections::HashMap;

use bt_kernel::{BtError, Result};

use super::compiler::SdeModelDef;

/// Return a preset model definition by name.
///
/// The caller can override `params` after receiving the definition.
/// Available presets: `"gbm"`, `"heston"`, `"merton"`, `"garch_jd"`.
pub fn preset(name: &str) -> Result<SdeModelDef> {
    match name {
        "gbm" => Ok(gbm()),
        "heston" => Ok(heston()),
        "merton" => Ok(merton()),
        "garch_jd" => Ok(garch_jd()),
        _ => Err(BtError::Strategy(format!(
            "unknown stochastic preset '{name}'. Available: gbm, heston, merton, garch_jd"
        ))),
    }
}

/// List all available preset names.
pub fn available_presets() -> &'static [&'static str] {
    &["gbm", "heston", "merton", "garch_jd"]
}

// ---------------------------------------------------------------------------
// Geometric Brownian Motion
// ---------------------------------------------------------------------------
// dS = μ·S·dt + σ·S·dW
//
// Simplest model. Constant drift and volatility.

fn gbm() -> SdeModelDef {
    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(),
    }
}

// ---------------------------------------------------------------------------
// Heston Stochastic Volatility
// ---------------------------------------------------------------------------
// dS  = μ·S·dt + √v·S·dW₁
// dv  = κ(θ − v)·dt + ξ·√v·dW₂
// corr(dW₁, dW₂) = ρ
//
// The state update for v uses an independent normal draw (randn()) to
// represent dW₂. Correlation with dW₁ (the Brownian used in the price step)
// is achieved via the Cholesky decomposition:
//   dW₂ = ρ·dW₁ + √(1−ρ²)·dZ  where dZ ⊥ dW₁
//
// NOTE: In this DSL the price-step Brownian (dW₁) is generated internally by
// the runner and is not directly accessible in expressions. To get correlated
// vol dynamics we use `ret` (the realised log-return) as a proxy for the
// direction of the price shock, and add an independent noise term.
//
// For a more rigorous correlation structure the user can override with their
// own expressions.

fn heston() -> SdeModelDef {
    SdeModelDef {
        name: "heston".into(),
        drift: "mu".into(),
        diffusion: "sqrt(v)".into(),
        jump_intensity: None,
        jump_size: None,
        state_vars: [("v".into(), 0.04)].into_iter().collect(),
        state_update: [(
            "v".into(),
            // Euler-Maruyama for the variance process:
            // v_{t+1} = v_t + κ(θ − v_t)·dt + ξ·√v_t·√dt·Z₂
            // Clamped to 0 via max() to prevent negative variance.
            "max(v + kappa * (theta - v) * dt + xi * sqrt(v) * sqrt(dt) * randn(), 0.0)".into(),
        )]
        .into_iter()
        .collect(),
        params: [
            ("mu".into(), 0.05),
            ("kappa".into(), 2.0),
            ("theta".into(), 0.04),
            ("xi".into(), 0.3),
        ]
        .into_iter()
        .collect(),
    }
}

// ---------------------------------------------------------------------------
// Merton Jump Diffusion
// ---------------------------------------------------------------------------
// dS = μ·S·dt + σ·S·dW + J·S·dN
//
// N is a Poisson process with intensity λ.
// J ~ Normal(μ_J, σ_J) are log-jump sizes.

fn merton() -> SdeModelDef {
    SdeModelDef {
        name: "merton".into(),
        drift: "mu".into(),
        diffusion: "sigma".into(),
        jump_intensity: Some("lambda".into()),
        jump_size: Some("normal(mu_j, sigma_j)".into()),
        state_vars: HashMap::new(),
        state_update: HashMap::new(),
        params: [
            ("mu".into(), 0.05),
            ("sigma".into(), 0.2),
            ("lambda".into(), 1.0),
            ("mu_j".into(), -0.05),
            ("sigma_j".into(), 0.08),
        ]
        .into_iter()
        .collect(),
    }
}

// ---------------------------------------------------------------------------
// GARCH(1,1) Jump Diffusion
// ---------------------------------------------------------------------------
// dS = μ·S·dt + √h·S·dW + J·S·dN
//
// h_{t+1} = ω + α·(r_t − μ)² + β·h_t     (GARCH variance update)
// N ~ Poisson(λ), J ~ Normal(μ_J, σ_J)

fn garch_jd() -> SdeModelDef {
    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(),
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stochastic::compiler::compile_sde;

    #[test]
    fn test_all_presets_compile() {
        for name in available_presets() {
            let def = preset(name).unwrap();
            let program = compile_sde(&def);
            assert!(
                program.is_ok(),
                "preset '{name}' failed to compile: {:?}",
                program.err()
            );
        }
    }

    #[test]
    fn test_unknown_preset_errors() {
        assert!(preset("foo").is_err());
    }

    #[test]
    fn test_preset_param_override() {
        let mut def = preset("gbm").unwrap();
        def.params.insert("mu".into(), 0.10);
        def.params.insert("sigma".into(), 0.30);
        let program = compile_sde(&def).unwrap();
        // Find mu param index and check value
        let mu_idx = program
            .param_names
            .iter()
            .position(|n| n == "mu")
            .unwrap();
        assert!((program.param_values[mu_idx] - 0.10).abs() < 1e-12);
    }
}
