//! Compile parsed SDE expressions into executable [`SdeOp`] tapes.
//!
//! The compiler resolves identifiers against known state variables, parameters,
//! and built-in names, then emits a flat register-machine instruction tape.

use std::collections::HashMap;

use bt_kernel::{BtError, Result};

use super::parser::{BinOp, SdeExpr};
use super::sde_ops::{Builtin, SdeOp};

// ---------------------------------------------------------------------------
// Compiled SDE program
// ---------------------------------------------------------------------------

/// A fully compiled SDE model ready for execution by the runner.
#[derive(Clone, Debug)]
pub struct SdeProgram {
    /// Drift expression tape: μ(S, t, state).
    pub drift: Vec<SdeOp>,
    /// Diffusion expression tape: σ(S, t, state).
    pub diffusion: Vec<SdeOp>,
    /// Jump intensity tape (optional): λ(S, t, state).
    pub jump_intensity: Option<Vec<SdeOp>>,
    /// Jump size tape (optional): J(S, t, state) — may contain `normal(...)`.
    pub jump_size: Option<Vec<SdeOp>>,
    /// State variable update tapes, executed after the price step.
    /// Each entry is `(state_index, tape)`.
    pub state_updates: Vec<(usize, Vec<SdeOp>)>,

    /// Ordered state variable names. Index 0 is always `"S"` (price).
    pub state_names: Vec<String>,
    /// Initial values for each state variable.
    pub state_inits: Vec<f64>,

    /// Ordered parameter names (for mapping to a flat `&[f64]` at runtime).
    pub param_names: Vec<String>,
    /// Parameter values.
    pub param_values: Vec<f64>,

    /// Maximum register count across all tapes (for pre-allocation).
    pub max_registers: usize,
}

// ---------------------------------------------------------------------------
// Compilation context
// ---------------------------------------------------------------------------

struct CompileCtx {
    /// state name → index
    state_map: HashMap<String, usize>,
    /// param name → index
    param_map: HashMap<String, usize>,
}

impl CompileCtx {
    fn new(state_names: &[String], param_names: &[String]) -> Self {
        let state_map: HashMap<String, usize> = state_names
            .iter()
            .enumerate()
            .map(|(i, n)| (n.clone(), i))
            .collect();
        let param_map: HashMap<String, usize> = param_names
            .iter()
            .enumerate()
            .map(|(i, n)| (n.clone(), i))
            .collect();
        Self {
            state_map,
            param_map,
        }
    }

    /// Resolve an identifier to an SdeOp.
    fn resolve_ident(&self, name: &str) -> Result<SdeOp> {
        // Built-ins first
        match name {
            "dt" => return Ok(SdeOp::LoadBuiltin(Builtin::Dt)),
            "t" => return Ok(SdeOp::LoadBuiltin(Builtin::T)),
            "step" | "step_idx" => return Ok(SdeOp::LoadBuiltin(Builtin::StepIdx)),
            "ret" | "log_ret" => return Ok(SdeOp::LoadBuiltin(Builtin::Ret)),
            _ => {}
        }

        // State variables (includes "S" at index 0)
        if let Some(&idx) = self.state_map.get(name) {
            return Ok(SdeOp::LoadState(idx));
        }

        // Parameters
        if let Some(&idx) = self.param_map.get(name) {
            return Ok(SdeOp::LoadParam(idx));
        }

        Err(BtError::Expression(format!(
            "unknown identifier '{name}'. Expected one of: state vars {:?}, params {:?}, or builtins (dt, t, step, ret)",
            self.state_map.keys().collect::<Vec<_>>(),
            self.param_map.keys().collect::<Vec<_>>(),
        )))
    }
}

// ---------------------------------------------------------------------------
// Expression → tape compiler
// ---------------------------------------------------------------------------

/// Emit SdeOp instructions for an expression, appending to `tape`.
/// Returns the register index of the result.
fn compile_node(expr: &SdeExpr, ctx: &CompileCtx, tape: &mut Vec<SdeOp>) -> Result<usize> {
    match expr {
        SdeExpr::Number(v) => {
            let idx = tape.len();
            tape.push(SdeOp::Literal(*v));
            Ok(idx)
        }

        SdeExpr::Ident(name) => {
            let op = ctx.resolve_ident(name)?;
            let idx = tape.len();
            tape.push(op);
            Ok(idx)
        }

        SdeExpr::Neg(inner) => {
            let a = compile_node(inner, ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Neg(a));
            Ok(idx)
        }

        SdeExpr::Not(inner) => {
            let a = compile_node(inner, ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Not(a));
            Ok(idx)
        }

        SdeExpr::BinOp { op, lhs, rhs } => {
            let a = compile_node(lhs, ctx, tape)?;
            let b = compile_node(rhs, ctx, tape)?;
            let idx = tape.len();
            let sde_op = match op {
                BinOp::Add => SdeOp::Add(a, b),
                BinOp::Sub => SdeOp::Sub(a, b),
                BinOp::Mul => SdeOp::Mul(a, b),
                BinOp::Div => SdeOp::Div(a, b),
                BinOp::Pow => SdeOp::Pow(a, b),
                BinOp::Gt => SdeOp::Gt(a, b),
                BinOp::Lt => SdeOp::Lt(a, b),
                BinOp::Gte => SdeOp::Gte(a, b),
                BinOp::Lte => SdeOp::Lte(a, b),
                BinOp::Eq => SdeOp::Eq(a, b),
                BinOp::And => SdeOp::And(a, b),
                BinOp::Or => SdeOp::Or(a, b),
            };
            tape.push(sde_op);
            Ok(idx)
        }

        SdeExpr::Call { name, args } => compile_call(name, args, ctx, tape),
    }
}

/// Compile a function call.
fn compile_call(
    name: &str,
    args: &[SdeExpr],
    ctx: &CompileCtx,
    tape: &mut Vec<SdeOp>,
) -> Result<usize> {
    match name {
        // Math functions (1 arg)
        "sqrt" => {
            expect_args(name, args, 1)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Sqrt(a));
            Ok(idx)
        }
        "abs" => {
            expect_args(name, args, 1)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Abs(a));
            Ok(idx)
        }
        "log" | "ln" => {
            expect_args(name, args, 1)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Log(a));
            Ok(idx)
        }
        "exp" => {
            expect_args(name, args, 1)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Exp(a));
            Ok(idx)
        }
        "floor" => {
            expect_args(name, args, 1)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Floor(a));
            Ok(idx)
        }

        // Math functions (2 args)
        "max" => {
            expect_args(name, args, 2)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let b = compile_node(&args[1], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Max(a, b));
            Ok(idx)
        }
        "min" => {
            expect_args(name, args, 2)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let b = compile_node(&args[1], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Min(a, b));
            Ok(idx)
        }
        "pow" => {
            expect_args(name, args, 2)?;
            let a = compile_node(&args[0], ctx, tape)?;
            let b = compile_node(&args[1], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::Pow(a, b));
            Ok(idx)
        }

        // Random: normal(mu, sigma) → mu + sigma * N(0,1)
        "normal" => {
            expect_args(name, args, 2)?;
            let mu = compile_node(&args[0], ctx, tape)?;
            let sigma = compile_node(&args[1], ctx, tape)?;
            let z_idx = tape.len();
            tape.push(SdeOp::RandNormal);
            let scaled_idx = tape.len();
            tape.push(SdeOp::Mul(sigma, z_idx));
            let result_idx = tape.len();
            tape.push(SdeOp::Add(mu, scaled_idx));
            Ok(result_idx)
        }

        // Random: uniform(lo, hi) → lo + (hi - lo) * U(0,1)
        "uniform" => {
            expect_args(name, args, 2)?;
            let lo = compile_node(&args[0], ctx, tape)?;
            let hi = compile_node(&args[1], ctx, tape)?;
            let u_idx = tape.len();
            tape.push(SdeOp::RandUniform);
            let range_idx = tape.len();
            tape.push(SdeOp::Sub(hi, lo));
            let scaled_idx = tape.len();
            tape.push(SdeOp::Mul(range_idx, u_idx));
            let result_idx = tape.len();
            tape.push(SdeOp::Add(lo, scaled_idx));
            Ok(result_idx)
        }

        // Random: randn() → N(0,1) shorthand
        "randn" => {
            expect_args(name, args, 0)?;
            let idx = tape.len();
            tape.push(SdeOp::RandNormal);
            Ok(idx)
        }

        // Conditional: if(cond, then, else)
        "if" => {
            expect_args(name, args, 3)?;
            let c = compile_node(&args[0], ctx, tape)?;
            let t = compile_node(&args[1], ctx, tape)?;
            let e = compile_node(&args[2], ctx, tape)?;
            let idx = tape.len();
            tape.push(SdeOp::IfElse {
                cond: c,
                then_val: t,
                else_val: e,
            });
            Ok(idx)
        }

        _ => Err(BtError::Expression(format!(
            "unknown function '{name}'. Available: sqrt, abs, log, ln, exp, floor, \
             max, min, pow, normal, uniform, randn, if"
        ))),
    }
}

fn expect_args(name: &str, args: &[SdeExpr], expected: usize) -> Result<()> {
    if args.len() != expected {
        Err(BtError::Expression(format!(
            "{name}() expects {expected} argument(s), got {}",
            args.len()
        )))
    } else {
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// SDE model definition as raw strings (from JSON/Python).
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct SdeModelDef {
    /// Human-readable name.
    #[serde(default = "default_name")]
    pub name: String,
    /// Drift expression string: μ(S, t, state, params).
    pub drift: String,
    /// Diffusion expression string: σ(S, t, state, params).
    pub diffusion: String,
    /// Jump intensity expression (optional): λ(S, t, state, params).
    #[serde(default)]
    pub jump_intensity: Option<String>,
    /// Jump size expression (optional): J — may use `normal(mu, sigma)`.
    #[serde(default)]
    pub jump_size: Option<String>,
    /// State variables with initial values. `"S"` is implicit (index 0).
    #[serde(default)]
    pub state_vars: HashMap<String, f64>,
    /// State variable update expressions. Key = var name, value = expression.
    #[serde(default)]
    pub state_update: HashMap<String, String>,
    /// Model parameters (name → value).
    #[serde(default)]
    pub params: HashMap<String, f64>,
}

fn default_name() -> String {
    "custom".into()
}

/// Compile an [`SdeModelDef`] into an executable [`SdeProgram`].
pub fn compile_sde(def: &SdeModelDef) -> Result<SdeProgram> {
    // Build ordered state names: S is always index 0
    let mut state_names = vec!["S".to_string()];
    let mut state_inits = vec![0.0]; // S0 set by runner at runtime

    // Sort for deterministic ordering
    let mut user_states: Vec<_> = def.state_vars.iter().collect();
    user_states.sort_by_key(|(k, _)| (*k).clone());
    for (name, init) in user_states {
        if name == "S" {
            continue; // already added
        }
        state_names.push(name.clone());
        state_inits.push(*init);
    }

    // Build ordered param names
    let mut param_entries: Vec<_> = def.params.iter().collect();
    param_entries.sort_by_key(|(k, _)| (*k).clone());
    let param_names: Vec<String> = param_entries.iter().map(|(k, _)| (*k).clone()).collect();
    let param_values: Vec<f64> = param_entries.iter().map(|(_, v)| **v).collect();

    let ctx = CompileCtx::new(&state_names, &param_names);

    // Parse and compile each expression
    let drift = compile_expr_str(&def.drift, &ctx, "drift")?;
    let diffusion = compile_expr_str(&def.diffusion, &ctx, "diffusion")?;

    let jump_intensity = match &def.jump_intensity {
        Some(s) => Some(compile_expr_str(s, &ctx, "jump_intensity")?),
        None => None,
    };
    let jump_size = match &def.jump_size {
        Some(s) => Some(compile_expr_str(s, &ctx, "jump_size")?),
        None => None,
    };

    // Validate: if jump_intensity is set, jump_size must be too
    if jump_intensity.is_some() && jump_size.is_none() {
        return Err(BtError::Expression(
            "jump_intensity is set but jump_size is missing".into(),
        ));
    }

    // Compile state updates
    let mut state_updates = Vec::new();
    for (name, expr_str) in &def.state_update {
        let idx = ctx
            .state_map
            .get(name)
            .copied()
            .ok_or_else(|| {
                BtError::Expression(format!(
                    "state_update references unknown state variable '{name}'"
                ))
            })?;
        if idx == 0 {
            return Err(BtError::Expression(
                "cannot define state_update for 'S' — price is updated by the Euler-Maruyama step"
                    .into(),
            ));
        }
        let tape = compile_expr_str(expr_str, &ctx, &format!("state_update[{name}]"))?;
        state_updates.push((idx, tape));
    }

    // Compute max registers needed
    let max_registers = [
        drift.len(),
        diffusion.len(),
        jump_intensity.as_ref().map(|t| t.len()).unwrap_or(0),
        jump_size.as_ref().map(|t| t.len()).unwrap_or(0),
    ]
    .into_iter()
    .chain(state_updates.iter().map(|(_, t)| t.len()))
    .max()
    .unwrap_or(1);

    Ok(SdeProgram {
        drift,
        diffusion,
        jump_intensity,
        jump_size,
        state_updates,
        state_names,
        state_inits,
        param_names,
        param_values,
        max_registers,
    })
}

fn compile_expr_str(input: &str, ctx: &CompileCtx, label: &str) -> Result<Vec<SdeOp>> {
    let ast = super::parser::parse_sde_expr(input).map_err(|e| {
        BtError::Expression(format!("error parsing {label} expression '{input}': {e}"))
    })?;
    let mut tape = Vec::new();
    compile_node(&ast, ctx, &mut tape)?;
    Ok(tape)
}

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

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

    fn make_gbm_def() -> 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.08), ("sigma".into(), 0.2)]
                .into_iter()
                .collect(),
        }
    }

    fn make_garch_jd_def() -> 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(),
        }
    }

    #[test]
    fn test_compile_gbm() {
        let def = make_gbm_def();
        let prog = compile_sde(&def).unwrap();

        assert_eq!(prog.state_names, vec!["S"]);
        assert!(prog.jump_intensity.is_none());
        assert!(prog.state_updates.is_empty());
        assert_eq!(prog.drift.len(), 1); // single LoadParam
        assert_eq!(prog.diffusion.len(), 1);
    }

    #[test]
    fn test_compile_garch_jd() {
        let def = make_garch_jd_def();
        let prog = compile_sde(&def).unwrap();

        assert_eq!(prog.state_names.len(), 2); // S, h
        assert!(prog.state_names.contains(&"S".to_string()));
        assert!(prog.state_names.contains(&"h".to_string()));
        assert!(prog.jump_intensity.is_some());
        assert!(prog.jump_size.is_some());
        assert_eq!(prog.state_updates.len(), 1);
    }

    #[test]
    fn test_jump_intensity_without_size_errors() {
        let mut def = make_gbm_def();
        def.jump_intensity = Some("5.0".into());
        assert!(compile_sde(&def).is_err());
    }

    #[test]
    fn test_unknown_ident_errors() {
        let mut def = make_gbm_def();
        def.drift = "unknown_var".into();
        assert!(compile_sde(&def).is_err());
    }

    #[test]
    fn test_state_update_for_s_errors() {
        let mut def = make_gbm_def();
        def.state_update.insert("S".into(), "S * 2".into());
        assert!(compile_sde(&def).is_err());
    }
}
