use std::collections::HashMap;
use std::sync::Arc;

use arrow::array::{Float64Array, TimestampNanosecondArray};
use bt_expr::{
    compile_expr, cross_sectional_mean_multi, cross_sectional_rank_multi,
    extract_cross_section_info, extract_symbol_ref_info, has_cross_sectional_ops,
    has_symbol_ref_ops, infer_expr_type, CrossSectionKind, DefaultExprEvaluator, DynPeriod, Expr,
    ExprEvaluator, ScalarValue, ValueType,
};

#[test]
fn compiler_constant_folding_and_cse_work() {
    let shared = Expr::Add(
        Box::new(Expr::Column("close".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(1.0))),
    );
    let expr = Expr::Add(Box::new(shared.clone()), Box::new(shared));
    let compiled = compile_expr(expr).expect("compile expression");

    // Nodes: close, 1.0, (close + 1.0), (sub + sub)
    assert_eq!(compiled.nodes.len(), 4);
}

#[test]
fn evaluator_handles_arithmetic_and_if_else() {
    let evaluator = DefaultExprEvaluator::new();

    let expr = Expr::IfElse(
        Box::new(Expr::Gt(
            Box::new(Expr::Column("close".to_string())),
            Box::new(Expr::Parameter("threshold".to_string())),
        )),
        Box::new(Expr::Column("close".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(0.0))),
    );
    let compiled = compile_expr(expr).expect("compile if/else");

    let mut inputs = HashMap::new();
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as _,
    );

    let mut params = HashMap::new();
    params.insert("threshold".to_string(), ScalarValue::Float64(1.5));

    let output = evaluator
        .evaluate(&compiled, &inputs, &params)
        .expect("evaluate expression");
    let output = output
        .get("output")
        .expect("single output")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("float output");

    assert_eq!(output.value(0), 0.0);
    assert_eq!(output.value(1), 2.0);
    assert_eq!(output.value(2), 3.0);
}

#[test]
fn evaluator_rolling_mean_follows_window_semantics() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::RollingMean(Box::new(Expr::Column("close".to_string())), DynPeriod::Fixed(2));
    let compiled = compile_expr(expr).expect("compile rolling mean");

    let mut inputs = HashMap::new();
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate expression");
    let output = output
        .get("output")
        .expect("single output")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("float output");

    assert!(output.value(0).is_nan());
    assert_eq!(output.value(1), 1.5);
    assert_eq!(output.value(2), 2.5);
}

#[test]
fn evaluator_rsi_wilder_smoothing() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::Rsi(Box::new(Expr::Column("close".to_string())), DynPeriod::Fixed(3));
    let compiled = compile_expr(expr).expect("compile rsi");

    // Prices: 44, 44.34, 44.09, 43.61, 44.33, 44.83, 45.10
    let prices = vec![44.0, 44.34, 44.09, 43.61, 44.33, 44.83, 45.10];
    let mut inputs = HashMap::new();
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(prices)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate rsi");
    let output = output
        .get("output")
        .expect("single output")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("float output");

    // First 3 values (indices 0..=2) should be NaN (warm-up)
    assert!(output.value(0).is_nan(), "index 0 should be NaN");
    assert!(output.value(1).is_nan(), "index 1 should be NaN");
    assert!(output.value(2).is_nan(), "index 2 should be NaN");

    // Index 3 is the first real RSI value
    let rsi_3 = output.value(3);
    assert!(!rsi_3.is_nan(), "index 3 should have a value");
    assert!((0.0..=100.0).contains(&rsi_3), "RSI should be in [0,100]");

    // Manual verification for period=3:
    // deltas: +0.34, -0.25, -0.48
    // gains: [0.34, 0, 0], losses: [0, 0.25, 0.48]
    // avg_gain = 0.34/3 = 0.1133..., avg_loss = 0.73/3 = 0.2433...
    // RS = 0.1133 / 0.2433 = 0.4657...
    // RSI = 100 - 100/(1 + 0.4657) = 31.78...
    assert!((rsi_3 - 31.78).abs() < 1.0, "RSI(3) ≈ 31.78, got {rsi_3}");

    // Remaining values should also be valid RSI
    for i in 4..7 {
        let v = output.value(i);
        assert!(!v.is_nan(), "index {i} should have a value");
        assert!((0.0..=100.0).contains(&v), "RSI should be in [0,100], got {v}");
    }

    // Index 6 (last): prices keep rising, RSI should be high
    assert!(output.value(6) > 60.0, "uptrend RSI should be > 60");
}

#[test]
fn evaluator_atr_wilder_smoothing() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::Atr(
        Box::new(Expr::Column("high".to_string())),
        Box::new(Expr::Column("low".to_string())),
        Box::new(Expr::Column("close".to_string())),
        DynPeriod::Fixed(3),
    );
    let compiled = compile_expr(expr).expect("compile atr");

    let high = vec![48.70, 48.72, 48.90, 48.87, 48.82, 49.05, 49.20];
    let low = vec![47.79, 48.14, 48.39, 48.37, 48.24, 48.64, 48.94];
    let close = vec![48.16, 48.61, 48.75, 48.63, 48.74, 49.03, 49.07];

    let mut inputs = HashMap::new();
    inputs.insert("high".to_string(), Arc::new(Float64Array::from(high)) as _);
    inputs.insert("low".to_string(), Arc::new(Float64Array::from(low)) as _);
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(close)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate atr");
    let output = output
        .get("output")
        .expect("single output")
        .as_any()
        .downcast_ref::<Float64Array>()
        .expect("float output");

    // First 3 values should be NaN (warm-up: index 0 has no prev close, + period-1 more)
    assert!(output.value(0).is_nan(), "index 0 should be NaN");
    assert!(output.value(1).is_nan(), "index 1 should be NaN");
    assert!(output.value(2).is_nan(), "index 2 should be NaN");

    // Index 3 is the first real ATR value
    let atr_3 = output.value(3);
    assert!(!atr_3.is_nan(), "index 3 should have a value");
    assert!(atr_3 > 0.0, "ATR should be positive, got {atr_3}");

    // Manual verification for period=3:
    // TR[1] = max(48.72-48.14, |48.72-48.16|, |48.14-48.16|) = max(0.58, 0.56, 0.02) = 0.58
    // TR[2] = max(48.90-48.39, |48.90-48.61|, |48.39-48.61|) = max(0.51, 0.29, 0.22) = 0.51
    // TR[3] = max(48.87-48.37, |48.87-48.75|, |48.37-48.75|) = max(0.50, 0.12, 0.38) = 0.50
    // ATR[3] = (0.58 + 0.51 + 0.50) / 3 = 0.53
    assert!((atr_3 - 0.53).abs() < 0.01, "ATR(3) ≈ 0.53, got {atr_3}");

    // Remaining values should all be valid
    for i in 4..7 {
        let v = output.value(i);
        assert!(!v.is_nan(), "index {i} should have a value");
        assert!(v > 0.0, "ATR should be positive, got {v}");
    }
}

#[test]
fn evaluator_linreg_slope_value_r2() {
    let evaluator = DefaultExprEvaluator::new();

    // Perfect linear data: y = 2*x + 1 → slope=2, value at end=9, R²=1.0
    let data = vec![1.0, 3.0, 5.0, 7.0, 9.0];
    let mut inputs = HashMap::new();
    inputs.insert("close".to_string(), Arc::new(Float64Array::from(data)) as _);

    // Test slope (window=3)
    let expr = Expr::LinRegSlope(Box::new(Expr::Column("close".to_string())), DynPeriod::Fixed(3));
    let compiled = compile_expr(expr).expect("compile linreg_slope");
    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate linreg_slope");
    let slope = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // First 2 values NaN (window=3, need 3 points)
    assert!(slope.value(0).is_nan());
    assert!(slope.value(1).is_nan());
    // Window [1,3,5]: slope = 2.0
    assert!(
        (slope.value(2) - 2.0).abs() < 1e-10,
        "slope should be 2.0, got {}",
        slope.value(2)
    );
    // Window [3,5,7]: slope = 2.0
    assert!(
        (slope.value(3) - 2.0).abs() < 1e-10,
        "slope should be 2.0, got {}",
        slope.value(3)
    );
    // Window [5,7,9]: slope = 2.0
    assert!(
        (slope.value(4) - 2.0).abs() < 1e-10,
        "slope should be 2.0, got {}",
        slope.value(4)
    );

    // Test value (window=3) — predicted value at end of window
    let expr = Expr::LinRegValue(Box::new(Expr::Column("close".to_string())), DynPeriod::Fixed(3));
    let compiled = compile_expr(expr).expect("compile linreg_value");
    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate linreg_value");
    let value = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // Window [1,3,5]: mean=3, predicted at end = mean + slope*(w-1)/2 = 3 + 2*1 = 5
    assert!(
        (value.value(2) - 5.0).abs() < 1e-10,
        "value should be 5.0, got {}",
        value.value(2)
    );
    // Window [5,7,9]: predicted at end = 7 + 2*1 = 9
    assert!(
        (value.value(4) - 9.0).abs() < 1e-10,
        "value should be 9.0, got {}",
        value.value(4)
    );

    // Test R² (window=3) — perfect linear: R²=1.0
    let expr = Expr::LinRegR2(Box::new(Expr::Column("close".to_string())), DynPeriod::Fixed(3));
    let compiled = compile_expr(expr).expect("compile linreg_r2");
    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate linreg_r2");
    let r2 = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert!(
        (r2.value(2) - 1.0).abs() < 1e-10,
        "R² should be 1.0 for perfect linear data, got {}",
        r2.value(2)
    );
    assert!(
        (r2.value(4) - 1.0).abs() < 1e-10,
        "R² should be 1.0 for perfect linear data, got {}",
        r2.value(4)
    );

    // Test with non-linear data: R² < 1.0
    let noisy = vec![1.0, 3.0, 2.0, 7.0, 5.0];
    let mut inputs2 = HashMap::new();
    inputs2.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(noisy)) as _,
    );

    let expr = Expr::LinRegR2(Box::new(Expr::Column("close".to_string())), DynPeriod::Fixed(3));
    let compiled = compile_expr(expr).expect("compile linreg_r2 noisy");
    let output = evaluator
        .evaluate(&compiled, &inputs2, &HashMap::new())
        .expect("evaluate linreg_r2 noisy");
    let r2_noisy = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // Window [1,3,2]: not perfectly linear, R² < 1
    let r2_val = r2_noisy.value(2);
    assert!(
        (0.0..1.0).contains(&r2_val),
        "R² should be in [0,1) for noisy data, got {}",
        r2_val
    );
}

#[test]
fn type_inference_validates_numeric_arithmetic() {
    let expr = Expr::Add(
        Box::new(Expr::Column("close".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(1.0))),
    );

    let columns = HashMap::from([("close".to_string(), ValueType::Float64)]);
    let ty = infer_expr_type(&expr, &columns, &HashMap::new()).expect("infer type");
    assert_eq!(ty, ValueType::Float64);
}

// ---------------------------------------------------------------------------
// Cross-sectional multi-symbol tests
// ---------------------------------------------------------------------------

#[test]
fn cross_sectional_mean_multi_computes_row_wise_mean() {
    // 3 symbols × 5 rows
    let a = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]);
    let b = Float64Array::from(vec![10.0, 20.0, 30.0, 40.0, 50.0]);
    let c = Float64Array::from(vec![100.0, 200.0, 300.0, 400.0, 500.0]);

    let results = cross_sectional_mean_multi(&[&a, &b, &c]);
    assert_eq!(results.len(), 3);

    // All symbols get the same mean at each row
    for result in &results {
        assert_eq!(result.len(), 5);
        assert!((result.value(0) - 37.0).abs() < 1e-10); // (1+10+100)/3
        assert!((result.value(1) - 74.0).abs() < 1e-10); // (2+20+200)/3
        assert!((result.value(2) - 111.0).abs() < 1e-10); // (3+30+300)/3
    }
}

#[test]
fn cross_sectional_mean_multi_handles_nan() {
    let a = Float64Array::from(vec![Some(1.0), Some(f64::NAN), Some(3.0)]);
    let b = Float64Array::from(vec![Some(5.0), Some(10.0), None]);

    let results = cross_sectional_mean_multi(&[&a, &b]);
    assert_eq!(results.len(), 2);

    // Row 0: (1+5)/2 = 3.0
    assert!((results[0].value(0) - 3.0).abs() < 1e-10);
    // Row 1: NaN is not finite, so only 10.0 contributes → mean = 10.0
    assert!((results[0].value(1) - 10.0).abs() < 1e-10);
    // Row 2: b is null, only a contributes → mean = 3.0
    assert!((results[0].value(2) - 3.0).abs() < 1e-10);
}

#[test]
fn cross_sectional_rank_multi_basic() {
    // 3 symbols × 4 rows
    let a = Float64Array::from(vec![10.0, 30.0, 50.0, 10.0]);
    let b = Float64Array::from(vec![20.0, 20.0, 40.0, 10.0]);
    let c = Float64Array::from(vec![30.0, 10.0, 30.0, 10.0]);

    let results = cross_sectional_rank_multi(&[&a, &b, &c]);
    assert_eq!(results.len(), 3);

    // Row 0: a=10 (rank 0.0), b=20 (rank 0.5), c=30 (rank 1.0)
    assert!((results[0].value(0) - 0.0).abs() < 1e-10);
    assert!((results[1].value(0) - 0.5).abs() < 1e-10);
    assert!((results[2].value(0) - 1.0).abs() < 1e-10);

    // Row 1: c=10 (rank 0.0), b=20 (rank 0.5), a=30 (rank 1.0)
    assert!((results[0].value(1) - 1.0).abs() < 1e-10);
    assert!((results[1].value(1) - 0.5).abs() < 1e-10);
    assert!((results[2].value(1) - 0.0).abs() < 1e-10);

    // Row 3: all tied at 10.0 → all get rank 0.5
    assert!((results[0].value(3) - 0.5).abs() < 1e-10);
    assert!((results[1].value(3) - 0.5).abs() < 1e-10);
    assert!((results[2].value(3) - 0.5).abs() < 1e-10);
}

#[test]
fn cross_sectional_rank_multi_handles_nan() {
    let a = Float64Array::from(vec![Some(10.0), None]);
    let b = Float64Array::from(vec![Some(20.0), Some(5.0)]);

    let results = cross_sectional_rank_multi(&[&a, &b]);

    // Row 0: a=10 (rank 0.0), b=20 (rank 1.0)
    assert!((results[0].value(0) - 0.0).abs() < 1e-10);
    assert!((results[1].value(0) - 1.0).abs() < 1e-10);

    // Row 1: a=null (NaN output), b=5.0 (single element → rank 0.5)
    assert!(results[0].value(1).is_nan());
    assert!((results[1].value(1) - 0.5).abs() < 1e-10);
}

#[test]
fn has_cross_sectional_ops_detects_correctly() {
    // Temporal-only signal
    let temporal = compile_expr(Expr::RollingMean(
        Box::new(Expr::Column("close".to_string())),
        DynPeriod::Fixed(20),
    ))
    .unwrap();
    assert!(!has_cross_sectional_ops(&temporal));

    // Cross-sectional signal
    let cs = compile_expr(Expr::CrossSectionalRank(Box::new(Expr::Column(
        "momentum".to_string(),
    ))))
    .unwrap();
    assert!(has_cross_sectional_ops(&cs));
}

#[test]
fn extract_cross_section_info_extracts_column() {
    let cs = compile_expr(Expr::CrossSectionalMean(Box::new(Expr::Column(
        "signal_x".to_string(),
    ))))
    .unwrap();
    let info = extract_cross_section_info(&cs).unwrap().unwrap();
    assert_eq!(info.op, CrossSectionKind::Mean);
    assert_eq!(info.input_column, "signal_x");

    let cs2 = compile_expr(Expr::CrossSectionalRank(Box::new(Expr::Column(
        "momentum".to_string(),
    ))))
    .unwrap();
    let info2 = extract_cross_section_info(&cs2).unwrap().unwrap();
    assert_eq!(info2.op, CrossSectionKind::Rank);
    assert_eq!(info2.input_column, "momentum");
}

#[test]
fn extract_cross_section_info_rejects_sub_expression() {
    // CrossSectionalMean(close + 1.0) should be rejected
    let cs = compile_expr(Expr::CrossSectionalMean(Box::new(Expr::Add(
        Box::new(Expr::Column("close".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(1.0))),
    ))))
    .unwrap();
    let result = extract_cross_section_info(&cs);
    assert!(result.is_err());
}

#[test]
fn cross_sectional_evaluator_errors_on_per_symbol_call() {
    let evaluator = DefaultExprEvaluator::new();
    let cs = compile_expr(Expr::CrossSectionalMean(Box::new(Expr::Column(
        "signal".to_string(),
    ))))
    .unwrap();

    let mut inputs = HashMap::new();
    inputs.insert(
        "signal".to_string(),
        Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as _,
    );

    let result = evaluator.evaluate(&cs, &inputs, &HashMap::new());
    assert!(result.is_err());
    assert!(result
        .unwrap_err()
        .to_string()
        .contains("orchestrator-level"));
}

// ---------------------------------------------------------------------------
// SymbolRef tests
// ---------------------------------------------------------------------------

#[test]
fn has_symbol_ref_ops_detects_correctly() {
    let temporal = compile_expr(Expr::RollingMean(
        Box::new(Expr::Column("close".to_string())),
        DynPeriod::Fixed(20),
    ))
    .unwrap();
    assert!(!has_symbol_ref_ops(&temporal));

    let sr = compile_expr(Expr::SymbolRef(
        "BTCUSDT".to_string(),
        Box::new(Expr::Column("close".to_string())),
    ))
    .unwrap();
    assert!(has_symbol_ref_ops(&sr));
    assert!(!has_cross_sectional_ops(&sr));
}

#[test]
fn extract_symbol_ref_info_extracts_column() {
    let sr = compile_expr(Expr::SymbolRef(
        "BTCUSDT".to_string(),
        Box::new(Expr::Column("momentum".to_string())),
    ))
    .unwrap();
    let info = extract_symbol_ref_info(&sr).unwrap().unwrap();
    assert_eq!(info.symbol, "BTCUSDT");
    assert_eq!(info.input_column, "momentum");
}

#[test]
fn extract_symbol_ref_info_rejects_sub_expression() {
    let sr = compile_expr(Expr::SymbolRef(
        "BTCUSDT".to_string(),
        Box::new(Expr::Add(
            Box::new(Expr::Column("close".to_string())),
            Box::new(Expr::Literal(ScalarValue::Float64(1.0))),
        )),
    ))
    .unwrap();
    let result = extract_symbol_ref_info(&sr);
    assert!(result.is_err());
}

#[test]
fn symbol_ref_evaluator_errors_on_per_symbol_call() {
    let evaluator = DefaultExprEvaluator::new();
    let sr = compile_expr(Expr::SymbolRef(
        "BTCUSDT".to_string(),
        Box::new(Expr::Column("close".to_string())),
    ))
    .unwrap();

    let mut inputs = HashMap::new();
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as _,
    );

    let result = evaluator.evaluate(&sr, &inputs, &HashMap::new());
    assert!(result.is_err());
    assert!(result
        .unwrap_err()
        .to_string()
        .contains("orchestrator-level"));
}

// ---------------------------------------------------------------------------
// Datetime extraction tests
// ---------------------------------------------------------------------------

/// Helper: create a nanosecond timestamp for a given UTC datetime.
fn utc_ns(year: i32, month: u32, day: u32, hour: u32, min: u32, sec: u32) -> i64 {
    use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
    let dt = NaiveDateTime::new(
        NaiveDate::from_ymd_opt(year, month, day).unwrap(),
        NaiveTime::from_hms_opt(hour, min, sec).unwrap(),
    );
    dt.and_utc().timestamp_nanos_opt().unwrap()
}

#[test]
fn datetime_hour_extracts_correctly() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::Hour(Box::new(Expr::Column("timestamp".to_string())));
    let compiled = compile_expr(expr).expect("compile Hour");

    // Timestamps: 00:00, 09:30, 14:15, 23:59 UTC
    let ts = vec![
        utc_ns(2024, 3, 15, 0, 0, 0),
        utc_ns(2024, 3, 15, 9, 30, 0),
        utc_ns(2024, 3, 15, 14, 15, 0),
        utc_ns(2024, 3, 15, 23, 59, 0),
    ];
    let mut inputs = HashMap::new();
    inputs.insert(
        "timestamp".to_string(),
        Arc::new(TimestampNanosecondArray::from(ts)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate Hour");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert_eq!(output.value(0), 0.0);
    assert_eq!(output.value(1), 9.0);
    assert_eq!(output.value(2), 14.0);
    assert_eq!(output.value(3), 23.0);
}

#[test]
fn datetime_minute_extracts_correctly() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::Minute(Box::new(Expr::Column("timestamp".to_string())));
    let compiled = compile_expr(expr).expect("compile Minute");

    let ts = vec![
        utc_ns(2024, 1, 1, 12, 0, 0),
        utc_ns(2024, 1, 1, 12, 30, 0),
        utc_ns(2024, 1, 1, 12, 59, 0),
    ];
    let mut inputs = HashMap::new();
    inputs.insert(
        "timestamp".to_string(),
        Arc::new(TimestampNanosecondArray::from(ts)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate Minute");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert_eq!(output.value(0), 0.0);
    assert_eq!(output.value(1), 30.0);
    assert_eq!(output.value(2), 59.0);
}

#[test]
fn datetime_day_of_week_extracts_correctly() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::DayOfWeek(Box::new(Expr::Column("timestamp".to_string())));
    let compiled = compile_expr(expr).expect("compile DayOfWeek");

    // 2024-01-01 = Monday (0), 2024-01-05 = Friday (4), 2024-01-07 = Sunday (6)
    let ts = vec![
        utc_ns(2024, 1, 1, 12, 0, 0), // Monday
        utc_ns(2024, 1, 5, 12, 0, 0), // Friday
        utc_ns(2024, 1, 7, 12, 0, 0), // Sunday
        utc_ns(2024, 1, 3, 12, 0, 0), // Wednesday
    ];
    let mut inputs = HashMap::new();
    inputs.insert(
        "timestamp".to_string(),
        Arc::new(TimestampNanosecondArray::from(ts)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate DayOfWeek");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert_eq!(output.value(0), 0.0); // Monday
    assert_eq!(output.value(1), 4.0); // Friday
    assert_eq!(output.value(2), 6.0); // Sunday
    assert_eq!(output.value(3), 2.0); // Wednesday
}

#[test]
fn datetime_month_extracts_correctly() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::Month(Box::new(Expr::Column("timestamp".to_string())));
    let compiled = compile_expr(expr).expect("compile Month");

    let ts = vec![
        utc_ns(2024, 1, 15, 0, 0, 0),
        utc_ns(2024, 6, 15, 0, 0, 0),
        utc_ns(2024, 12, 31, 23, 59, 59),
    ];
    let mut inputs = HashMap::new();
    inputs.insert(
        "timestamp".to_string(),
        Arc::new(TimestampNanosecondArray::from(ts)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate Month");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert_eq!(output.value(0), 1.0);
    assert_eq!(output.value(1), 6.0);
    assert_eq!(output.value(2), 12.0);
}

#[test]
fn datetime_day_of_month_extracts_correctly() {
    let evaluator = DefaultExprEvaluator::new();
    let expr = Expr::DayOfMonth(Box::new(Expr::Column("timestamp".to_string())));
    let compiled = compile_expr(expr).expect("compile DayOfMonth");

    let ts = vec![
        utc_ns(2024, 1, 1, 0, 0, 0),
        utc_ns(2024, 2, 29, 12, 0, 0), // leap year
        utc_ns(2024, 3, 31, 23, 59, 59),
    ];
    let mut inputs = HashMap::new();
    inputs.insert(
        "timestamp".to_string(),
        Arc::new(TimestampNanosecondArray::from(ts)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate DayOfMonth");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert_eq!(output.value(0), 1.0);
    assert_eq!(output.value(1), 29.0);
    assert_eq!(output.value(2), 31.0);
}

#[test]
fn datetime_composable_with_arithmetic_and_filters() {
    // Test: filter to only trade during US market hours (14-21 UTC)
    let evaluator = DefaultExprEvaluator::new();

    let hour = Expr::Hour(Box::new(Expr::Column("timestamp".to_string())));
    // US hours: hour >= 14 AND hour < 21
    let in_us_hours = Expr::And(
        Box::new(Expr::Gt(
            Box::new(hour.clone()),
            Box::new(Expr::Literal(ScalarValue::Float64(13.5))),
        )),
        Box::new(Expr::Lt(
            Box::new(hour),
            Box::new(Expr::Literal(ScalarValue::Float64(21.0))),
        )),
    );
    // signal = if_else(in_us_hours, close, 0.0)
    let signal = Expr::IfElse(
        Box::new(in_us_hours),
        Box::new(Expr::Column("close".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(0.0))),
    );
    let compiled = compile_expr(signal).expect("compile session filter");

    let ts = vec![
        utc_ns(2024, 1, 15, 10, 0, 0),  // 10 UTC — outside US hours
        utc_ns(2024, 1, 15, 14, 0, 0),  // 14 UTC — inside US hours
        utc_ns(2024, 1, 15, 18, 30, 0), // 18:30 UTC — inside
        utc_ns(2024, 1, 15, 22, 0, 0),  // 22 UTC — outside
    ];
    let closes = vec![100.0, 101.0, 102.0, 103.0];

    let mut inputs = HashMap::new();
    inputs.insert(
        "timestamp".to_string(),
        Arc::new(TimestampNanosecondArray::from(ts)) as _,
    );
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(closes)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate session filter");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert_eq!(output.value(0), 0.0); // outside — zeroed
    assert_eq!(output.value(1), 101.0); // inside
    assert_eq!(output.value(2), 102.0); // inside
    assert_eq!(output.value(3), 0.0); // outside — zeroed
}

#[test]
fn datetime_day_of_week_filter_weekdays_only() {
    // Filter: only trade on weekdays (Mon-Fri, i.e. day_of_week < 5)
    let evaluator = DefaultExprEvaluator::new();

    let dow = Expr::DayOfWeek(Box::new(Expr::Column("timestamp".to_string())));
    let is_weekday = Expr::Lt(
        Box::new(dow),
        Box::new(Expr::Literal(ScalarValue::Float64(5.0))),
    );
    let signal = Expr::IfElse(
        Box::new(is_weekday),
        Box::new(Expr::Column("close".to_string())),
        Box::new(Expr::Literal(ScalarValue::Float64(0.0))),
    );
    let compiled = compile_expr(signal).expect("compile weekday filter");

    // 2024-01-01=Mon, 2024-01-06=Sat, 2024-01-07=Sun, 2024-01-08=Mon
    let ts = vec![
        utc_ns(2024, 1, 1, 12, 0, 0),
        utc_ns(2024, 1, 6, 12, 0, 0),
        utc_ns(2024, 1, 7, 12, 0, 0),
        utc_ns(2024, 1, 8, 12, 0, 0),
    ];
    let closes = vec![50.0, 51.0, 52.0, 53.0];

    let mut inputs = HashMap::new();
    inputs.insert(
        "timestamp".to_string(),
        Arc::new(TimestampNanosecondArray::from(ts)) as _,
    );
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(closes)) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate weekday filter");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    assert_eq!(output.value(0), 50.0); // Monday — trade
    assert_eq!(output.value(1), 0.0); // Saturday — skip
    assert_eq!(output.value(2), 0.0); // Sunday — skip
    assert_eq!(output.value(3), 53.0); // Monday — trade
}

// ---------------------------------------------------------------------------
// Scan (stateful fold) tests
// ---------------------------------------------------------------------------

#[test]
fn test_scan_cumsum() {
    // Cumulative sum via scan: state = running sum, output = running sum
    let evaluator = DefaultExprEvaluator::new();

    let expr = Expr::Scan {
        state_names: vec!["total".to_string()],
        init_exprs: vec![Box::new(Expr::Literal(ScalarValue::Float64(0.0)))],
        update_names: vec!["total".to_string()],
        update_exprs: vec![Box::new(Expr::Add(
            Box::new(Expr::ScanPrev("total".to_string())),
            Box::new(Expr::Column("value".to_string())),
        ))],
        output: "total".to_string(),
    };
    let compiled = compile_expr(expr).expect("compile scan cumsum");

    let mut inputs = HashMap::new();
    inputs.insert(
        "value".to_string(),
        Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0])) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate scan cumsum");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // cumsum: 1, 3, 6, 10, 15
    assert_eq!(output.value(0), 1.0);
    assert_eq!(output.value(1), 3.0);
    assert_eq!(output.value(2), 6.0);
    assert_eq!(output.value(3), 10.0);
    assert_eq!(output.value(4), 15.0);
}

#[test]
fn test_scan_ema() {
    // EMA via scan: ema[t] = alpha * close[t] + (1 - alpha) * ema[t-1]
    let evaluator = DefaultExprEvaluator::new();
    let alpha = 0.1;

    let expr = Expr::Scan {
        state_names: vec!["ema".to_string()],
        init_exprs: vec![Box::new(Expr::Column("close".to_string()))],
        update_names: vec!["ema".to_string()],
        update_exprs: vec![Box::new(Expr::Add(
            Box::new(Expr::Mul(
                Box::new(Expr::Literal(ScalarValue::Float64(alpha))),
                Box::new(Expr::Column("close".to_string())),
            )),
            Box::new(Expr::Mul(
                Box::new(Expr::Literal(ScalarValue::Float64(1.0 - alpha))),
                Box::new(Expr::ScanPrev("ema".to_string())),
            )),
        ))],
        output: "ema".to_string(),
    };
    let compiled = compile_expr(expr).expect("compile scan ema");

    let close_data = vec![100.0, 102.0, 101.0, 103.0, 105.0];
    let mut inputs = HashMap::new();
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(close_data.clone())) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate scan ema");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // Hand-compute EMA:
    // t=0: alpha*100 + (1-alpha)*100 = 100.0 (init is close[0]=100)
    // t=1: 0.1*102 + 0.9*100.0 = 100.2
    // t=2: 0.1*101 + 0.9*100.2 = 100.28
    // t=3: 0.1*103 + 0.9*100.28 = 100.552
    // t=4: 0.1*105 + 0.9*100.552 = 100.9968
    let mut expected = Vec::new();
    let mut ema = close_data[0];
    for &c in &close_data {
        ema = alpha * c + (1.0 - alpha) * ema;
        expected.push(ema);
    }

    for (i, exp) in expected.iter().enumerate() {
        assert!(
            (output.value(i) - exp).abs() < 1e-10,
            "EMA mismatch at t={}: got {}, expected {}",
            i,
            output.value(i),
            exp
        );
    }
}

#[test]
fn test_scan_with_parameters() {
    // scan that uses Parameter references for alpha/beta
    let evaluator = DefaultExprEvaluator::new();

    let expr = Expr::Scan {
        state_names: vec!["ema".to_string()],
        init_exprs: vec![Box::new(Expr::Column("close".to_string()))],
        update_names: vec!["ema".to_string()],
        update_exprs: vec![Box::new(Expr::Add(
            Box::new(Expr::Mul(
                Box::new(Expr::Parameter("alpha".to_string())),
                Box::new(Expr::Column("close".to_string())),
            )),
            Box::new(Expr::Mul(
                Box::new(Expr::Sub(
                    Box::new(Expr::Literal(ScalarValue::Float64(1.0))),
                    Box::new(Expr::Parameter("alpha".to_string())),
                )),
                Box::new(Expr::ScanPrev("ema".to_string())),
            )),
        ))],
        output: "ema".to_string(),
    };
    let compiled = compile_expr(expr).expect("compile scan with params");

    let mut inputs = HashMap::new();
    inputs.insert(
        "close".to_string(),
        Arc::new(Float64Array::from(vec![100.0, 110.0, 105.0])) as _,
    );

    let mut params = HashMap::new();
    params.insert("alpha".to_string(), ScalarValue::Float64(0.5));

    let output = evaluator
        .evaluate(&compiled, &inputs, &params)
        .expect("evaluate scan with params");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // alpha=0.5: ema[0]=0.5*100+0.5*100=100, ema[1]=0.5*110+0.5*100=105, ema[2]=0.5*105+0.5*105=105
    assert!((output.value(0) - 100.0).abs() < 1e-10);
    assert!((output.value(1) - 105.0).abs() < 1e-10);
    assert!((output.value(2) - 105.0).abs() < 1e-10);
}

#[test]
fn test_scan_kalman() {
    // Kalman filter: x[t] = x[t-1] + k*(obs - x[t-1])
    // where k = p_pred / (p_pred + R), p_pred = p[t-1] + Q
    let evaluator = DefaultExprEvaluator::new();
    let q = 1e-2;
    let r = 1.0;

    let expr = Expr::Scan {
        state_names: vec!["x".to_string(), "p".to_string()],
        init_exprs: vec![
            Box::new(Expr::Column("obs".to_string())),
            Box::new(Expr::Literal(ScalarValue::Float64(1.0))),
        ],
        update_names: vec![
            "p_pred".to_string(),
            "k".to_string(),
            "x".to_string(),
            "p".to_string(),
        ],
        update_exprs: vec![
            // p_pred = p[t-1] + Q
            Box::new(Expr::Add(
                Box::new(Expr::ScanPrev("p".to_string())),
                Box::new(Expr::Literal(ScalarValue::Float64(q))),
            )),
            // k = p_pred / (p_pred + R)
            Box::new(Expr::Div(
                Box::new(Expr::ScanVar("p_pred".to_string())),
                Box::new(Expr::Add(
                    Box::new(Expr::ScanVar("p_pred".to_string())),
                    Box::new(Expr::Literal(ScalarValue::Float64(r))),
                )),
            )),
            // x = x[t-1] + k * (obs - x[t-1])
            Box::new(Expr::Add(
                Box::new(Expr::ScanPrev("x".to_string())),
                Box::new(Expr::Mul(
                    Box::new(Expr::ScanVar("k".to_string())),
                    Box::new(Expr::Sub(
                        Box::new(Expr::Column("obs".to_string())),
                        Box::new(Expr::ScanPrev("x".to_string())),
                    )),
                )),
            )),
            // p = (1 - k) * p_pred
            Box::new(Expr::Mul(
                Box::new(Expr::Sub(
                    Box::new(Expr::Literal(ScalarValue::Float64(1.0))),
                    Box::new(Expr::ScanVar("k".to_string())),
                )),
                Box::new(Expr::ScanVar("p_pred".to_string())),
            )),
        ],
        output: "x".to_string(),
    };
    let compiled = compile_expr(expr).expect("compile kalman scan");

    let observations = vec![100.0, 102.0, 101.0, 103.0, 105.0, 104.0, 106.0, 108.0];
    let mut inputs = HashMap::new();
    inputs.insert(
        "obs".to_string(),
        Arc::new(Float64Array::from(observations.clone())) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate kalman");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // Verify against hand-computed Kalman filter
    let mut x = observations[0];
    let mut p = 1.0;
    let mut expected = Vec::new();
    for &obs in &observations {
        let p_pred = p + q;
        let k = p_pred / (p_pred + r);
        x = x + k * (obs - x);
        p = (1.0 - k) * p_pred;
        expected.push(x);
    }

    for (i, exp) in expected.iter().enumerate() {
        assert!(
            (output.value(i) - exp).abs() < 1e-10,
            "Kalman mismatch at t={}: got {}, expected {}",
            i,
            output.value(i),
            exp
        );
    }

    // The Kalman filter should smooth the signal — check it's between min and max of obs
    let min_obs = observations.iter().cloned().fold(f64::INFINITY, f64::min);
    let max_obs = observations
        .iter()
        .cloned()
        .fold(f64::NEG_INFINITY, f64::max);
    for i in 0..observations.len() {
        assert!(
            output.value(i) >= min_obs - 1.0 && output.value(i) <= max_obs + 1.0,
            "Kalman output out of expected range at t={}",
            i
        );
    }
}

#[test]
fn test_scan_serde_roundtrip() {
    // Verify Scan serializes/deserializes correctly via serde_json
    let scan_expr = Expr::Scan {
        state_names: vec!["x".to_string()],
        init_exprs: vec![Box::new(Expr::Column("close".to_string()))],
        update_names: vec!["x".to_string()],
        update_exprs: vec![Box::new(Expr::Add(
            Box::new(Expr::ScanPrev("x".to_string())),
            Box::new(Expr::Column("close".to_string())),
        ))],
        output: "x".to_string(),
    };

    let json = serde_json::to_string(&scan_expr).expect("serialize Scan");
    let deserialized: Expr = serde_json::from_str(&json).expect("deserialize Scan");
    assert_eq!(scan_expr, deserialized);

    // Also test ScanPrev/ScanVar standalone serialization
    let prev = Expr::ScanPrev("x".to_string());
    let json = serde_json::to_string(&prev).expect("serialize ScanPrev");
    let de: Expr = serde_json::from_str(&json).expect("deserialize ScanPrev");
    assert_eq!(prev, de);

    let var = Expr::ScanVar("k".to_string());
    let json = serde_json::to_string(&var).expect("serialize ScanVar");
    let de: Expr = serde_json::from_str(&json).expect("deserialize ScanVar");
    assert_eq!(var, de);
}

#[test]
fn test_scan_if_else_inside() {
    // Test IfElse inside scan: clamp accumulator to max 10
    let evaluator = DefaultExprEvaluator::new();

    let expr = Expr::Scan {
        state_names: vec!["acc".to_string()],
        init_exprs: vec![Box::new(Expr::Literal(ScalarValue::Float64(0.0)))],
        update_names: vec!["raw".to_string(), "acc".to_string()],
        update_exprs: vec![
            // raw = prev(acc) + value
            Box::new(Expr::Add(
                Box::new(Expr::ScanPrev("acc".to_string())),
                Box::new(Expr::Column("value".to_string())),
            )),
            // acc = if raw > 10 then 10 else raw
            Box::new(Expr::IfElse(
                Box::new(Expr::Gt(
                    Box::new(Expr::ScanVar("raw".to_string())),
                    Box::new(Expr::Literal(ScalarValue::Float64(10.0))),
                )),
                Box::new(Expr::Literal(ScalarValue::Float64(10.0))),
                Box::new(Expr::ScanVar("raw".to_string())),
            )),
        ],
        output: "acc".to_string(),
    };
    let compiled = compile_expr(expr).expect("compile scan with if_else");

    let mut inputs = HashMap::new();
    inputs.insert(
        "value".to_string(),
        Arc::new(Float64Array::from(vec![3.0, 3.0, 3.0, 3.0, 3.0])) as _,
    );

    let output = evaluator
        .evaluate(&compiled, &inputs, &HashMap::new())
        .expect("evaluate scan with if_else");
    let output = output
        .get("output")
        .unwrap()
        .as_any()
        .downcast_ref::<Float64Array>()
        .unwrap();

    // t=0: raw=0+3=3, acc=3
    // t=1: raw=3+3=6, acc=6
    // t=2: raw=6+3=9, acc=9
    // t=3: raw=9+3=12>10, acc=10
    // t=4: raw=10+3=13>10, acc=10
    assert_eq!(output.value(0), 3.0);
    assert_eq!(output.value(1), 6.0);
    assert_eq!(output.value(2), 9.0);
    assert_eq!(output.value(3), 10.0);
    assert_eq!(output.value(4), 10.0);
}
