use std::collections::{BTreeMap, HashMap};

use bt_kernel::{BtError, Result};

use crate::expr::{DynPeriod, Expr, ScalarValue};

#[derive(Clone, Debug, PartialEq)]
pub enum CompiledOp {
    Column(String),
    Literal(ScalarValue),
    Parameter(String),

    Add,
    Sub,
    Mul,
    Div,

    Gt,
    Lt,
    Eq,
    And,
    Or,
    Not,

    Lag {
        periods: usize,
    },
    Lead {
        periods: usize,
    },
    RollingMean {
        window: usize,
    },
    RollingStd {
        window: usize,
    },
    RollingSum {
        window: usize,
    },
    RollingMin {
        window: usize,
    },
    RollingMax {
        window: usize,
    },
    EwmMean {
        span: f64,
    },
    Diff {
        periods: usize,
    },
    PctChange {
        periods: usize,
    },
    CumSum,
    CumProd,
    Rank,
    ZScore {
        window: usize,
    },
    Rsi {
        period: usize,
    },
    Atr {
        period: usize,
    },
    LinRegSlope {
        window: usize,
    },
    LinRegValue {
        window: usize,
    },
    LinRegR2 {
        window: usize,
    },

    // -- Moving Averages --
    Dema {
        period: usize,
    },
    Tema {
        period: usize,
    },
    Wma {
        window: usize,
    },
    Hma {
        window: usize,
    },
    Kama {
        period: usize,
    },

    // -- Momentum --
    Roc {
        period: usize,
    },
    Macd {
        fast: usize,
        slow: usize,
    },
    MacdSignal {
        fast: usize,
        slow: usize,
        signal: usize,
    },
    MacdHist {
        fast: usize,
        slow: usize,
        signal: usize,
    },
    StochK {
        period: usize,
    },
    WilliamsR {
        period: usize,
    },
    Cci {
        period: usize,
    },
    Adx {
        period: usize,
    },

    // -- Volatility --
    TrueRange,
    Natr {
        period: usize,
    },
    BollingerUpper {
        window: usize,
        num_std: f64,
    },
    BollingerLower {
        window: usize,
        num_std: f64,
    },
    BollingerWidth {
        window: usize,
        num_std: f64,
    },
    KeltnerUpper {
        period: usize,
        multiplier: f64,
    },
    KeltnerLower {
        period: usize,
        multiplier: f64,
    },
    SuperTrend {
        period: usize,
        multiplier: f64,
    },

    // -- Volume --
    Obv,
    Vwap,
    AdLine,
    Mfi {
        period: usize,
    },

    // -- Signals --
    CrossAbove,
    CrossBelow,

    // -- Statistics --
    RollingMedian {
        window: usize,
    },

    // -- Trend --
    ParabolicSar {
        af_step: f64,
        af_max: f64,
    },

    IfElse,

    CrossSectionalMean,
    CrossSectionalRank,

    SymbolRef {
        symbol: String,
    },

    Hour,
    Minute,
    DayOfWeek,
    Month,
    DayOfMonth,

    Function {
        name: String,
    },

    Scan(CompiledScan),
}

/// Scalar instruction for the scan VM — executed once per row.
#[derive(Clone, Debug, PartialEq)]
pub enum ScanOp {
    Literal(f64),
    /// Load pre-evaluated data column at current row: data_columns[i][t]
    LoadData(usize),
    /// Load resolved parameter value: param_values[i]
    LoadParam(usize),
    /// Load state variable value from previous row: prev_state[i]
    LoadPrev(usize),
    /// Load a register computed earlier in the current step: registers[i]
    LoadVar(usize),
    Add(usize, usize),
    Sub(usize, usize),
    Mul(usize, usize),
    Div(usize, usize),
    Neg(usize),
    Abs(usize),
    Sqrt(usize),
    Log(usize),
    Exp(usize),
    Max(usize, usize),
    Min(usize, usize),
    Gt(usize, usize),
    Lt(usize, usize),
    Eq(usize, usize),
    And(usize, usize),
    Or(usize, usize),
    Not(usize),
    IfElse {
        cond: usize,
        then_val: usize,
        else_val: usize,
    },
}

/// A compiled scan (fold) — all metadata needed by the evaluator VM.
#[derive(Clone, Debug, PartialEq)]
pub struct CompiledScan {
    pub num_states: usize,
    /// Indices in the outer DAG for initial state values.
    pub init_node_indices: Vec<usize>,
    /// Column names whose pre-evaluated arrays are fed as data to the VM.
    pub data_column_names: Vec<String>,
    /// Indices in the outer DAG for each data column (pre-evaluated).
    pub data_node_indices: Vec<usize>,
    /// Parameter names referenced inside the scan body.
    pub param_names: Vec<String>,
    /// Flat instruction tape for the scan VM.
    pub instructions: Vec<ScanOp>,
    /// Register index that holds the result of each update step.
    pub step_result_registers: Vec<usize>,
    /// (update_step_idx, state_idx) — which update steps write back to state.
    pub state_writeback: Vec<(usize, usize)>,
    /// Register index whose value is the scan output at each row.
    pub output_register: usize,
    /// Total number of registers needed by the VM.
    pub num_registers: usize,
}

#[derive(Clone, Debug, PartialEq)]
pub struct CompiledNode {
    pub op: CompiledOp,
    pub inputs: Vec<usize>,
}

#[derive(Clone, Debug, PartialEq)]
pub struct CompiledExpr {
    pub nodes: Vec<CompiledNode>,
    pub output_indices: Vec<usize>,
    pub output_names: Vec<String>,
}

impl CompiledExpr {
    pub fn single_output_index(&self) -> Result<usize> {
        if self.output_indices.len() != 1 {
            return Err(BtError::Expression(format!(
                "expected exactly one output, got {}",
                self.output_indices.len()
            )));
        }

        Ok(self.output_indices[0])
    }
}

pub fn compile_expr(expr: Expr) -> Result<CompiledExpr> {
    let mut named = HashMap::new();
    named.insert("output".to_string(), expr);
    compile_named_exprs(&named)
}

#[tracing::instrument(level = "debug", skip_all)]
pub fn compile_named_exprs(exprs: &HashMap<String, Expr>) -> Result<CompiledExpr> {
    if exprs.is_empty() {
        return Err(BtError::Expression(
            "cannot compile empty expression set".to_string(),
        ));
    }

    let mut nodes = Vec::new();
    let mut cse = HashMap::new();
    let mut output_indices = Vec::new();
    let mut output_names = Vec::new();

    let sorted: BTreeMap<String, Expr> = exprs
        .iter()
        .map(|(name, expr)| (name.clone(), expr.clone()))
        .collect();

    for (name, expr) in sorted {
        let folded = constant_fold(expr);
        let index = compile_node(folded, &mut nodes, &mut cse)?;
        output_names.push(name);
        output_indices.push(index);
    }

    Ok(CompiledExpr {
        nodes,
        output_indices,
        output_names,
    })
}

fn compile_node(
    expr: Expr,
    nodes: &mut Vec<CompiledNode>,
    cse: &mut HashMap<String, usize>,
) -> Result<usize> {
    // Special case: Scan bypasses normal decompose
    if let Expr::Scan { .. } = &expr {
        return compile_scan_node(expr, nodes, cse);
    }

    let (op, children) = decompose(expr);
    let mut inputs = Vec::with_capacity(children.len());
    for child in children {
        inputs.push(compile_node(child, nodes, cse)?);
    }

    let key = format!("{op:?}|{inputs:?}");
    if let Some(existing) = cse.get(&key) {
        return Ok(*existing);
    }

    let index = nodes.len();
    nodes.push(CompiledNode {
        op: op.clone(),
        inputs,
    });
    cse.insert(key, index);
    Ok(index)
}

// ---------------------------------------------------------------------------
// Scan compiler
// ---------------------------------------------------------------------------

fn compile_scan_node(
    expr: Expr,
    nodes: &mut Vec<CompiledNode>,
    cse: &mut HashMap<String, usize>,
) -> Result<usize> {
    let (state_names, init_exprs, update_names, update_exprs, output) = match expr {
        Expr::Scan {
            state_names,
            init_exprs,
            update_names,
            update_exprs,
            output,
        } => (state_names, init_exprs, update_names, update_exprs, output),
        _ => unreachable!(),
    };

    let num_states = state_names.len();

    // 1. Compile init exprs as outer DAG nodes
    let mut init_node_indices = Vec::with_capacity(num_states);
    for init_expr in &init_exprs {
        let idx = compile_node(*init_expr.clone(), nodes, cse)?;
        init_node_indices.push(idx);
    }

    // 2. Collect data columns and parameters referenced in update_exprs
    let mut data_col_set: Vec<String> = Vec::new();
    let mut param_set: Vec<String> = Vec::new();
    for ue in &update_exprs {
        collect_scan_data_columns(ue, &mut data_col_set);
        collect_scan_parameters(ue, &mut param_set);
    }
    // Also collect from init_exprs (they reference columns too)
    for ie in &init_exprs {
        collect_scan_data_columns(ie, &mut data_col_set);
    }
    data_col_set.dedup();
    param_set.dedup();

    // 3. Compile data columns as outer DAG nodes
    let mut data_node_indices = Vec::with_capacity(data_col_set.len());
    for col_name in &data_col_set {
        let idx = compile_node(Expr::Column(col_name.clone()), nodes, cse)?;
        data_node_indices.push(idx);
    }

    // 4. Build the scan instruction tape
    let mut instructions: Vec<ScanOp> = Vec::new();
    let mut step_result_registers: Vec<usize> = Vec::new();
    let mut state_writeback: Vec<(usize, usize)> = Vec::new();
    let mut hoisted_exprs: Vec<Expr> = Vec::new();
    let base_data_count = data_col_set.len();

    // Map update_name -> step index for ScanVar resolution
    let mut update_name_to_step: HashMap<String, usize> = HashMap::new();

    for (step_idx, ue) in update_exprs.iter().enumerate() {
        let reg = lower_scan_expr(
            ue,
            &state_names,
            &update_names,
            &update_name_to_step,
            &step_result_registers,
            &data_col_set,
            &param_set,
            &mut instructions,
            &mut hoisted_exprs,
            base_data_count,
        )?;
        step_result_registers.push(reg);
        update_name_to_step.insert(update_names[step_idx].clone(), step_idx);

        // Check if this update step writes back to a state variable
        if let Some(state_idx) = state_names
            .iter()
            .position(|n| n == &update_names[step_idx])
        {
            state_writeback.push((step_idx, state_idx));
        }
    }

    // 4b. Compile hoisted expressions as outer DAG nodes
    for (i, hoisted_expr) in hoisted_exprs.iter().enumerate() {
        let idx = compile_node(hoisted_expr.clone(), nodes, cse)?;
        data_node_indices.push(idx);
        data_col_set.push(format!("__hoisted_{i}"));
    }

    // 5. Find the output register
    let output_register = if let Some(step_idx) = update_name_to_step.get(&output) {
        step_result_registers[*step_idx]
    } else {
        return Err(BtError::Expression(format!(
            "scan output '{}' is not in update_names: {:?}",
            output, update_names
        )));
    };

    let num_registers = instructions.len();

    let scan = CompiledScan {
        num_states,
        init_node_indices: init_node_indices.clone(),
        data_column_names: data_col_set,
        data_node_indices: data_node_indices.clone(),
        param_names: param_set,
        instructions,
        step_result_registers,
        state_writeback,
        output_register,
        num_registers,
    };

    // Collect all outer DAG inputs: init nodes + data nodes
    let mut all_inputs = init_node_indices;
    all_inputs.extend(data_node_indices);

    let index = nodes.len();
    nodes.push(CompiledNode {
        op: CompiledOp::Scan(scan),
        inputs: all_inputs,
    });
    Ok(index)
}

/// Check if an expression tree contains any ScanPrev or ScanVar references.
fn contains_scan_refs(expr: &Expr) -> bool {
    match expr {
        Expr::ScanPrev(_) | Expr::ScanVar(_) => true,
        Expr::Column(_) | Expr::Literal(_) | Expr::Parameter(_) => false,
        Expr::Add(l, r)
        | Expr::Sub(l, r)
        | Expr::Mul(l, r)
        | Expr::Div(l, r)
        | Expr::Gt(l, r)
        | Expr::Lt(l, r)
        | Expr::Eq(l, r)
        | Expr::And(l, r)
        | Expr::Or(l, r) => contains_scan_refs(l) || contains_scan_refs(r),
        Expr::Not(inner) => contains_scan_refs(inner),
        Expr::IfElse(c, t, e) => {
            contains_scan_refs(c) || contains_scan_refs(t) || contains_scan_refs(e)
        }
        Expr::Function(_, args) => args.iter().any(contains_scan_refs),
        // A nested Scan is a self-contained computation — its internal
        // ScanPrev/ScanVar are scoped to it and don't reference outer state.
        // So from the outer scan's perspective, it has no scan refs.
        Expr::Scan { .. } => false,
        Expr::Lag(inner, _)
        | Expr::Lead(inner, _)
        | Expr::RollingMean(inner, _)
        | Expr::RollingStd(inner, _)
        | Expr::RollingSum(inner, _)
        | Expr::RollingMin(inner, _)
        | Expr::RollingMax(inner, _)
        | Expr::EwmMean(inner, _)
        | Expr::Diff(inner, _)
        | Expr::PctChange(inner, _)
        | Expr::CumSum(inner)
        | Expr::CumProd(inner)
        | Expr::Rank(inner)
        | Expr::ZScore(inner, _)
        | Expr::Rsi(inner, _)
        | Expr::LinRegSlope(inner, _)
        | Expr::LinRegValue(inner, _)
        | Expr::LinRegR2(inner, _)
        | Expr::Dema(inner, _)
        | Expr::Tema(inner, _)
        | Expr::Wma(inner, _)
        | Expr::Hma(inner, _)
        | Expr::Kama(inner, _)
        | Expr::Roc(inner, _)
        | Expr::Macd(inner, _, _)
        | Expr::MacdSignal(inner, _, _, _)
        | Expr::MacdHist(inner, _, _, _)
        | Expr::BollingerUpper(inner, _, _)
        | Expr::BollingerLower(inner, _, _)
        | Expr::BollingerWidth(inner, _, _)
        | Expr::RollingMedian(inner, _)
        | Expr::CrossSectionalMean(inner)
        | Expr::CrossSectionalRank(inner)
        | Expr::SymbolRef(_, inner)
        | Expr::Hour(inner)
        | Expr::Minute(inner)
        | Expr::DayOfWeek(inner)
        | Expr::Month(inner)
        | Expr::DayOfMonth(inner) => contains_scan_refs(inner),
        Expr::CrossAbove(a, b) | Expr::CrossBelow(a, b) | Expr::Obv(a, b) => {
            contains_scan_refs(a) || contains_scan_refs(b)
        }
        Expr::Atr(h, l, c, _)
        | Expr::TrueRange(h, l, c)
        | Expr::Natr(h, l, c, _)
        | Expr::StochK(h, l, c, _)
        | Expr::WilliamsR(h, l, c, _)
        | Expr::Cci(h, l, c, _)
        | Expr::Adx(h, l, c, _)
        | Expr::SuperTrend(h, l, c, _, _)
        | Expr::KeltnerUpper(h, l, c, _, _)
        | Expr::KeltnerLower(h, l, c, _, _) => {
            contains_scan_refs(h) || contains_scan_refs(l) || contains_scan_refs(c)
        }
        Expr::Vwap(h, l, c, v) | Expr::AdLine(h, l, c, v) | Expr::Mfi(h, l, c, v, _) => {
            contains_scan_refs(h)
                || contains_scan_refs(l)
                || contains_scan_refs(c)
                || contains_scan_refs(v)
        }
        Expr::ParabolicSar(h, l, _, _) => contains_scan_refs(h) || contains_scan_refs(l),
    }
}

/// Recursively lower a scan update expression to ScanOp instructions.
/// Returns the register index holding this sub-expression's result.
///
/// Sub-expressions that are not natively supported by the scan VM but do NOT
/// reference any `ScanPrev`/`ScanVar` are "hoisted" — pushed into
/// `hoisted_exprs` so they can be compiled as outer DAG nodes and referenced
/// via `LoadData`.
fn lower_scan_expr(
    expr: &Expr,
    state_names: &[String],
    update_names: &[String],
    update_name_to_step: &HashMap<String, usize>,
    step_result_registers: &[usize],
    data_columns: &[String],
    param_names: &[String],
    instructions: &mut Vec<ScanOp>,
    hoisted_exprs: &mut Vec<Expr>,
    base_data_count: usize,
) -> Result<usize> {
    match expr {
        Expr::ScanPrev(name) => {
            let state_idx = state_names.iter().position(|n| n == name).ok_or_else(|| {
                BtError::Expression(format!(
                    "ScanPrev('{}') references unknown state variable; known: {:?}",
                    name, state_names
                ))
            })?;
            let reg = instructions.len();
            instructions.push(ScanOp::LoadPrev(state_idx));
            Ok(reg)
        }
        Expr::ScanVar(name) => {
            let step_idx = update_name_to_step.get(name).ok_or_else(|| {
                BtError::Expression(format!(
                    "ScanVar('{}') references variable not yet computed in this scan step; \
                     ensure it appears earlier in the update dict. Known: {:?}",
                    name,
                    update_name_to_step.keys().collect::<Vec<_>>()
                ))
            })?;
            let source_reg = step_result_registers[*step_idx];
            let reg = instructions.len();
            instructions.push(ScanOp::LoadVar(source_reg));
            Ok(reg)
        }
        Expr::Column(name) => {
            let col_idx = data_columns.iter().position(|n| n == name).ok_or_else(|| {
                BtError::Expression(format!(
                    "scan references column '{}' but it wasn't collected; known: {:?}",
                    name, data_columns
                ))
            })?;
            let reg = instructions.len();
            instructions.push(ScanOp::LoadData(col_idx));
            Ok(reg)
        }
        Expr::Parameter(name) => {
            let param_idx = param_names.iter().position(|n| n == name).ok_or_else(|| {
                BtError::Expression(format!(
                    "scan references parameter '{}' but it wasn't collected; known: {:?}",
                    name, param_names
                ))
            })?;
            let reg = instructions.len();
            instructions.push(ScanOp::LoadParam(param_idx));
            Ok(reg)
        }
        Expr::Literal(sv) => {
            let val = sv.as_f64().ok_or_else(|| {
                BtError::Expression(format!(
                    "scan body only supports numeric literals, got {:?}",
                    sv
                ))
            })?;
            let reg = instructions.len();
            instructions.push(ScanOp::Literal(val));
            Ok(reg)
        }
        Expr::Add(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Add(l, r));
            Ok(reg)
        }
        Expr::Sub(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Sub(l, r));
            Ok(reg)
        }
        Expr::Mul(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Mul(l, r));
            Ok(reg)
        }
        Expr::Div(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Div(l, r));
            Ok(reg)
        }
        Expr::Gt(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Gt(l, r));
            Ok(reg)
        }
        Expr::Lt(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Lt(l, r));
            Ok(reg)
        }
        Expr::Eq(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Eq(l, r));
            Ok(reg)
        }
        Expr::And(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::And(l, r));
            Ok(reg)
        }
        Expr::Or(lhs, rhs) => {
            let l = lower_scan_expr(
                lhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let r = lower_scan_expr(
                rhs,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Or(l, r));
            Ok(reg)
        }
        Expr::Not(inner) => {
            let i = lower_scan_expr(
                inner,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::Not(i));
            Ok(reg)
        }
        Expr::IfElse(cond, then_val, else_val) => {
            let c = lower_scan_expr(
                cond,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let t = lower_scan_expr(
                then_val,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let e = lower_scan_expr(
                else_val,
                state_names,
                update_names,
                update_name_to_step,
                step_result_registers,
                data_columns,
                param_names,
                instructions,
                hoisted_exprs,
                base_data_count,
            )?;
            let reg = instructions.len();
            instructions.push(ScanOp::IfElse {
                cond: c,
                then_val: t,
                else_val: e,
            });
            Ok(reg)
        }
        Expr::Function(name, args) => {
            // Support built-in math functions inside scan
            match (name.as_str(), args.len()) {
                ("abs", 1) => {
                    let a = lower_scan_expr(
                        &args[0],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let reg = instructions.len();
                    instructions.push(ScanOp::Abs(a));
                    Ok(reg)
                }
                ("sqrt", 1) => {
                    let a = lower_scan_expr(
                        &args[0],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let reg = instructions.len();
                    instructions.push(ScanOp::Sqrt(a));
                    Ok(reg)
                }
                ("log", 1) => {
                    let a = lower_scan_expr(
                        &args[0],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let reg = instructions.len();
                    instructions.push(ScanOp::Log(a));
                    Ok(reg)
                }
                ("exp", 1) => {
                    let a = lower_scan_expr(
                        &args[0],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let reg = instructions.len();
                    instructions.push(ScanOp::Exp(a));
                    Ok(reg)
                }
                ("max", 2) => {
                    let a = lower_scan_expr(
                        &args[0],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let b = lower_scan_expr(
                        &args[1],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let reg = instructions.len();
                    instructions.push(ScanOp::Max(a, b));
                    Ok(reg)
                }
                ("min", 2) => {
                    let a = lower_scan_expr(
                        &args[0],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let b = lower_scan_expr(
                        &args[1],
                        state_names,
                        update_names,
                        update_name_to_step,
                        step_result_registers,
                        data_columns,
                        param_names,
                        instructions,
                        hoisted_exprs,
                        base_data_count,
                    )?;
                    let reg = instructions.len();
                    instructions.push(ScanOp::Min(a, b));
                    Ok(reg)
                }
                _ => Err(BtError::Expression(format!(
                    "function '{}' with {} args is not supported inside scan body; \
                     only scalar math functions (abs, sqrt, log, exp, max, min) are allowed",
                    name,
                    args.len()
                ))),
            }
        }
        // Any other expression type (windowed ops, nested scans, etc.)
        // If it doesn't reference ScanPrev/ScanVar, hoist it to the outer DAG
        // and reference the pre-evaluated result via LoadData.
        _ => {
            if contains_scan_refs(expr) {
                return Err(BtError::Expression(format!(
                    "expression type {:?} contains ScanPrev/ScanVar references \
                     but is not a supported scan primitive; windowed/vectorized \
                     operations cannot depend on scan state",
                    std::mem::discriminant(expr)
                )));
            }
            let hoisted_idx = hoisted_exprs.len();
            hoisted_exprs.push(expr.clone());
            let reg = instructions.len();
            instructions.push(ScanOp::LoadData(base_data_count + hoisted_idx));
            Ok(reg)
        }
    }
}

/// Collect column names from an expression, skipping ScanPrev/ScanVar.
fn collect_scan_data_columns(expr: &Expr, out: &mut Vec<String>) {
    match expr {
        Expr::Column(name) => {
            if !out.contains(name) {
                out.push(name.clone());
            }
        }
        Expr::ScanPrev(_) | Expr::ScanVar(_) | Expr::Literal(_) | Expr::Parameter(_) => {}
        Expr::Add(l, r)
        | Expr::Sub(l, r)
        | Expr::Mul(l, r)
        | Expr::Div(l, r)
        | Expr::Gt(l, r)
        | Expr::Lt(l, r)
        | Expr::Eq(l, r)
        | Expr::And(l, r)
        | Expr::Or(l, r) => {
            collect_scan_data_columns(l, out);
            collect_scan_data_columns(r, out);
        }
        Expr::Not(inner) => collect_scan_data_columns(inner, out),
        Expr::IfElse(c, t, e) => {
            collect_scan_data_columns(c, out);
            collect_scan_data_columns(t, out);
            collect_scan_data_columns(e, out);
        }
        Expr::Function(_, args) => {
            for a in args {
                collect_scan_data_columns(a, out);
            }
        }
        _ => {} // windowed ops etc. — will be rejected at lower time
    }
}

/// Collect parameter names from an expression, skipping ScanPrev/ScanVar.
fn collect_scan_parameters(expr: &Expr, out: &mut Vec<String>) {
    match expr {
        Expr::Parameter(name) => {
            if !out.contains(name) {
                out.push(name.clone());
            }
        }
        Expr::ScanPrev(_) | Expr::ScanVar(_) | Expr::Literal(_) | Expr::Column(_) => {}
        Expr::Add(l, r)
        | Expr::Sub(l, r)
        | Expr::Mul(l, r)
        | Expr::Div(l, r)
        | Expr::Gt(l, r)
        | Expr::Lt(l, r)
        | Expr::Eq(l, r)
        | Expr::And(l, r)
        | Expr::Or(l, r) => {
            collect_scan_parameters(l, out);
            collect_scan_parameters(r, out);
        }
        Expr::Not(inner) => collect_scan_parameters(inner, out),
        Expr::IfElse(c, t, e) => {
            collect_scan_parameters(c, out);
            collect_scan_parameters(t, out);
            collect_scan_parameters(e, out);
        }
        Expr::Function(_, args) => {
            for a in args {
                collect_scan_parameters(a, out);
            }
        }
        _ => {}
    }
}

fn decompose(expr: Expr) -> (CompiledOp, Vec<Expr>) {
    match expr {
        Expr::Column(name) => (CompiledOp::Column(name), Vec::new()),
        Expr::Literal(value) => (CompiledOp::Literal(value), Vec::new()),
        Expr::Parameter(name) => (CompiledOp::Parameter(name), Vec::new()),

        Expr::Add(lhs, rhs) => (CompiledOp::Add, vec![*lhs, *rhs]),
        Expr::Sub(lhs, rhs) => (CompiledOp::Sub, vec![*lhs, *rhs]),
        Expr::Mul(lhs, rhs) => (CompiledOp::Mul, vec![*lhs, *rhs]),
        Expr::Div(lhs, rhs) => (CompiledOp::Div, vec![*lhs, *rhs]),

        Expr::Gt(lhs, rhs) => (CompiledOp::Gt, vec![*lhs, *rhs]),
        Expr::Lt(lhs, rhs) => (CompiledOp::Lt, vec![*lhs, *rhs]),
        Expr::Eq(lhs, rhs) => (CompiledOp::Eq, vec![*lhs, *rhs]),
        Expr::And(lhs, rhs) => (CompiledOp::And, vec![*lhs, *rhs]),
        Expr::Or(lhs, rhs) => (CompiledOp::Or, vec![*lhs, *rhs]),
        Expr::Not(value) => (CompiledOp::Not, vec![*value]),

        Expr::Lag(value, periods) => (
            CompiledOp::Lag {
                periods: periods
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Lead(value, periods) => (
            CompiledOp::Lead {
                periods: periods
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::RollingMean(value, window) => (
            CompiledOp::RollingMean {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::RollingStd(value, window) => (
            CompiledOp::RollingStd {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::RollingSum(value, window) => (
            CompiledOp::RollingSum {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::RollingMin(value, window) => (
            CompiledOp::RollingMin {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::RollingMax(value, window) => (
            CompiledOp::RollingMax {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::EwmMean(value, span) => (
            CompiledOp::EwmMean {
                span: span.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*value],
        ),
        Expr::Diff(value, periods) => (
            CompiledOp::Diff {
                periods: periods
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::PctChange(value, periods) => (
            CompiledOp::PctChange {
                periods: periods
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::CumSum(value) => (CompiledOp::CumSum, vec![*value]),
        Expr::CumProd(value) => (CompiledOp::CumProd, vec![*value]),
        Expr::Rank(value) => (CompiledOp::Rank, vec![*value]),
        Expr::ZScore(value, window) => (
            CompiledOp::ZScore {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Rsi(value, period) => (
            CompiledOp::Rsi {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Atr(high, low, close, period) => (
            CompiledOp::Atr {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*high, *low, *close],
        ),
        Expr::LinRegSlope(value, window) => (
            CompiledOp::LinRegSlope {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::LinRegValue(value, window) => (
            CompiledOp::LinRegValue {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::LinRegR2(value, window) => (
            CompiledOp::LinRegR2 {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),

        // -- Moving Averages --
        Expr::Dema(value, period) => (
            CompiledOp::Dema {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Tema(value, period) => (
            CompiledOp::Tema {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Wma(value, window) => (
            CompiledOp::Wma {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Hma(value, window) => (
            CompiledOp::Hma {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Kama(value, period) => (
            CompiledOp::Kama {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),

        // -- Momentum --
        Expr::Roc(value, period) => (
            CompiledOp::Roc {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::Macd(value, fast, slow) => (
            CompiledOp::Macd {
                fast: fast
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                slow: slow
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::MacdSignal(value, fast, slow, signal) => (
            CompiledOp::MacdSignal {
                fast: fast
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                slow: slow
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                signal: signal
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::MacdHist(value, fast, slow, signal) => (
            CompiledOp::MacdHist {
                fast: fast
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                slow: slow
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                signal: signal
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),
        Expr::StochK(high, low, close, period) => (
            CompiledOp::StochK {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*high, *low, *close],
        ),
        Expr::WilliamsR(high, low, close, period) => (
            CompiledOp::WilliamsR {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*high, *low, *close],
        ),
        Expr::Cci(high, low, close, period) => (
            CompiledOp::Cci {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*high, *low, *close],
        ),
        Expr::Adx(high, low, close, period) => (
            CompiledOp::Adx {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*high, *low, *close],
        ),

        // -- Volatility --
        Expr::TrueRange(high, low, close) => (CompiledOp::TrueRange, vec![*high, *low, *close]),
        Expr::Natr(high, low, close, period) => (
            CompiledOp::Natr {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*high, *low, *close],
        ),
        Expr::BollingerUpper(value, window, num_std) => (
            CompiledOp::BollingerUpper {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                num_std: num_std.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*value],
        ),
        Expr::BollingerLower(value, window, num_std) => (
            CompiledOp::BollingerLower {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                num_std: num_std.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*value],
        ),
        Expr::BollingerWidth(value, window, num_std) => (
            CompiledOp::BollingerWidth {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                num_std: num_std.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*value],
        ),
        Expr::KeltnerUpper(high, low, close, period, multiplier) => (
            CompiledOp::KeltnerUpper {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                multiplier: multiplier.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*high, *low, *close],
        ),
        Expr::KeltnerLower(high, low, close, period, multiplier) => (
            CompiledOp::KeltnerLower {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                multiplier: multiplier.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*high, *low, *close],
        ),
        Expr::SuperTrend(high, low, close, period, multiplier) => (
            CompiledOp::SuperTrend {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
                multiplier: multiplier.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*high, *low, *close],
        ),

        // -- Volume --
        Expr::Obv(close, volume) => (CompiledOp::Obv, vec![*close, *volume]),
        Expr::Vwap(high, low, close, volume) => {
            (CompiledOp::Vwap, vec![*high, *low, *close, *volume])
        }
        Expr::AdLine(high, low, close, volume) => {
            (CompiledOp::AdLine, vec![*high, *low, *close, *volume])
        }
        Expr::Mfi(high, low, close, volume, period) => (
            CompiledOp::Mfi {
                period: period
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*high, *low, *close, *volume],
        ),

        // -- Signals --
        Expr::CrossAbove(a, b) => (CompiledOp::CrossAbove, vec![*a, *b]),
        Expr::CrossBelow(a, b) => (CompiledOp::CrossBelow, vec![*a, *b]),

        // -- Statistics --
        Expr::RollingMedian(value, window) => (
            CompiledOp::RollingMedian {
                window: window
                    .as_fixed()
                    .expect("unresolved DynPeriod -- call resolve_dynamic_periods first"),
            },
            vec![*value],
        ),

        // -- Trend --
        Expr::ParabolicSar(high, low, af_step, af_max) => (
            CompiledOp::ParabolicSar {
                af_step: af_step.as_fixed().expect("unresolved DynFloat"),
                af_max: af_max.as_fixed().expect("unresolved DynFloat"),
            },
            vec![*high, *low],
        ),

        Expr::IfElse(condition, then_value, else_value) => (
            CompiledOp::IfElse,
            vec![*condition, *then_value, *else_value],
        ),

        Expr::CrossSectionalMean(value) => (CompiledOp::CrossSectionalMean, vec![*value]),
        Expr::CrossSectionalRank(value) => (CompiledOp::CrossSectionalRank, vec![*value]),
        Expr::SymbolRef(symbol, inner) => (CompiledOp::SymbolRef { symbol }, vec![*inner]),

        Expr::Hour(value) => (CompiledOp::Hour, vec![*value]),
        Expr::Minute(value) => (CompiledOp::Minute, vec![*value]),
        Expr::DayOfWeek(value) => (CompiledOp::DayOfWeek, vec![*value]),
        Expr::Month(value) => (CompiledOp::Month, vec![*value]),
        Expr::DayOfMonth(value) => (CompiledOp::DayOfMonth, vec![*value]),

        Expr::Function(name, args) => (CompiledOp::Function { name }, args),

        // Scan is handled by compile_scan_node before decompose is called.
        Expr::Scan { .. } => unreachable!("Scan should be intercepted by compile_node"),
        Expr::ScanPrev(name) => {
            panic!("ScanPrev('{name}') found outside of Scan context — this is a bug")
        }
        Expr::ScanVar(name) => {
            panic!("ScanVar('{name}') found outside of Scan context — this is a bug")
        }
    }
}

fn constant_fold(expr: Expr) -> Expr {
    match expr {
        Expr::Add(lhs, rhs) => fold_binary_numeric(*lhs, *rhs, |a, b| a + b, Expr::Add),
        Expr::Sub(lhs, rhs) => fold_binary_numeric(*lhs, *rhs, |a, b| a - b, Expr::Sub),
        Expr::Mul(lhs, rhs) => fold_binary_numeric(*lhs, *rhs, |a, b| a * b, Expr::Mul),
        Expr::Div(lhs, rhs) => fold_binary_numeric(
            *lhs,
            *rhs,
            |a, b| if b == 0.0 { f64::NAN } else { a / b },
            Expr::Div,
        ),
        Expr::Gt(lhs, rhs) => fold_binary_cmp(*lhs, *rhs, |a, b| a > b, Expr::Gt),
        Expr::Lt(lhs, rhs) => fold_binary_cmp(*lhs, *rhs, |a, b| a < b, Expr::Lt),
        Expr::Eq(lhs, rhs) => {
            let lhs = constant_fold(*lhs);
            let rhs = constant_fold(*rhs);
            match (&lhs, &rhs) {
                (Expr::Literal(lhs), Expr::Literal(rhs)) => {
                    Expr::Literal(ScalarValue::Bool(lhs == rhs))
                }
                _ => Expr::Eq(Box::new(lhs), Box::new(rhs)),
            }
        }
        Expr::And(lhs, rhs) => fold_binary_bool(*lhs, *rhs, |a, b| a && b, Expr::And),
        Expr::Or(lhs, rhs) => fold_binary_bool(*lhs, *rhs, |a, b| a || b, Expr::Or),
        Expr::Not(value) => {
            let value = constant_fold(*value);
            match value {
                Expr::Literal(ScalarValue::Bool(inner)) => Expr::Literal(ScalarValue::Bool(!inner)),
                _ => Expr::Not(Box::new(value)),
            }
        }
        Expr::IfElse(condition, then_value, else_value) => {
            let condition = constant_fold(*condition);
            let then_value = constant_fold(*then_value);
            let else_value = constant_fold(*else_value);
            match condition {
                Expr::Literal(ScalarValue::Bool(value)) => {
                    if value {
                        then_value
                    } else {
                        else_value
                    }
                }
                _ => Expr::IfElse(
                    Box::new(condition),
                    Box::new(then_value),
                    Box::new(else_value),
                ),
            }
        }
        Expr::Lag(value, DynPeriod::Fixed(0)) => constant_fold(*value),
        Expr::Lead(value, DynPeriod::Fixed(0)) => constant_fold(*value),
        Expr::Lag(value, periods) => Expr::Lag(Box::new(constant_fold(*value)), periods),
        Expr::Lead(value, periods) => Expr::Lead(Box::new(constant_fold(*value)), periods),
        Expr::RollingMean(value, window) => {
            Expr::RollingMean(Box::new(constant_fold(*value)), window)
        }
        Expr::RollingStd(value, window) => {
            Expr::RollingStd(Box::new(constant_fold(*value)), window)
        }
        Expr::RollingSum(value, window) => {
            Expr::RollingSum(Box::new(constant_fold(*value)), window)
        }
        Expr::RollingMin(value, window) => {
            Expr::RollingMin(Box::new(constant_fold(*value)), window)
        }
        Expr::RollingMax(value, window) => {
            Expr::RollingMax(Box::new(constant_fold(*value)), window)
        }
        Expr::EwmMean(value, span) => Expr::EwmMean(Box::new(constant_fold(*value)), span),
        Expr::Diff(value, periods) => Expr::Diff(Box::new(constant_fold(*value)), periods),
        Expr::PctChange(value, periods) => {
            Expr::PctChange(Box::new(constant_fold(*value)), periods)
        }
        Expr::CumSum(value) => Expr::CumSum(Box::new(constant_fold(*value))),
        Expr::CumProd(value) => Expr::CumProd(Box::new(constant_fold(*value))),
        Expr::Rank(value) => Expr::Rank(Box::new(constant_fold(*value))),
        Expr::ZScore(value, window) => Expr::ZScore(Box::new(constant_fold(*value)), window),
        Expr::Rsi(value, period) => Expr::Rsi(Box::new(constant_fold(*value)), period),
        Expr::Atr(high, low, close, period) => Expr::Atr(
            Box::new(constant_fold(*high)),
            Box::new(constant_fold(*low)),
            Box::new(constant_fold(*close)),
            period,
        ),
        Expr::LinRegSlope(value, window) => {
            Expr::LinRegSlope(Box::new(constant_fold(*value)), window)
        }
        Expr::LinRegValue(value, window) => {
            Expr::LinRegValue(Box::new(constant_fold(*value)), window)
        }
        Expr::LinRegR2(value, window) => Expr::LinRegR2(Box::new(constant_fold(*value)), window),
        Expr::CrossSectionalMean(value) => {
            Expr::CrossSectionalMean(Box::new(constant_fold(*value)))
        }
        Expr::CrossSectionalRank(value) => {
            Expr::CrossSectionalRank(Box::new(constant_fold(*value)))
        }
        Expr::SymbolRef(symbol, inner) => Expr::SymbolRef(symbol, Box::new(constant_fold(*inner))),
        Expr::Hour(value) => Expr::Hour(Box::new(constant_fold(*value))),
        Expr::Minute(value) => Expr::Minute(Box::new(constant_fold(*value))),
        Expr::DayOfWeek(value) => Expr::DayOfWeek(Box::new(constant_fold(*value))),
        Expr::Month(value) => Expr::Month(Box::new(constant_fold(*value))),
        Expr::DayOfMonth(value) => Expr::DayOfMonth(Box::new(constant_fold(*value))),
        Expr::Function(name, args) => {
            Expr::Function(name, args.into_iter().map(constant_fold).collect())
        }
        Expr::Scan {
            state_names,
            init_exprs,
            update_names,
            update_exprs,
            output,
        } => Expr::Scan {
            state_names,
            init_exprs: init_exprs
                .into_iter()
                .map(|e| Box::new(constant_fold(*e)))
                .collect(),
            update_names,
            update_exprs: update_exprs
                .into_iter()
                .map(|e| Box::new(constant_fold(*e)))
                .collect(),
            output,
        },
        Expr::ScanPrev(_) | Expr::ScanVar(_) => expr,

        // -- New single-input indicators --
        Expr::Dema(v, p) => Expr::Dema(Box::new(constant_fold(*v)), p),
        Expr::Tema(v, p) => Expr::Tema(Box::new(constant_fold(*v)), p),
        Expr::Wma(v, w) => Expr::Wma(Box::new(constant_fold(*v)), w),
        Expr::Hma(v, w) => Expr::Hma(Box::new(constant_fold(*v)), w),
        Expr::Kama(v, p) => Expr::Kama(Box::new(constant_fold(*v)), p),
        Expr::Roc(v, p) => Expr::Roc(Box::new(constant_fold(*v)), p),
        Expr::Macd(v, f, s) => Expr::Macd(Box::new(constant_fold(*v)), f, s),
        Expr::MacdSignal(v, f, s, sig) => Expr::MacdSignal(Box::new(constant_fold(*v)), f, s, sig),
        Expr::MacdHist(v, f, s, sig) => Expr::MacdHist(Box::new(constant_fold(*v)), f, s, sig),
        Expr::BollingerUpper(v, w, n) => Expr::BollingerUpper(Box::new(constant_fold(*v)), w, n),
        Expr::BollingerLower(v, w, n) => Expr::BollingerLower(Box::new(constant_fold(*v)), w, n),
        Expr::BollingerWidth(v, w, n) => Expr::BollingerWidth(Box::new(constant_fold(*v)), w, n),
        Expr::RollingMedian(v, w) => Expr::RollingMedian(Box::new(constant_fold(*v)), w),

        // -- New two-input indicators --
        Expr::CrossAbove(a, b) => {
            Expr::CrossAbove(Box::new(constant_fold(*a)), Box::new(constant_fold(*b)))
        }
        Expr::CrossBelow(a, b) => {
            Expr::CrossBelow(Box::new(constant_fold(*a)), Box::new(constant_fold(*b)))
        }
        Expr::Obv(c, v) => Expr::Obv(Box::new(constant_fold(*c)), Box::new(constant_fold(*v))),

        // -- New HLC indicators --
        Expr::TrueRange(h, l, c) => Expr::TrueRange(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
        ),
        Expr::Natr(h, l, c, p) => Expr::Natr(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
        ),
        Expr::StochK(h, l, c, p) => Expr::StochK(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
        ),
        Expr::WilliamsR(h, l, c, p) => Expr::WilliamsR(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
        ),
        Expr::Cci(h, l, c, p) => Expr::Cci(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
        ),
        Expr::Adx(h, l, c, p) => Expr::Adx(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
        ),
        Expr::SuperTrend(h, l, c, p, m) => Expr::SuperTrend(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
            m,
        ),
        Expr::KeltnerUpper(h, l, c, p, m) => Expr::KeltnerUpper(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
            m,
        ),
        Expr::KeltnerLower(h, l, c, p, m) => Expr::KeltnerLower(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            p,
            m,
        ),

        // -- New HLCV indicators --
        Expr::Vwap(h, l, c, v) => Expr::Vwap(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            Box::new(constant_fold(*v)),
        ),
        Expr::AdLine(h, l, c, v) => Expr::AdLine(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            Box::new(constant_fold(*v)),
        ),
        Expr::Mfi(h, l, c, v, p) => Expr::Mfi(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            Box::new(constant_fold(*c)),
            Box::new(constant_fold(*v)),
            p,
        ),

        // -- New HL-only indicators --
        Expr::ParabolicSar(h, l, s, m) => Expr::ParabolicSar(
            Box::new(constant_fold(*h)),
            Box::new(constant_fold(*l)),
            s,
            m,
        ),

        Expr::Column(_) | Expr::Literal(_) | Expr::Parameter(_) => expr,
    }
}

fn fold_binary_numeric(
    lhs: Expr,
    rhs: Expr,
    op: impl Fn(f64, f64) -> f64,
    make_expr: impl Fn(Box<Expr>, Box<Expr>) -> Expr,
) -> Expr {
    let lhs = constant_fold(lhs);
    let rhs = constant_fold(rhs);
    match (&lhs, &rhs) {
        (Expr::Literal(lhs), Expr::Literal(rhs)) => match (lhs.as_f64(), rhs.as_f64()) {
            (Some(lhs), Some(rhs)) => Expr::Literal(ScalarValue::Float64(op(lhs, rhs))),
            _ => make_expr(
                Box::new(Expr::Literal(lhs.clone())),
                Box::new(Expr::Literal(rhs.clone())),
            ),
        },
        _ => make_expr(Box::new(lhs), Box::new(rhs)),
    }
}

fn fold_binary_cmp(
    lhs: Expr,
    rhs: Expr,
    op: impl Fn(f64, f64) -> bool,
    make_expr: impl Fn(Box<Expr>, Box<Expr>) -> Expr,
) -> Expr {
    let lhs = constant_fold(lhs);
    let rhs = constant_fold(rhs);
    match (&lhs, &rhs) {
        (Expr::Literal(lhs), Expr::Literal(rhs)) => match (lhs.as_f64(), rhs.as_f64()) {
            (Some(lhs), Some(rhs)) => Expr::Literal(ScalarValue::Bool(op(lhs, rhs))),
            _ => make_expr(
                Box::new(Expr::Literal(lhs.clone())),
                Box::new(Expr::Literal(rhs.clone())),
            ),
        },
        _ => make_expr(Box::new(lhs), Box::new(rhs)),
    }
}

fn fold_binary_bool(
    lhs: Expr,
    rhs: Expr,
    op: impl Fn(bool, bool) -> bool,
    make_expr: impl Fn(Box<Expr>, Box<Expr>) -> Expr,
) -> Expr {
    let lhs = constant_fold(lhs);
    let rhs = constant_fold(rhs);
    match (&lhs, &rhs) {
        (Expr::Literal(ScalarValue::Bool(lhs)), Expr::Literal(ScalarValue::Bool(rhs))) => {
            Expr::Literal(ScalarValue::Bool(op(*lhs, *rhs)))
        }
        _ => make_expr(Box::new(lhs), Box::new(rhs)),
    }
}

// ---------------------------------------------------------------------------
// Cross-sectional helpers
// ---------------------------------------------------------------------------

/// Returns `true` if any node in the compiled expression is a cross-sectional op.
pub fn has_cross_sectional_ops(expr: &CompiledExpr) -> bool {
    expr.nodes.iter().any(|node| {
        matches!(
            node.op,
            CompiledOp::CrossSectionalMean | CompiledOp::CrossSectionalRank
        )
    })
}

/// Describes a cross-sectional operation within a compiled expression.
#[derive(Clone, Debug, PartialEq)]
pub struct CrossSectionInfo {
    /// The kind of cross-sectional operation.
    pub op: CrossSectionKind,
    /// The name of the input column (signal or bar column) that the op wraps.
    pub input_column: String,
}

#[derive(Clone, Debug, PartialEq)]
pub enum CrossSectionKind {
    Mean,
    Rank,
}

/// Extract cross-section info from a compiled expression that contains exactly
/// one cross-sectional op wrapping a single Column node.
///
/// Returns `None` if the expression contains no cross-sectional ops.
/// Returns `Err` if the cross-sectional op's input is not a plain Column reference.
pub fn extract_cross_section_info(expr: &CompiledExpr) -> Result<Option<CrossSectionInfo>> {
    for node in &expr.nodes {
        let kind = match &node.op {
            CompiledOp::CrossSectionalMean => CrossSectionKind::Mean,
            CompiledOp::CrossSectionalRank => CrossSectionKind::Rank,
            _ => continue,
        };

        if node.inputs.len() != 1 {
            return Err(BtError::Expression(
                "cross-sectional op must have exactly one input".to_string(),
            ));
        }

        let input_node = &expr.nodes[node.inputs[0]];
        match &input_node.op {
            CompiledOp::Column(name) => {
                return Ok(Some(CrossSectionInfo {
                    op: kind,
                    input_column: name.clone(),
                }));
            }
            _ => {
                return Err(BtError::Expression(
                    "cross-sectional op argument must be a column or signal name, \
                     not a sub-expression; define the sub-expression as a separate signal"
                        .to_string(),
                ));
            }
        }
    }

    Ok(None)
}

// ---------------------------------------------------------------------------
// Symbol-ref helpers
// ---------------------------------------------------------------------------

/// Returns `true` if any node in the compiled expression is a SymbolRef op.
pub fn has_symbol_ref_ops(expr: &CompiledExpr) -> bool {
    expr.nodes
        .iter()
        .any(|node| matches!(node.op, CompiledOp::SymbolRef { .. }))
}

/// Describes a symbol-reference operation within a compiled expression.
#[derive(Clone, Debug, PartialEq)]
pub struct SymbolRefInfo {
    /// The target symbol name to reference.
    pub symbol: String,
    /// The name of the column to look up from the target symbol's environment.
    pub input_column: String,
}

/// Extract symbol-ref info from a compiled expression whose **root** node
/// is a `SymbolRef` op wrapping a single `Column` node (a "pure" symbol-ref).
///
/// Returns `None` if the root node is not a `SymbolRef` (even if the expression
/// contains `SymbolRef` ops deeper in the tree — those are "composite" and
/// must be handled differently by the orchestrator).
pub fn extract_symbol_ref_info(expr: &CompiledExpr) -> Result<Option<SymbolRefInfo>> {
    // Only check the root (last) node — not arbitrary nodes in the tree.
    let root = match expr.nodes.last() {
        Some(node) => node,
        None => return Ok(None),
    };

    let symbol = match &root.op {
        CompiledOp::SymbolRef { symbol } => symbol.clone(),
        _ => return Ok(None), // Root is not SymbolRef → not a pure symbol-ref signal.
    };

    if root.inputs.len() != 1 {
        return Err(BtError::Expression(
            "SymbolRef op must have exactly one input".to_string(),
        ));
    }

    let input_node = &expr.nodes[root.inputs[0]];
    match &input_node.op {
        CompiledOp::Column(name) => Ok(Some(SymbolRefInfo {
            symbol,
            input_column: name.clone(),
        })),
        _ => Err(BtError::Expression(
            "SymbolRef argument must be a column or signal name, \
             not a sub-expression; define the sub-expression as a separate signal"
                .to_string(),
        )),
    }
}
