use std::borrow::Cow;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, RwLock};

use arrow::array::{
    Array, ArrayRef, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, StringArray,
    TimestampNanosecondArray, UInt32Array, UInt64Array, UInt8Array,
};
use bt_kernel::{BtError, Result};
use chrono::{DateTime, Datelike, Timelike};

use crate::compiler::{CompiledExpr, CompiledOp, CompiledScan, ScanOp};
use crate::expr::ScalarValue;
use crate::functions::FunctionRegistry;

/// Cache key: (op discriminant + period/params, input data pointer).
/// The data pointer identifies which symbol's column is being processed —
/// same Arc<dyn Array> pointer = same underlying data = same result.
type CacheKey = (u64, usize);

/// Thread-safe indicator result cache for batch/sweep operations.
/// First evaluation computes and stores; subsequent evaluations hit the cache.
///
/// The cache is keyed by a hash of (op_type, period, input_array_pointer).
/// Since all strategies in a batch share the same aligned data (same Arc pointers),
/// EMA(12) on symbol 1's close will always produce the same cache key.
pub struct IndicatorCache {
    store: RwLock<HashMap<CacheKey, ArrayRef>>,
}

impl Default for IndicatorCache {
    fn default() -> Self {
        Self::new()
    }
}

impl IndicatorCache {
    pub fn new() -> Self {
        Self {
            store: RwLock::new(HashMap::new()),
        }
    }

    fn make_key(op: &CompiledOp, inputs: &[ArrayRef]) -> CacheKey {
        use std::hash::{Hash, Hasher};
        let mut hasher = std::collections::hash_map::DefaultHasher::new();

        // Hash the op discriminant + params
        std::mem::discriminant(op).hash(&mut hasher);
        // Hash op-specific params (period, span, etc.)
        match op {
            CompiledOp::EwmMean { span } => span.to_bits().hash(&mut hasher),
            CompiledOp::RollingMean { window }
            | CompiledOp::RollingStd { window }
            | CompiledOp::RollingSum { window }
            | CompiledOp::RollingMin { window }
            | CompiledOp::RollingMax { window }
            | CompiledOp::Wma { window }
            | CompiledOp::Hma { window } => window.hash(&mut hasher),
            CompiledOp::ZScore { window } => window.hash(&mut hasher),
            CompiledOp::Rsi { period }
            | CompiledOp::Dema { period }
            | CompiledOp::Tema { period }
            | CompiledOp::Kama { period } => period.hash(&mut hasher),
            CompiledOp::Roc { period } => period.hash(&mut hasher),
            CompiledOp::Atr { period } => period.hash(&mut hasher),
            CompiledOp::LinRegSlope { window }
            | CompiledOp::LinRegValue { window }
            | CompiledOp::LinRegR2 { window } => window.hash(&mut hasher),
            CompiledOp::Lag { periods } | CompiledOp::Lead { periods } => periods.hash(&mut hasher),
            CompiledOp::Diff { periods } | CompiledOp::PctChange { periods } => {
                periods.hash(&mut hasher)
            }
            CompiledOp::RollingMedian { window } => window.hash(&mut hasher),
            _ => {}
        }
        let op_hash = hasher.finish();

        // Use input data pointer(s) as identity — same Arc = same data
        let mut data_hash: usize = 0;
        for input in inputs {
            data_hash ^= Arc::as_ptr(input) as *const () as usize;
        }

        (op_hash, data_hash)
    }

    fn is_cacheable(op: &CompiledOp) -> bool {
        matches!(
            op,
            CompiledOp::EwmMean { .. }
                | CompiledOp::RollingMean { .. }
                | CompiledOp::RollingStd { .. }
                | CompiledOp::RollingSum { .. }
                | CompiledOp::RollingMin { .. }
                | CompiledOp::RollingMax { .. }
                | CompiledOp::ZScore { .. }
                | CompiledOp::Rsi { .. }
                | CompiledOp::Dema { .. }
                | CompiledOp::Tema { .. }
                | CompiledOp::Wma { .. }
                | CompiledOp::Hma { .. }
                | CompiledOp::Kama { .. }
                | CompiledOp::Roc { .. }
                | CompiledOp::Atr { .. }
                | CompiledOp::LinRegSlope { .. }
                | CompiledOp::LinRegValue { .. }
                | CompiledOp::LinRegR2 { .. }
                | CompiledOp::Lag { .. }
                | CompiledOp::Lead { .. }
                | CompiledOp::Diff { .. }
                | CompiledOp::PctChange { .. }
                | CompiledOp::RollingMedian { .. }
        )
    }

    fn get(&self, key: &CacheKey) -> Option<ArrayRef> {
        self.store.read().ok()?.get(key).cloned()
    }

    fn insert(&self, key: CacheKey, value: ArrayRef) {
        if let Ok(mut w) = self.store.write() {
            w.insert(key, value);
        }
    }
}

pub trait ExprEvaluator: Send + Sync {
    fn evaluate(
        &self,
        expr: &CompiledExpr,
        inputs: &HashMap<String, ArrayRef>,
        params: &HashMap<String, ScalarValue>,
    ) -> Result<HashMap<String, ArrayRef>>;
}

pub struct DefaultExprEvaluator {
    functions: FunctionRegistry,
    cache: Option<Arc<IndicatorCache>>,
}

impl Clone for DefaultExprEvaluator {
    fn clone(&self) -> Self {
        Self {
            functions: self.functions.clone(),
            cache: self.cache.clone(),
        }
    }
}

impl Default for DefaultExprEvaluator {
    fn default() -> Self {
        Self::new()
    }
}

impl DefaultExprEvaluator {
    pub fn new() -> Self {
        Self {
            functions: FunctionRegistry::with_builtins(),
            cache: None,
        }
    }

    /// Create an evaluator with a shared indicator cache (for batch/sweep).
    pub fn with_cache(cache: Arc<IndicatorCache>) -> Self {
        Self {
            functions: FunctionRegistry::with_builtins(),
            cache: Some(cache),
        }
    }
}

impl ExprEvaluator for DefaultExprEvaluator {
    #[tracing::instrument(level = "trace", skip_all)]
    fn evaluate(
        &self,
        expr: &CompiledExpr,
        inputs: &HashMap<String, ArrayRef>,
        params: &HashMap<String, ScalarValue>,
    ) -> Result<HashMap<String, ArrayRef>> {
        if expr.nodes.is_empty() {
            return Err(BtError::Expression(
                "cannot evaluate empty compiled expression".to_string(),
            ));
        }

        let len = validate_and_resolve_len(inputs)?;
        let mut values: Vec<Option<ArrayRef>> = vec![None; expr.nodes.len()];

        for (index, node) in expr.nodes.iter().enumerate() {
            let mut input_values = Vec::with_capacity(node.inputs.len());
            for input_index in &node.inputs {
                let input = values
                    .get(*input_index)
                    .and_then(Clone::clone)
                    .ok_or_else(|| {
                        BtError::Expression(format!(
                            "invalid compiled graph: missing input node {input_index}"
                        ))
                    })?;
                input_values.push(input);
            }

            let value = self.evaluate_node(&node.op, &input_values, inputs, params, len)?;
            if value.len() != len {
                return Err(BtError::Expression(format!(
                    "expression node {index} produced length {}, expected {len}",
                    value.len()
                )));
            }

            values[index] = Some(value);
        }

        let mut out = HashMap::new();
        for (name, index) in expr.output_names.iter().zip(&expr.output_indices) {
            let value = values.get(*index).and_then(Clone::clone).ok_or_else(|| {
                BtError::Expression(format!("missing output node {index} for '{name}'"))
            })?;
            out.insert(name.clone(), value);
        }

        Ok(out)
    }
}

impl DefaultExprEvaluator {
    fn evaluate_node(
        &self,
        op: &CompiledOp,
        inputs: &[ArrayRef],
        env: &HashMap<String, ArrayRef>,
        params: &HashMap<String, ScalarValue>,
        len: usize,
    ) -> Result<ArrayRef> {
        // Check indicator cache before computing
        if let Some(ref cache) = self.cache {
            if IndicatorCache::is_cacheable(op) {
                let key = IndicatorCache::make_key(op, inputs);
                if let Some(cached) = cache.get(&key) {
                    return Ok(cached);
                }
                // Compute, cache, and return
                let result = self.evaluate_node_inner(op, inputs, env, params, len)?;
                cache.insert(key, result.clone());
                return Ok(result);
            }
        }
        self.evaluate_node_inner(op, inputs, env, params, len)
    }

    fn evaluate_node_inner(
        &self,
        op: &CompiledOp,
        inputs: &[ArrayRef],
        env: &HashMap<String, ArrayRef>,
        params: &HashMap<String, ScalarValue>,
        len: usize,
    ) -> Result<ArrayRef> {
        match op {
            CompiledOp::Column(name) => env
                .get(name)
                .cloned()
                .ok_or_else(|| BtError::Expression(format!("unknown input column '{name}'"))),
            CompiledOp::Literal(value) => scalar_to_array(value, len),
            CompiledOp::Parameter(name) => {
                let value = params.get(name).ok_or_else(|| {
                    BtError::Expression(format!("missing expression parameter '{name}'"))
                })?;
                scalar_to_array(value, len)
            }

            CompiledOp::Add => eval_numeric_binary(inputs, |lhs, rhs| lhs + rhs),
            CompiledOp::Sub => eval_numeric_binary(inputs, |lhs, rhs| lhs - rhs),
            CompiledOp::Mul => eval_numeric_binary(inputs, |lhs, rhs| lhs * rhs),
            CompiledOp::Div => {
                eval_numeric_binary(
                    inputs,
                    |lhs, rhs| {
                        if rhs == 0.0 {
                            f64::NAN
                        } else {
                            lhs / rhs
                        }
                    },
                )
            }

            CompiledOp::Gt => eval_comparison_binary(inputs, |lhs, rhs| lhs > rhs),
            CompiledOp::Lt => eval_comparison_binary(inputs, |lhs, rhs| lhs < rhs),
            CompiledOp::Eq => eval_equality_binary(inputs),
            CompiledOp::And => eval_bool_binary(inputs, |lhs, rhs| lhs && rhs),
            CompiledOp::Or => eval_bool_binary(inputs, |lhs, rhs| lhs || rhs),
            CompiledOp::Not => {
                let input = inputs
                    .first()
                    .ok_or_else(|| BtError::Expression("NOT expects one argument".to_string()))?;
                let input = coerce_bool(input)?;
                if input.null_count() == 0 {
                    let buf = !input.values();
                    return Ok(Arc::new(BooleanArray::new(buf, None)));
                }
                let mut out = Vec::with_capacity(input.len());
                for idx in 0..input.len() {
                    if input.is_null(idx) {
                        out.push(None);
                    } else {
                        out.push(Some(!input.value(idx)));
                    }
                }
                Ok(Arc::new(BooleanArray::from(out)))
            }

            CompiledOp::Lag { periods } => shift_array(inputs, *periods, true),
            CompiledOp::Lead { periods } => shift_array(inputs, *periods, false),
            CompiledOp::RollingMean { window } => rolling(inputs, *window, RollingKind::Mean),
            CompiledOp::RollingStd { window } => rolling(inputs, *window, RollingKind::Std),
            CompiledOp::RollingSum { window } => rolling(inputs, *window, RollingKind::Sum),
            CompiledOp::RollingMin { window } => rolling(inputs, *window, RollingKind::Min),
            CompiledOp::RollingMax { window } => rolling(inputs, *window, RollingKind::Max),
            CompiledOp::EwmMean { span } => ewm_mean(inputs, *span),
            CompiledOp::Diff { periods } => diff(inputs, *periods),
            CompiledOp::PctChange { periods } => pct_change(inputs, *periods),
            CompiledOp::CumSum => cum_sum(inputs),
            CompiledOp::CumProd => cum_prod(inputs),
            CompiledOp::Rank => rank(inputs),
            CompiledOp::ZScore { window } => zscore(inputs, *window),
            CompiledOp::Rsi { period } => rsi(inputs, *period),
            CompiledOp::Atr { period } => atr(inputs, *period),
            CompiledOp::LinRegSlope { window } => linreg(inputs, *window, LinRegOutput::Slope),
            CompiledOp::LinRegValue { window } => linreg(inputs, *window, LinRegOutput::Value),
            CompiledOp::LinRegR2 { window } => linreg(inputs, *window, LinRegOutput::R2),

            // -- New indicators (delegated to indicators module) --
            CompiledOp::Dema { period } => crate::indicators::dema(inputs, *period),
            CompiledOp::Tema { period } => crate::indicators::tema(inputs, *period),
            CompiledOp::Wma { window } => crate::indicators::wma(inputs, *window),
            CompiledOp::Hma { window } => crate::indicators::hma(inputs, *window),
            CompiledOp::Kama { period } => crate::indicators::kama(inputs, *period),
            CompiledOp::Roc { period } => crate::indicators::roc(inputs, *period),
            CompiledOp::Macd { fast, slow } => crate::indicators::macd(inputs, *fast, *slow),
            CompiledOp::MacdSignal { fast, slow, signal } => {
                crate::indicators::macd_signal(inputs, *fast, *slow, *signal)
            }
            CompiledOp::MacdHist { fast, slow, signal } => {
                crate::indicators::macd_hist(inputs, *fast, *slow, *signal)
            }
            CompiledOp::StochK { period } => crate::indicators::stoch_k(inputs, *period),
            CompiledOp::WilliamsR { period } => crate::indicators::williams_r(inputs, *period),
            CompiledOp::Cci { period } => crate::indicators::cci(inputs, *period),
            CompiledOp::Adx { period } => crate::indicators::adx(inputs, *period),
            CompiledOp::TrueRange => crate::indicators::true_range(inputs),
            CompiledOp::Natr { period } => crate::indicators::natr(inputs, *period),
            CompiledOp::BollingerUpper { window, num_std } => {
                crate::indicators::bollinger_upper(inputs, *window, *num_std)
            }
            CompiledOp::BollingerLower { window, num_std } => {
                crate::indicators::bollinger_lower(inputs, *window, *num_std)
            }
            CompiledOp::BollingerWidth { window, num_std } => {
                crate::indicators::bollinger_width(inputs, *window, *num_std)
            }
            CompiledOp::KeltnerUpper { period, multiplier } => {
                crate::indicators::keltner_upper(inputs, *period, *multiplier)
            }
            CompiledOp::KeltnerLower { period, multiplier } => {
                crate::indicators::keltner_lower(inputs, *period, *multiplier)
            }
            CompiledOp::SuperTrend { period, multiplier } => {
                crate::indicators::supertrend(inputs, *period, *multiplier)
            }
            CompiledOp::Obv => crate::indicators::obv(inputs),
            CompiledOp::Vwap => crate::indicators::vwap(inputs),
            CompiledOp::AdLine => crate::indicators::ad_line(inputs),
            CompiledOp::Mfi { period } => crate::indicators::mfi(inputs, *period),
            CompiledOp::CrossAbove => crate::indicators::cross_above(inputs),
            CompiledOp::CrossBelow => crate::indicators::cross_below(inputs),
            CompiledOp::RollingMedian { window } => {
                crate::indicators::rolling_median(inputs, *window)
            }
            CompiledOp::ParabolicSar { af_step, af_max } => {
                crate::indicators::parabolic_sar(inputs, *af_step, *af_max)
            }

            CompiledOp::IfElse => if_else(inputs),

            CompiledOp::Hour => extract_datetime_component(inputs, DatetimeComponent::Hour),
            CompiledOp::Minute => extract_datetime_component(inputs, DatetimeComponent::Minute),
            CompiledOp::DayOfWeek => {
                extract_datetime_component(inputs, DatetimeComponent::DayOfWeek)
            }
            CompiledOp::Month => extract_datetime_component(inputs, DatetimeComponent::Month),
            CompiledOp::DayOfMonth => {
                extract_datetime_component(inputs, DatetimeComponent::DayOfMonth)
            }

            CompiledOp::CrossSectionalMean => Err(BtError::Expression(
                "CrossSectionalMean requires orchestrator-level multi-symbol handling; \
                 this op should not be evaluated per-symbol"
                    .to_string(),
            )),
            CompiledOp::CrossSectionalRank => Err(BtError::Expression(
                "CrossSectionalRank requires orchestrator-level multi-symbol handling; \
                 this op should not be evaluated per-symbol"
                    .to_string(),
            )),
            CompiledOp::SymbolRef { .. } => Err(BtError::Expression(
                "SymbolRef requires orchestrator-level multi-symbol handling; \
                 this op should not be evaluated per-symbol"
                    .to_string(),
            )),

            CompiledOp::Function { name } => self.functions.call(name, inputs),

            CompiledOp::Scan(ref scan) => evaluate_scan(scan, inputs, params, len),
        }
    }
}

// ---------------------------------------------------------------------------
// Scan evaluator VM
// ---------------------------------------------------------------------------

fn evaluate_scan(
    scan: &CompiledScan,
    inputs: &[ArrayRef],
    params: &HashMap<String, ScalarValue>,
    len: usize,
) -> Result<ArrayRef> {
    // inputs layout: [init_0, init_1, ..., init_N, data_0, data_1, ..., data_M]
    let num_states = scan.num_states;

    // 1. Extract initial state values from the first element of each init array
    let mut prev_state = Vec::with_capacity(num_states);
    for i in 0..num_states {
        let init_arr = coerce_f64(&inputs[i])?;
        let val = if init_arr.is_null(0) {
            0.0
        } else {
            init_arr.value(0)
        };
        prev_state.push(val);
    }

    // 2. Extract data slices (pre-evaluated arrays)
    // First collect all coerced arrays so they live long enough for slicing.
    let mut data_arrays: Vec<Vec<f64>> = Vec::with_capacity(scan.data_column_names.len());
    for i in 0..scan.data_column_names.len() {
        let arr = coerce_f64(&inputs[num_states + i])?;
        if arr.null_count() == 0 {
            data_arrays.push(arr.values().as_ref().to_vec());
        } else {
            let mut vals = Vec::with_capacity(len);
            for j in 0..len {
                if arr.is_null(j) {
                    vals.push(0.0);
                } else {
                    vals.push(arr.value(j));
                }
            }
            data_arrays.push(vals);
        }
    }

    // 3. Resolve parameter values
    let mut param_vals = Vec::with_capacity(scan.param_names.len());
    for pname in &scan.param_names {
        let sv = params.get(pname).ok_or_else(|| {
            BtError::Expression(format!("scan references missing parameter '{pname}'"))
        })?;
        let val = sv.as_f64().ok_or_else(|| {
            BtError::Expression(format!(
                "scan parameter '{pname}' must be numeric, got {:?}",
                sv
            ))
        })?;
        param_vals.push(val);
    }

    // 4. Allocate registers and output
    let num_regs = scan.num_registers;
    let mut registers = vec![0.0f64; num_regs];
    let mut output = Vec::with_capacity(len);

    // 5. Main VM loop
    for t in 0..len {
        for (ip, op) in scan.instructions.iter().enumerate() {
            registers[ip] = match op {
                ScanOp::Literal(v) => *v,
                ScanOp::LoadData(i) => data_arrays[*i][t],
                ScanOp::LoadParam(i) => param_vals[*i],
                ScanOp::LoadPrev(i) => prev_state[*i],
                ScanOp::LoadVar(r) => registers[*r],
                ScanOp::Add(a, b) => registers[*a] + registers[*b],
                ScanOp::Sub(a, b) => registers[*a] - registers[*b],
                ScanOp::Mul(a, b) => registers[*a] * registers[*b],
                ScanOp::Div(a, b) => {
                    let denom = registers[*b];
                    if denom == 0.0 {
                        f64::NAN
                    } else {
                        registers[*a] / denom
                    }
                }
                ScanOp::Neg(a) => -registers[*a],
                ScanOp::Abs(a) => registers[*a].abs(),
                ScanOp::Sqrt(a) => registers[*a].sqrt(),
                ScanOp::Log(a) => registers[*a].ln(),
                ScanOp::Exp(a) => registers[*a].exp(),
                ScanOp::Max(a, b) => registers[*a].max(registers[*b]),
                ScanOp::Min(a, b) => registers[*a].min(registers[*b]),
                ScanOp::Gt(a, b) => {
                    if registers[*a] > registers[*b] {
                        1.0
                    } else {
                        0.0
                    }
                }
                ScanOp::Lt(a, b) => {
                    if registers[*a] < registers[*b] {
                        1.0
                    } else {
                        0.0
                    }
                }
                ScanOp::Eq(a, b) => {
                    if (registers[*a] - registers[*b]).abs() < f64::EPSILON {
                        1.0
                    } else {
                        0.0
                    }
                }
                ScanOp::And(a, b) => {
                    if registers[*a] != 0.0 && registers[*b] != 0.0 {
                        1.0
                    } else {
                        0.0
                    }
                }
                ScanOp::Or(a, b) => {
                    if registers[*a] != 0.0 || registers[*b] != 0.0 {
                        1.0
                    } else {
                        0.0
                    }
                }
                ScanOp::Not(a) => {
                    if registers[*a] == 0.0 {
                        1.0
                    } else {
                        0.0
                    }
                }
                ScanOp::IfElse {
                    cond,
                    then_val,
                    else_val,
                } => {
                    if registers[*cond] != 0.0 {
                        registers[*then_val]
                    } else {
                        registers[*else_val]
                    }
                }
            };
        }

        output.push(registers[scan.output_register]);

        // Writeback state for next iteration
        for &(step_idx, state_idx) in &scan.state_writeback {
            prev_state[state_idx] = registers[scan.step_result_registers[step_idx]];
        }
    }

    Ok(Arc::new(Float64Array::from(output)))
}

fn validate_and_resolve_len(inputs: &HashMap<String, ArrayRef>) -> Result<usize> {
    if inputs.is_empty() {
        return Ok(1);
    }

    let len = inputs
        .values()
        .next()
        .map(|array| array.len())
        .ok_or_else(|| {
            BtError::Expression("failed to resolve expression input length".to_string())
        })?;

    for (name, array) in inputs {
        if array.len() != len {
            return Err(BtError::Expression(format!(
                "input column '{name}' has length {}, expected {len}",
                array.len()
            )));
        }
    }

    Ok(len)
}

fn scalar_to_array(value: &ScalarValue, len: usize) -> Result<ArrayRef> {
    let out: ArrayRef = match value {
        ScalarValue::Null => Arc::new(Float64Array::from(vec![None; len])),
        ScalarValue::NaN => Arc::new(Float64Array::from(vec![f64::NAN; len])),
        ScalarValue::Bool(inner) => Arc::new(BooleanArray::from(vec![*inner; len])),
        ScalarValue::Int64(inner) => Arc::new(Int64Array::from(vec![*inner; len])),
        ScalarValue::Float64(inner) => Arc::new(Float64Array::from(vec![*inner; len])),
        ScalarValue::Utf8(inner) => Arc::new(StringArray::from(vec![inner.as_str(); len])),
    };
    Ok(out)
}

fn eval_numeric_binary(inputs: &[ArrayRef], op: impl Fn(f64, f64) -> f64) -> Result<ArrayRef> {
    let (lhs, rhs) = expect_binary(inputs)?;
    let lhs = coerce_f64(lhs)?;
    let rhs = coerce_f64(rhs)?;
    let len = lhs.len();
    if len != rhs.len() {
        return Err(BtError::Expression(format!(
            "numeric binary op length mismatch: {} vs {}",
            len,
            rhs.len()
        )));
    }

    // Fast path: no nulls — raw slice access, no Option wrapping.
    if lhs.null_count() == 0 && rhs.null_count() == 0 {
        let lv = lhs.values();
        let rv = rhs.values();
        let mut out = vec![0.0_f64; len];
        for i in 0..len {
            out[i] = op(lv[i], rv[i]);
        }
        return Ok(Arc::new(Float64Array::from(out)));
    }

    let mut out = Vec::with_capacity(len);
    for idx in 0..len {
        if lhs.is_null(idx) || rhs.is_null(idx) {
            out.push(None);
        } else {
            out.push(Some(op(lhs.value(idx), rhs.value(idx))));
        }
    }
    Ok(Arc::new(Float64Array::from(out)))
}

fn eval_comparison_binary(inputs: &[ArrayRef], op: impl Fn(f64, f64) -> bool) -> Result<ArrayRef> {
    let (lhs, rhs) = expect_binary(inputs)?;
    let lhs = coerce_f64(lhs)?;
    let rhs = coerce_f64(rhs)?;
    let len = lhs.len();
    if len != rhs.len() {
        return Err(BtError::Expression(format!(
            "comparison length mismatch: {} vs {}",
            len,
            rhs.len()
        )));
    }

    // Fast path: no nulls — raw slice access, BooleanBuffer directly.
    if lhs.null_count() == 0 && rhs.null_count() == 0 {
        let lv = lhs.values();
        let rv = rhs.values();
        let out: Vec<bool> = (0..len).map(|i| op(lv[i], rv[i])).collect();
        return Ok(Arc::new(BooleanArray::new(
            arrow::buffer::BooleanBuffer::from(out),
            None,
        )));
    }

    let mut out = Vec::with_capacity(len);
    for idx in 0..len {
        if lhs.is_null(idx) || rhs.is_null(idx) {
            out.push(None);
        } else {
            out.push(Some(op(lhs.value(idx), rhs.value(idx))));
        }
    }
    Ok(Arc::new(BooleanArray::from(out)))
}

fn eval_equality_binary(inputs: &[ArrayRef]) -> Result<ArrayRef> {
    let (lhs, rhs) = expect_binary(inputs)?;

    if lhs.data_type() == rhs.data_type() {
        if let (Some(lhs), Some(rhs)) = (
            lhs.as_any().downcast_ref::<Float64Array>(),
            rhs.as_any().downcast_ref::<Float64Array>(),
        ) {
            let mut out = Vec::with_capacity(lhs.len());
            for idx in 0..lhs.len() {
                if lhs.is_null(idx) || rhs.is_null(idx) {
                    out.push(None);
                } else {
                    out.push(Some(lhs.value(idx) == rhs.value(idx)));
                }
            }
            return Ok(Arc::new(BooleanArray::from(out)));
        }

        if let (Some(lhs), Some(rhs)) = (
            lhs.as_any().downcast_ref::<BooleanArray>(),
            rhs.as_any().downcast_ref::<BooleanArray>(),
        ) {
            let mut out = Vec::with_capacity(lhs.len());
            for idx in 0..lhs.len() {
                if lhs.is_null(idx) || rhs.is_null(idx) {
                    out.push(None);
                } else {
                    out.push(Some(lhs.value(idx) == rhs.value(idx)));
                }
            }
            return Ok(Arc::new(BooleanArray::from(out)));
        }

        if let (Some(lhs), Some(rhs)) = (
            lhs.as_any().downcast_ref::<StringArray>(),
            rhs.as_any().downcast_ref::<StringArray>(),
        ) {
            let mut out = Vec::with_capacity(lhs.len());
            for idx in 0..lhs.len() {
                if lhs.is_null(idx) || rhs.is_null(idx) {
                    out.push(None);
                } else {
                    out.push(Some(lhs.value(idx) == rhs.value(idx)));
                }
            }
            return Ok(Arc::new(BooleanArray::from(out)));
        }
    }

    eval_comparison_binary(inputs, |lhs, rhs| lhs == rhs)
}

fn eval_bool_binary(inputs: &[ArrayRef], op: impl Fn(bool, bool) -> bool) -> Result<ArrayRef> {
    let (lhs, rhs) = expect_binary(inputs)?;
    let lhs = coerce_bool(lhs)?;
    let rhs = coerce_bool(rhs)?;
    let len = lhs.len();
    if len != rhs.len() {
        return Err(BtError::Expression(format!(
            "bool binary op length mismatch: {} vs {}",
            len,
            rhs.len()
        )));
    }

    // Fast path: no nulls — use Arrow's bitwise buffer operations.
    if lhs.null_count() == 0 && rhs.null_count() == 0 {
        use arrow::buffer::BooleanBuffer;
        let lv = lhs.values();
        let rv = rhs.values();

        // Detect AND/OR by testing op(true,true), op(true,false), op(false,true)
        let tt = op(true, true);
        let tf = op(true, false);
        let ft = op(false, true);

        let buf = if tt && !tf && !ft {
            // AND
            lv & rv
        } else if tt && tf && ft {
            // OR
            lv | rv
        } else {
            // Generic fallback
            let out: Vec<bool> = (0..len).map(|i| op(lv.value(i), rv.value(i))).collect();
            BooleanBuffer::from(out)
        };
        return Ok(Arc::new(BooleanArray::new(buf, None)));
    }

    let mut out = Vec::with_capacity(len);
    for idx in 0..len {
        if lhs.is_null(idx) || rhs.is_null(idx) {
            out.push(None);
        } else {
            out.push(Some(op(lhs.value(idx), rhs.value(idx))));
        }
    }

    Ok(Arc::new(BooleanArray::from(out)))
}

fn shift_array(inputs: &[ArrayRef], periods: usize, lag: bool) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("shift op expects one argument".to_string()))?;
    if periods == 0 {
        return Ok(input.clone());
    }

    if let Some(input) = input.as_any().downcast_ref::<Float64Array>() {
        let mut out = Vec::with_capacity(input.len());
        for idx in 0..input.len() {
            let source = if lag {
                idx.checked_sub(periods)
            } else {
                idx.checked_add(periods)
                    .filter(|value| *value < input.len())
            };

            match source {
                Some(source) if !input.is_null(source) => out.push(Some(input.value(source))),
                _ => out.push(None),
            }
        }
        return Ok(Arc::new(Float64Array::from(out)));
    }

    if let Some(input) = input.as_any().downcast_ref::<BooleanArray>() {
        let mut out = Vec::with_capacity(input.len());
        for idx in 0..input.len() {
            let source = if lag {
                idx.checked_sub(periods)
            } else {
                idx.checked_add(periods)
                    .filter(|value| *value < input.len())
            };

            match source {
                Some(source) if !input.is_null(source) => out.push(Some(input.value(source))),
                _ => out.push(None),
            }
        }
        return Ok(Arc::new(BooleanArray::from(out)));
    }

    Err(BtError::Expression(format!(
        "shift op unsupported for {:?}",
        input.data_type()
    )))
}

#[derive(Copy, Clone)]
enum RollingKind {
    Mean,
    Std,
    Sum,
    Min,
    Max,
}

/// Read a f64 value, treating null and NaN as "bad" (excluded from aggregation).
#[inline(always)]
fn read_f64_val(input: &Float64Array, idx: usize) -> (f64, bool) {
    if input.is_null(idx) {
        (0.0, true)
    } else {
        let v = input.value(idx);
        (v, v.is_nan())
    }
}

/// O(n) rolling aggregation. Previous implementation was O(n*window).
fn rolling(inputs: &[ArrayRef], window: usize, kind: RollingKind) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("rolling op expects one argument".to_string()))?;
    if window == 0 {
        return Err(BtError::Expression(
            "rolling window must be > 0".to_string(),
        ));
    }

    let input = coerce_f64(input)?;
    let len = input.len();
    let w = window as f64;
    let mut out = Vec::with_capacity(len);

    match kind {
        RollingKind::Sum | RollingKind::Mean => {
            let divide = matches!(kind, RollingKind::Mean);
            let mut sum = 0.0;
            let mut bad_count = 0usize;
            for idx in 0..len {
                let (val, is_bad) = read_f64_val(&input, idx);
                if is_bad {
                    bad_count += 1;
                } else {
                    sum += val;
                }
                if idx >= window {
                    let (old, was_bad) = read_f64_val(&input, idx - window);
                    if was_bad {
                        bad_count -= 1;
                    } else {
                        sum -= old;
                    }
                }
                if idx + 1 < window || bad_count > 0 {
                    out.push(Some(f64::NAN));
                } else if divide {
                    out.push(Some(sum / w));
                } else {
                    out.push(Some(sum));
                }
            }
        }
        RollingKind::Std => {
            let mut sum = 0.0;
            let mut sum_sq = 0.0;
            let mut bad_count = 0usize;
            for idx in 0..len {
                let (val, is_bad) = read_f64_val(&input, idx);
                if is_bad {
                    bad_count += 1;
                } else {
                    sum += val;
                    sum_sq += val * val;
                }
                if idx >= window {
                    let (old, was_bad) = read_f64_val(&input, idx - window);
                    if was_bad {
                        bad_count -= 1;
                    } else {
                        sum -= old;
                        sum_sq -= old * old;
                    }
                }
                if idx + 1 < window || bad_count > 0 {
                    out.push(Some(f64::NAN));
                } else {
                    let mean = sum / w;
                    let variance = sum_sq / w - mean * mean;
                    out.push(Some(variance.max(0.0).sqrt()));
                }
            }
        }
        RollingKind::Min => {
            let mut deque: VecDeque<usize> = VecDeque::new();
            for idx in 0..len {
                // Evict front if outside window.
                while let Some(&front) = deque.front() {
                    if front + window <= idx {
                        deque.pop_front();
                    } else {
                        break;
                    }
                }
                let (val, is_bad) = read_f64_val(&input, idx);
                if !is_bad {
                    // Maintain monotone increasing deque for min.
                    while let Some(&back) = deque.back() {
                        if input.value(back) >= val {
                            deque.pop_back();
                        } else {
                            break;
                        }
                    }
                    deque.push_back(idx);
                }
                if idx + 1 < window || deque.is_empty() {
                    out.push(Some(f64::NAN));
                } else {
                    out.push(Some(input.value(*deque.front().unwrap())));
                }
            }
        }
        RollingKind::Max => {
            let mut deque: VecDeque<usize> = VecDeque::new();
            for idx in 0..len {
                while let Some(&front) = deque.front() {
                    if front + window <= idx {
                        deque.pop_front();
                    } else {
                        break;
                    }
                }
                let (val, is_bad) = read_f64_val(&input, idx);
                if !is_bad {
                    // Maintain monotone decreasing deque for max.
                    while let Some(&back) = deque.back() {
                        if input.value(back) <= val {
                            deque.pop_back();
                        } else {
                            break;
                        }
                    }
                    deque.push_back(idx);
                }
                if idx + 1 < window || deque.is_empty() {
                    out.push(Some(f64::NAN));
                } else {
                    out.push(Some(input.value(*deque.front().unwrap())));
                }
            }
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

fn ewm_mean(inputs: &[ArrayRef], span: f64) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("ewm_mean expects one argument".to_string()))?;
    if span <= 0.0 {
        return Err(BtError::Expression("ewm span must be > 0".to_string()));
    }

    let input = coerce_f64(input)?;
    let alpha = 2.0 / (span + 1.0);
    let one_minus_alpha = 1.0 - alpha;

    // Fast path: no nulls — direct slice access, no Option wrapping.
    // Note: inputs may still contain NaN floats (e.g. Roc warmup bars) even
    // when null_count() == 0, so we must skip NaN values the same way the
    // slow path does instead of letting them poison `last`.
    if input.null_count() == 0 {
        let vals = input.values();
        let len = vals.len();
        let mut out = vec![f64::NAN; len];
        let mut last = f64::NAN;
        for i in 0..len {
            let v = vals[i];
            if v.is_nan() {
                continue;
            }
            if last.is_nan() {
                last = v;
            } else {
                last = alpha * v + one_minus_alpha * last;
            }
            out[i] = last;
        }
        return Ok(Arc::new(Float64Array::from(out)));
    }

    let mut out = Vec::with_capacity(input.len());
    let mut last = f64::NAN;
    for idx in 0..input.len() {
        if input.is_null(idx) {
            out.push(Some(f64::NAN));
            continue;
        }

        let value = input.value(idx);
        if last.is_nan() {
            last = value;
        } else {
            last = alpha * value + one_minus_alpha * last;
        }
        out.push(Some(last));
    }

    Ok(Arc::new(Float64Array::from(out)))
}

fn diff(inputs: &[ArrayRef], periods: usize) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("diff expects one argument".to_string()))?;
    let input = coerce_f64(input)?;

    let mut out = Vec::with_capacity(input.len());
    for idx in 0..input.len() {
        if idx < periods || input.is_null(idx) || input.is_null(idx - periods) {
            out.push(Some(f64::NAN));
        } else {
            out.push(Some(input.value(idx) - input.value(idx - periods)));
        }
    }
    Ok(Arc::new(Float64Array::from(out)))
}

fn pct_change(inputs: &[ArrayRef], periods: usize) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("pct_change expects one argument".to_string()))?;
    let input = coerce_f64(input)?;

    let mut out = Vec::with_capacity(input.len());
    for idx in 0..input.len() {
        if idx < periods || input.is_null(idx) || input.is_null(idx - periods) {
            out.push(Some(f64::NAN));
            continue;
        }

        let prev = input.value(idx - periods);
        if prev == 0.0 {
            out.push(Some(f64::NAN));
        } else {
            out.push(Some((input.value(idx) - prev) / prev));
        }
    }
    Ok(Arc::new(Float64Array::from(out)))
}

fn cum_sum(inputs: &[ArrayRef]) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("cumsum expects one argument".to_string()))?;
    let input = coerce_f64(input)?;

    let mut out = Vec::with_capacity(input.len());
    let mut running = 0.0_f64;
    for idx in 0..input.len() {
        if input.is_null(idx) {
            out.push(Some(f64::NAN));
            continue;
        }
        running += input.value(idx);
        out.push(Some(running));
    }
    Ok(Arc::new(Float64Array::from(out)))
}

fn cum_prod(inputs: &[ArrayRef]) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("cumprod expects one argument".to_string()))?;
    let input = coerce_f64(input)?;

    let mut out = Vec::with_capacity(input.len());
    let mut running = 1.0_f64;
    for idx in 0..input.len() {
        if input.is_null(idx) {
            out.push(Some(f64::NAN));
            continue;
        }
        running *= input.value(idx);
        out.push(Some(running));
    }
    Ok(Arc::new(Float64Array::from(out)))
}

fn rank(inputs: &[ArrayRef]) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("rank expects one argument".to_string()))?;
    let input = coerce_f64(input)?;

    let mut ranked = vec![f64::NAN; input.len()];
    let mut values = Vec::new();
    for idx in 0..input.len() {
        if !input.is_null(idx) {
            values.push((idx, input.value(idx)));
        }
    }

    values.sort_by(|(_, lhs): &(usize, f64), (_, rhs): &(usize, f64)| lhs.total_cmp(rhs));
    for (rank, (index, _)) in values.into_iter().enumerate() {
        ranked[index] = rank as f64 + 1.0;
    }

    Ok(Arc::new(Float64Array::from(
        ranked.into_iter().map(Some).collect::<Vec<_>>(),
    )))
}

/// Single-pass O(n) zscore: computes rolling mean and std together.
fn zscore(inputs: &[ArrayRef], window: usize) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("zscore expects one argument".to_string()))?;
    if window == 0 {
        return Err(BtError::Expression("zscore window must be > 0".to_string()));
    }

    let values = coerce_f64(input)?;

    // Fast path: no null bitmap — access raw &[f64] directly.
    if values.null_count() == 0 {
        return zscore_no_nulls(values.values(), window);
    }

    let len = values.len();
    let w = window as f64;
    let mut out = Vec::with_capacity(len);

    let mut sum = 0.0;
    let mut sum_sq = 0.0;
    let mut bad_count = 0usize;

    for idx in 0..len {
        let (val, is_bad) = read_f64_val(&values, idx);
        if is_bad {
            bad_count += 1;
        } else {
            sum += val;
            sum_sq += val * val;
        }
        if idx >= window {
            let (old, was_bad) = read_f64_val(&values, idx - window);
            if was_bad {
                bad_count -= 1;
            } else {
                sum -= old;
                sum_sq -= old * old;
            }
        }
        if idx + 1 < window || bad_count > 0 || is_bad {
            out.push(Some(f64::NAN));
        } else {
            let mean = sum / w;
            let variance = sum_sq / w - mean * mean;
            let std = variance.max(0.0).sqrt();
            if std == 0.0 {
                out.push(Some(f64::NAN));
            } else {
                out.push(Some((val - mean) / std));
            }
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

/// Fast-path zscore: direct slice access, no null bitmap checks, no Option wrapping.
fn zscore_no_nulls(vals: &[f64], window: usize) -> Result<ArrayRef> {
    let len = vals.len();
    let w = window as f64;
    let mut out = vec![f64::NAN; len];

    let mut sum = 0.0;
    let mut sum_sq = 0.0;
    let mut nan_count = 0usize;

    for idx in 0..len {
        let val = vals[idx];
        if val.is_nan() {
            nan_count += 1;
        } else {
            sum += val;
            sum_sq += val * val;
        }
        if idx >= window {
            let old = vals[idx - window];
            if old.is_nan() {
                nan_count -= 1;
            } else {
                sum -= old;
                sum_sq -= old * old;
            }
        }
        if idx + 1 >= window && nan_count == 0 && !val.is_nan() {
            let mean = sum / w;
            let variance = sum_sq / w - mean * mean;
            let std = variance.max(0.0).sqrt();
            if std != 0.0 {
                out[idx] = (val - mean) / std;
            }
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

/// Single-pass O(n) RSI using Wilder's smoothing.
/// Input: one numeric array (typically close prices).
/// Output: Float64 in [0, 100], with first `period` values as NaN.
fn rsi(inputs: &[ArrayRef], period: usize) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("rsi expects one argument".to_string()))?;
    if period == 0 {
        return Err(BtError::Expression("rsi period must be > 0".to_string()));
    }

    let values = coerce_f64(input)?;

    // Fast path: no nulls
    if values.null_count() == 0 {
        return rsi_no_nulls(values.values(), period);
    }

    let len = values.len();
    let mut out = Vec::with_capacity(len);

    if len == 0 {
        return Ok(Arc::new(Float64Array::from(out)));
    }

    // First element has no delta → NaN
    out.push(Some(f64::NAN));

    let mut sum_gain = 0.0;
    let mut sum_loss = 0.0;
    let mut avg_gain: f64;
    let mut avg_loss: f64;
    let p = period as f64;

    // Warm-up: accumulate first `period` deltas (indices 1..=period)
    for i in 1..len.min(period + 1) {
        if values.is_null(i) || values.is_null(i - 1) {
            out.push(Some(f64::NAN));
            continue;
        }
        let delta = values.value(i) - values.value(i - 1);
        if delta > 0.0 {
            sum_gain += delta;
        } else {
            sum_loss += -delta;
        }
        out.push(Some(f64::NAN));
    }

    if len <= period {
        return Ok(Arc::new(Float64Array::from(out)));
    }

    // First real RSI at index = period
    avg_gain = sum_gain / p;
    avg_loss = sum_loss / p;
    let last = out.last_mut().unwrap();
    if avg_loss == 0.0 {
        *last = Some(100.0);
    } else {
        let rs = avg_gain / avg_loss;
        *last = Some(100.0 - 100.0 / (1.0 + rs));
    }

    // Wilder's smoothing for remaining values
    for i in (period + 1)..len {
        if values.is_null(i) || values.is_null(i - 1) {
            out.push(Some(f64::NAN));
            continue;
        }
        let delta = values.value(i) - values.value(i - 1);
        let (gain, loss) = if delta > 0.0 {
            (delta, 0.0)
        } else {
            (0.0, -delta)
        };
        avg_gain = (avg_gain * (p - 1.0) + gain) / p;
        avg_loss = (avg_loss * (p - 1.0) + loss) / p;
        if avg_loss == 0.0 {
            out.push(Some(100.0));
        } else {
            let rs = avg_gain / avg_loss;
            out.push(Some(100.0 - 100.0 / (1.0 + rs)));
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

/// Fast-path RSI: direct slice access, no null checks.
fn rsi_no_nulls(vals: &[f64], period: usize) -> Result<ArrayRef> {
    let len = vals.len();
    let mut out = vec![f64::NAN; len];
    let p = period as f64;

    if len <= period {
        return Ok(Arc::new(Float64Array::from(out)));
    }

    let mut sum_gain = 0.0;
    let mut sum_loss = 0.0;

    for i in 1..=period {
        let delta = vals[i] - vals[i - 1];
        if delta > 0.0 {
            sum_gain += delta;
        } else {
            sum_loss += -delta;
        }
    }

    let mut avg_gain = sum_gain / p;
    let mut avg_loss = sum_loss / p;

    if avg_loss == 0.0 {
        out[period] = 100.0;
    } else {
        let rs = avg_gain / avg_loss;
        out[period] = 100.0 - 100.0 / (1.0 + rs);
    }

    for i in (period + 1)..len {
        let delta = vals[i] - vals[i - 1];
        let (gain, loss) = if delta > 0.0 {
            (delta, 0.0)
        } else {
            (0.0, -delta)
        };
        avg_gain = (avg_gain * (p - 1.0) + gain) / p;
        avg_loss = (avg_loss * (p - 1.0) + loss) / p;
        if avg_loss == 0.0 {
            out[i] = 100.0;
        } else {
            let rs = avg_gain / avg_loss;
            out[i] = 100.0 - 100.0 / (1.0 + rs);
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

/// Single-pass O(n) ATR using Wilder's smoothing.
/// Inputs: [high, low, close] — three numeric arrays.
/// True Range = max(H-L, |H-prevClose|, |L-prevClose|)
/// ATR = Wilder's smoothing of TR over `period`.
fn atr(inputs: &[ArrayRef], period: usize) -> Result<ArrayRef> {
    if inputs.len() != 3 {
        return Err(BtError::Expression(format!(
            "atr expects 3 arguments (high, low, close), got {}",
            inputs.len()
        )));
    }
    if period == 0 {
        return Err(BtError::Expression("atr period must be > 0".to_string()));
    }

    let high = coerce_f64(&inputs[0])?;
    let low = coerce_f64(&inputs[1])?;
    let close = coerce_f64(&inputs[2])?;
    let len = high.len();

    if low.len() != len || close.len() != len {
        return Err(BtError::Expression(
            "atr: high, low, close arrays must have the same length".to_string(),
        ));
    }

    // Fast path: no nulls
    if high.null_count() == 0 && low.null_count() == 0 && close.null_count() == 0 {
        return atr_no_nulls(high.values(), low.values(), close.values(), period);
    }

    let mut out = Vec::with_capacity(len);

    if len == 0 {
        return Ok(Arc::new(Float64Array::from(out)));
    }

    // Index 0: no previous close, TR = H - L
    out.push(Some(f64::NAN));

    let mut tr_sum = 0.0;
    let mut atr_val: f64;
    let p = period as f64;

    // Warm-up: compute TR and accumulate for first `period` bars (indices 1..=period)
    for i in 1..len.min(period + 1) {
        if high.is_null(i) || low.is_null(i) || close.is_null(i) || close.is_null(i - 1) {
            out.push(Some(f64::NAN));
            continue;
        }
        let h = high.value(i);
        let l = low.value(i);
        let pc = close.value(i - 1);
        let tr = (h - l).max((h - pc).abs()).max((l - pc).abs());
        tr_sum += tr;
        out.push(Some(f64::NAN));
    }

    if len <= period {
        return Ok(Arc::new(Float64Array::from(out)));
    }

    // First ATR at index = period
    atr_val = tr_sum / p;
    let last = out.last_mut().unwrap();
    *last = Some(atr_val);

    // Wilder's smoothing for remaining values
    for i in (period + 1)..len {
        if high.is_null(i) || low.is_null(i) || close.is_null(i) || close.is_null(i - 1) {
            out.push(Some(f64::NAN));
            continue;
        }
        let h = high.value(i);
        let l = low.value(i);
        let pc = close.value(i - 1);
        let tr = (h - l).max((h - pc).abs()).max((l - pc).abs());
        atr_val = (atr_val * (p - 1.0) + tr) / p;
        out.push(Some(atr_val));
    }

    Ok(Arc::new(Float64Array::from(out)))
}

/// Fast-path ATR: direct slice access, no null checks.
fn atr_no_nulls(h: &[f64], l: &[f64], c: &[f64], period: usize) -> Result<ArrayRef> {
    let len = h.len();
    let mut out = vec![f64::NAN; len];
    let p = period as f64;

    if len <= period {
        return Ok(Arc::new(Float64Array::from(out)));
    }

    let mut tr_sum = 0.0;
    for i in 1..=period {
        let tr = (h[i] - l[i])
            .max((h[i] - c[i - 1]).abs())
            .max((l[i] - c[i - 1]).abs());
        tr_sum += tr;
    }

    let mut atr_val = tr_sum / p;
    out[period] = atr_val;

    for i in (period + 1)..len {
        let tr = (h[i] - l[i])
            .max((h[i] - c[i - 1]).abs())
            .max((l[i] - c[i - 1]).abs());
        atr_val = (atr_val * (p - 1.0) + tr) / p;
        out[i] = atr_val;
    }

    Ok(Arc::new(Float64Array::from(out)))
}

// ---------------------------------------------------------------------------
// Rolling linear regression: slope, predicted value, R²
// ---------------------------------------------------------------------------

#[derive(Copy, Clone)]
enum LinRegOutput {
    Slope,
    Value,
    R2,
}

/// Single-pass O(n) rolling linear regression.
///
/// Fits y = a + b*x with x = 0, 1, ..., w-1 within each window.
/// Maintains rolling sums: sy (Σy), sxy (Σ(local_x * y)), sy2 (Σy² for R²).
///
/// Update rule when sliding window (remove y_old, add y_new):
///   sxy = sxy - sy + y_old + (w-1) * y_new   (before updating sy)
///   sy  = sy - y_old + y_new
///   sy2 = sy2 - y_old² + y_new²
fn linreg(inputs: &[ArrayRef], window: usize, output: LinRegOutput) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("linreg expects one argument".to_string()))?;
    if window < 2 {
        return Err(BtError::Expression(
            "linreg window must be >= 2".to_string(),
        ));
    }

    let values = coerce_f64(input)?;

    if values.null_count() == 0 {
        return linreg_no_nulls(values.values(), window, output);
    }

    // Nullable path: same algorithm as no-nulls but with bad-value tracking.
    // For simplicity and correctness, delegate to no-nulls after extracting
    // values (treating null/NaN uniformly as NaN).
    let len = values.len();
    let mut vals = vec![0.0_f64; len];
    for i in 0..len {
        if values.is_null(i) {
            vals[i] = f64::NAN;
        } else {
            vals[i] = values.value(i);
        }
    }
    linreg_no_nulls(&vals, window, output)
}

/// Fast-path rolling linear regression: direct slice access, no null checks.
fn linreg_no_nulls(vals: &[f64], window: usize, output: LinRegOutput) -> Result<ArrayRef> {
    let len = vals.len();
    let w = window as f64;
    let sx = w * (w - 1.0) / 2.0;
    let sx2 = w * (w - 1.0) * (2.0 * w - 1.0) / 6.0;
    let denom = w * sx2 - sx * sx;
    let wm1 = (window - 1) as f64;

    let mut out = vec![f64::NAN; len];

    // Build initial window [0..window)
    if len < window {
        return Ok(Arc::new(Float64Array::from(out)));
    }

    let mut sy = 0.0;
    let mut sxy = 0.0;
    let mut sy2 = 0.0;
    let mut nan_count = 0usize;

    // Warm up: fill first window
    for i in 0..window {
        let v = vals[i];
        if v.is_nan() {
            nan_count += 1;
        } else {
            sy += v;
            sxy += i as f64 * v;
            sy2 += v * v;
        }
    }

    if nan_count == 0 {
        emit_linreg(
            &mut out,
            window - 1,
            sy,
            sxy,
            sy2,
            w,
            sx,
            denom,
            wm1,
            output,
        );
    }

    // Slide window
    for idx in window..len {
        let new_val = vals[idx];
        let old_val = vals[idx - window];

        // Remove old value (was at local position 0)
        if old_val.is_nan() {
            nan_count -= 1;
        } else {
            // Shift: all remaining local positions decrease by 1
            // sxy -= (sy - old_val) means: subtract 1*y for each y still in window
            sxy -= sy - old_val;
            sy -= old_val;
            sy2 -= old_val * old_val;
        }

        // Add new value at local position w-1
        if new_val.is_nan() {
            nan_count += 1;
        } else {
            sy += new_val;
            sxy += wm1 * new_val;
            sy2 += new_val * new_val;
        }

        if nan_count == 0 {
            emit_linreg(&mut out, idx, sy, sxy, sy2, w, sx, denom, wm1, output);
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

#[inline(always)]
fn emit_linreg(
    out: &mut [f64],
    idx: usize,
    sy: f64,
    sxy: f64,
    sy2: f64,
    w: f64,
    sx: f64,
    denom: f64,
    wm1: f64,
    output: LinRegOutput,
) {
    let slope = (w * sxy - sx * sy) / denom;
    match output {
        LinRegOutput::Slope => out[idx] = slope,
        LinRegOutput::Value => {
            let mean = sy / w;
            out[idx] = mean + slope * wm1 / 2.0;
        }
        LinRegOutput::R2 => {
            let ss_tot = w * sy2 - sy * sy;
            if ss_tot.abs() < 1e-30 {
                // constant series — R² undefined
            } else {
                let num = w * sxy - sx * sy;
                let r2 = (num * num) / (denom * ss_tot);
                out[idx] = r2.clamp(0.0, 1.0);
            }
        }
    }
}

fn if_else(inputs: &[ArrayRef]) -> Result<ArrayRef> {
    if inputs.len() != 3 {
        return Err(BtError::Expression(format!(
            "if_else expects 3 arguments, got {}",
            inputs.len()
        )));
    }

    let condition = coerce_bool(&inputs[0])?;
    if let (Ok(then_values), Ok(else_values)) = (coerce_f64(&inputs[1]), coerce_f64(&inputs[2]))
        as (Result<Cow<'_, Float64Array>>, Result<Cow<'_, Float64Array>>)
    {
        let len = condition.len();
        if then_values.len() != len || else_values.len() != len {
            return Err(BtError::Expression(
                "if_else argument length mismatch".to_string(),
            ));
        }

        // Fast path: no nulls in any input.
        if condition.null_count() == 0
            && then_values.null_count() == 0
            && else_values.null_count() == 0
        {
            let tv = then_values.values();
            let ev = else_values.values();
            let mut out = vec![0.0_f64; len];
            for i in 0..len {
                out[i] = if condition.value(i) { tv[i] } else { ev[i] };
            }
            return Ok(Arc::new(Float64Array::from(out)));
        }

        let mut out = Vec::with_capacity(len);
        for idx in 0..len {
            if condition.is_null(idx) {
                out.push(None);
            } else if condition.value(idx) {
                out.push(if then_values.is_null(idx) {
                    None
                } else {
                    Some(then_values.value(idx))
                });
            } else {
                out.push(if else_values.is_null(idx) {
                    None
                } else {
                    Some(else_values.value(idx))
                });
            }
        }

        return Ok(Arc::new(Float64Array::from(out)));
    }

    Err(BtError::Expression(
        "if_else currently supports numeric branches only".to_string(),
    ))
}

fn expect_binary(inputs: &[ArrayRef]) -> Result<(&ArrayRef, &ArrayRef)> {
    if inputs.len() != 2 {
        return Err(BtError::Expression(format!(
            "binary operation expected 2 arguments, got {}",
            inputs.len()
        )));
    }

    Ok((&inputs[0], &inputs[1]))
}

fn coerce_bool(array: &ArrayRef) -> Result<BooleanArray> {
    if let Some(array) = array.as_any().downcast_ref::<BooleanArray>() {
        return Ok(array.clone());
    }

    Err(BtError::Expression(format!(
        "expected boolean array, got {:?}",
        array.data_type()
    )))
}

// ---------------------------------------------------------------------------
// Cross-sectional multi-symbol operations
// ---------------------------------------------------------------------------

/// Compute cross-sectional mean at each row across all symbols.
///
/// For each row index, computes the mean of the values across all symbol arrays
/// (skipping NaN/null). Returns one array per symbol, each filled with the
/// same cross-sectional mean at each row.
pub fn cross_sectional_mean_multi(symbol_arrays: &[&Float64Array]) -> Vec<Float64Array> {
    if symbol_arrays.is_empty() {
        return Vec::new();
    }
    let len = symbol_arrays[0].len();
    let n_symbols = symbol_arrays.len();

    let mut means = Vec::with_capacity(len);
    for row in 0..len {
        let mut sum = 0.0_f64;
        let mut count = 0_u32;
        for arr in symbol_arrays {
            if !arr.is_null(row) {
                let v = arr.value(row);
                if v.is_finite() {
                    sum += v;
                    count += 1;
                }
            }
        }
        means.push(if count == 0 {
            f64::NAN
        } else {
            sum / count as f64
        });
    }

    let mean_array = Float64Array::from(means);
    vec![mean_array; n_symbols]
}

/// Compute cross-sectional fractional rank at each row across all symbols.
///
/// For each row index, ranks the values across all symbol arrays. Ties receive
/// the average rank. Output is normalized to \[0, 1\] where 0 = lowest, 1 = highest.
/// NaN inputs produce NaN output.
pub fn cross_sectional_rank_multi(symbol_arrays: &[&Float64Array]) -> Vec<Float64Array> {
    if symbol_arrays.is_empty() {
        return Vec::new();
    }
    let len = symbol_arrays[0].len();
    let n_symbols = symbol_arrays.len();

    let mut result: Vec<Vec<f64>> = vec![vec![f64::NAN; len]; n_symbols];

    for row in 0..len {
        // Collect (value, symbol_index) for non-null finite values
        let mut vals: Vec<(f64, usize)> = Vec::new();
        for (si, arr) in symbol_arrays.iter().enumerate() {
            if !arr.is_null(row) {
                let v = arr.value(row);
                if v.is_finite() {
                    vals.push((v, si));
                }
            }
        }

        if vals.is_empty() {
            continue; // all NaN for this row
        }

        // Sort by value
        vals.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

        let n = vals.len() as f64;
        // Assign fractional ranks handling ties
        let mut i = 0;
        while i < vals.len() {
            let mut j = i;
            while j < vals.len() && vals[j].0 == vals[i].0 {
                j += 1;
            }
            // Average rank for tied values, normalized to [0, 1]
            let avg_rank = (i + j - 1) as f64 / 2.0;
            let frac_rank = if n > 1.0 { avg_rank / (n - 1.0) } else { 0.5 };
            for k in i..j {
                result[vals[k].1][row] = frac_rank;
            }
            i = j;
        }
    }

    result.into_iter().map(Float64Array::from).collect()
}

// ---------------------------------------------------------------------------
// Datetime component extraction
// ---------------------------------------------------------------------------

#[derive(Copy, Clone)]
enum DatetimeComponent {
    Hour,
    Minute,
    DayOfWeek,
    Month,
    DayOfMonth,
}

/// Extract a datetime component from a timestamp nanosecond column.
///
/// Accepts either `TimestampNanosecond` or numeric arrays (interpreted as
/// nanoseconds since Unix epoch). Returns `Float64Array` with the extracted
/// component as a float for consistency with the rest of the expression engine.
///
/// Component ranges:
/// - Hour: 0–23
/// - Minute: 0–59
/// - DayOfWeek: 0=Monday, 6=Sunday (ISO weekday)
/// - Month: 1–12
/// - DayOfMonth: 1–31
fn extract_datetime_component(
    inputs: &[ArrayRef],
    component: DatetimeComponent,
) -> Result<ArrayRef> {
    let input = inputs
        .first()
        .ok_or_else(|| BtError::Expression("datetime op expects one argument".to_string()))?;

    // Try direct TimestampNanosecondArray first, then fall back to coercing
    // to i64 (covers cases where the timestamp arrived as Int64 or Float64).
    let nanos: Vec<Option<i64>> =
        if let Some(ts) = input.as_any().downcast_ref::<TimestampNanosecondArray>() {
            (0..ts.len())
                .map(|i| {
                    if ts.is_null(i) {
                        None
                    } else {
                        Some(ts.value(i))
                    }
                })
                .collect()
        } else if let Some(arr) = input.as_any().downcast_ref::<Int64Array>() {
            (0..arr.len())
                .map(|i| {
                    if arr.is_null(i) {
                        None
                    } else {
                        Some(arr.value(i))
                    }
                })
                .collect()
        } else if let Some(arr) = input.as_any().downcast_ref::<Float64Array>() {
            (0..arr.len())
                .map(|i| {
                    if arr.is_null(i) || arr.value(i).is_nan() {
                        None
                    } else {
                        Some(arr.value(i) as i64)
                    }
                })
                .collect()
        } else {
            return Err(BtError::Expression(format!(
                "datetime op requires a timestamp column, got {:?}",
                input.data_type()
            )));
        };

    let mut out = Vec::with_capacity(nanos.len());
    for ns_opt in nanos {
        match ns_opt {
            None => out.push(None),
            Some(ns) => {
                let secs = ns.div_euclid(1_000_000_000);
                let subsec_nanos = ns.rem_euclid(1_000_000_000) as u32;
                let dt = DateTime::from_timestamp(secs, subsec_nanos).ok_or_else(|| {
                    BtError::Expression(format!("invalid timestamp nanoseconds: {ns}"))
                })?;
                let value = match component {
                    DatetimeComponent::Hour => dt.hour() as f64,
                    DatetimeComponent::Minute => dt.minute() as f64,
                    DatetimeComponent::DayOfWeek => {
                        // chrono: Monday=0 .. Sunday=6 (num_days_from_monday)
                        dt.weekday().num_days_from_monday() as f64
                    }
                    DatetimeComponent::Month => dt.month() as f64,
                    DatetimeComponent::DayOfMonth => dt.day() as f64,
                };
                out.push(Some(value));
            }
        }
    }

    Ok(Arc::new(Float64Array::from(out)))
}

pub(crate) fn coerce_f64(array: &ArrayRef) -> Result<Cow<'_, Float64Array>> {
    if let Some(array) = array.as_any().downcast_ref::<Float64Array>() {
        return Ok(Cow::Borrowed(array));
    }

    let mut out = Vec::with_capacity(array.len());
    if let Some(array) = array.as_any().downcast_ref::<Float32Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<Int64Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<Int32Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<UInt64Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<UInt32Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<UInt8Array>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else if let Some(array) = array.as_any().downcast_ref::<TimestampNanosecondArray>() {
        for idx in 0..array.len() {
            out.push(if array.is_null(idx) {
                None
            } else {
                Some(array.value(idx) as f64)
            });
        }
    } else {
        return Err(BtError::Expression(format!(
            "expected numeric array, got {:?}",
            array.data_type()
        )));
    }

    Ok(Cow::Owned(Float64Array::from(out)))
}
