//! Information-Driven Bars from Marcos Lopez de Prado's
//! "Advances in Financial Machine Learning" (Chapter 2).
//!
//! Provides 8 bar-sampling methods:
//! - Standard: tick bars, volume bars, dollar bars
//! - Information-driven: tick/volume imbalance bars, tick/volume run bars
//!
//! All functions are pure Rust + PyO3, no external dependencies.
//! Designed for tick-by-tick data (millions of points).

use pyo3::prelude::*;

// ===========================================================================
// Bar type
// ===========================================================================

/// A single OHLCV bar sampled from tick data.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct Bar {
    #[pyo3(get)]
    pub timestamp: f64,
    #[pyo3(get)]
    pub open: f64,
    #[pyo3(get)]
    pub high: f64,
    #[pyo3(get)]
    pub low: f64,
    #[pyo3(get)]
    pub close: f64,
    #[pyo3(get)]
    pub volume: f64,
    #[pyo3(get)]
    pub vwap: f64,
    #[pyo3(get)]
    pub n_ticks: u64,
}

#[pymethods]
impl Bar {
    fn __repr__(&self) -> String {
        format!(
            "Bar(ts={:.3}, o={:.6}, h={:.6}, l={:.6}, c={:.6}, vol={:.2}, vwap={:.6}, ticks={})",
            self.timestamp, self.open, self.high, self.low, self.close,
            self.volume, self.vwap, self.n_ticks
        )
    }
}

// ===========================================================================
// Accumulator helper
// ===========================================================================

/// Accumulates OHLCV + VWAP state for building a bar.
struct BarAccumulator {
    first_ts: f64,
    open: f64,
    high: f64,
    low: f64,
    close: f64,
    volume: f64,
    pv_sum: f64, // sum(price * volume) for VWAP
    n_ticks: u64,
}

impl BarAccumulator {
    #[inline]
    fn new() -> Self {
        Self {
            first_ts: 0.0,
            open: 0.0,
            high: f64::NEG_INFINITY,
            low: f64::INFINITY,
            close: 0.0,
            volume: 0.0,
            pv_sum: 0.0,
            n_ticks: 0,
        }
    }

    #[inline]
    fn update(&mut self, ts: f64, price: f64, vol: f64) {
        if self.n_ticks == 0 {
            self.first_ts = ts;
            self.open = price;
            self.high = price;
            self.low = price;
        } else {
            if price > self.high {
                self.high = price;
            }
            if price < self.low {
                self.low = price;
            }
        }
        self.close = price;
        self.volume += vol;
        self.pv_sum += price * vol;
        self.n_ticks += 1;
    }

    #[inline]
    fn emit(&self) -> Bar {
        let vwap = if self.volume > 0.0 {
            self.pv_sum / self.volume
        } else {
            self.close
        };
        Bar {
            timestamp: self.first_ts,
            open: self.open,
            high: self.high,
            low: self.low,
            close: self.close,
            volume: self.volume,
            vwap,
            n_ticks: self.n_ticks,
        }
    }

    #[inline]
    fn is_empty(&self) -> bool {
        self.n_ticks == 0
    }

    #[inline]
    fn reset(&mut self) {
        self.first_ts = 0.0;
        self.open = 0.0;
        self.high = f64::NEG_INFINITY;
        self.low = f64::INFINITY;
        self.close = 0.0;
        self.volume = 0.0;
        self.pv_sum = 0.0;
        self.n_ticks = 0;
    }
}

// ===========================================================================
// Validation helper
// ===========================================================================

#[inline]
fn validate_inputs(
    timestamps: &[f64],
    prices: &[f64],
    volumes: &[f64],
) -> PyResult<()> {
    if timestamps.len() != prices.len() || timestamps.len() != volumes.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "timestamps, prices, and volumes must have the same length",
        ));
    }
    Ok(())
}

// ===========================================================================
// Tick direction helper
// ===========================================================================

/// Compute tick direction: sign(price_t - price_{t-1}).
/// b_0 = 1 by convention. If price unchanged, carry forward previous direction.
#[inline]
fn tick_direction(price: f64, prev_price: f64, prev_b: f64) -> f64 {
    let diff = price - prev_price;
    if diff > 0.0 {
        1.0
    } else if diff < 0.0 {
        -1.0
    } else {
        prev_b
    }
}

// ===========================================================================
// EWM helper for expected bar length / threshold
// ===========================================================================

/// Exponentially weighted moving average update.
/// span = most recent bar length, alpha = 2/(span+1).
#[inline]
fn ewm_update(prev_estimate: f64, new_value: f64, span: f64) -> f64 {
    if span <= 0.0 {
        return new_value;
    }
    let alpha = 2.0 / (span + 1.0);
    alpha * new_value + (1.0 - alpha) * prev_estimate
}

// ===========================================================================
// Standard bars
// ===========================================================================

/// Sample a new bar every `threshold` ticks.
///
/// Args:
///     timestamps: tick timestamps (f64, e.g. unix epoch seconds)
///     prices: tick prices
///     volumes: tick volumes
///     threshold: number of ticks per bar
#[pyfunction]
pub fn tick_bars(
    timestamps: Vec<f64>,
    prices: Vec<f64>,
    volumes: Vec<f64>,
    threshold: u64,
) -> PyResult<Vec<Bar>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    validate_inputs(&timestamps, &prices, &volumes)?;
    let n = timestamps.len();
    if n == 0 || threshold == 0 {
        return Ok(vec![]);
    }

    let mut bars = Vec::with_capacity(n / threshold as usize + 1);
    let mut acc = BarAccumulator::new();

    for i in 0..n {
        acc.update(timestamps[i], prices[i], volumes[i]);
        if acc.n_ticks >= threshold {
            bars.push(acc.emit());
            acc.reset();
        }
    }

    // Emit final partial bar if any ticks remain
    if !acc.is_empty() {
        bars.push(acc.emit());
    }

    Ok(bars)
}

/// Sample a new bar every time cumulative volume exceeds `threshold`.
///
/// Args:
///     timestamps: tick timestamps
///     prices: tick prices
///     volumes: tick volumes
///     threshold: cumulative volume threshold per bar
#[pyfunction]
pub fn volume_bars(
    timestamps: Vec<f64>,
    prices: Vec<f64>,
    volumes: Vec<f64>,
    threshold: f64,
) -> PyResult<Vec<Bar>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    validate_inputs(&timestamps, &prices, &volumes)?;
    let n = timestamps.len();
    if n == 0 || threshold <= 0.0 {
        return Ok(vec![]);
    }

    let mut bars = Vec::with_capacity(n / 100 + 1);
    let mut acc = BarAccumulator::new();

    for i in 0..n {
        acc.update(timestamps[i], prices[i], volumes[i]);
        if acc.volume >= threshold {
            bars.push(acc.emit());
            acc.reset();
        }
    }

    if !acc.is_empty() {
        bars.push(acc.emit());
    }

    Ok(bars)
}

/// Sample a new bar every time cumulative dollar volume (price * volume) exceeds `threshold`.
///
/// Args:
///     timestamps: tick timestamps
///     prices: tick prices
///     volumes: tick volumes
///     threshold: cumulative dollar volume threshold per bar
#[pyfunction]
pub fn dollar_bars(
    timestamps: Vec<f64>,
    prices: Vec<f64>,
    volumes: Vec<f64>,
    threshold: f64,
) -> PyResult<Vec<Bar>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    validate_inputs(&timestamps, &prices, &volumes)?;
    let n = timestamps.len();
    if n == 0 || threshold <= 0.0 {
        return Ok(vec![]);
    }

    let mut bars = Vec::with_capacity(n / 100 + 1);
    let mut acc = BarAccumulator::new();
    let mut cum_dollar: f64 = 0.0;

    for i in 0..n {
        acc.update(timestamps[i], prices[i], volumes[i]);
        cum_dollar += prices[i] * volumes[i];
        if cum_dollar >= threshold {
            bars.push(acc.emit());
            acc.reset();
            cum_dollar = 0.0;
        }
    }

    if !acc.is_empty() {
        bars.push(acc.emit());
    }

    Ok(bars)
}

// ===========================================================================
// Information-driven bars: Imbalance bars (Chapter 2.3.2)
// ===========================================================================

/// Tick Imbalance Bars (TIBs) — Chapter 2.3.2.1.
///
/// Samples a new bar when the absolute cumulative tick imbalance
/// exceeds a dynamically estimated threshold.
///
/// - b_t = sign(price_t - price_{t-1}), b_0 = 1
/// - theta_t = sum(b_i) from start of current bar
/// - New bar when |theta_t| >= expected_T
/// - expected_T updated via EWM of bar lengths
///
/// Args:
///     timestamps: tick timestamps
///     prices: tick prices
///     volumes: tick volumes
///     initial_estimate: initial expected |imbalance| threshold (E_0[T])
#[pyfunction]
pub fn tick_imbalance_bars(
    timestamps: Vec<f64>,
    prices: Vec<f64>,
    volumes: Vec<f64>,
    initial_estimate: f64,
) -> PyResult<Vec<Bar>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    validate_inputs(&timestamps, &prices, &volumes)?;
    let n = timestamps.len();
    if n == 0 || initial_estimate <= 0.0 {
        return Ok(vec![]);
    }

    let mut bars = Vec::with_capacity(n / 100 + 1);
    let mut acc = BarAccumulator::new();
    let mut theta: f64 = 0.0; // cumulative tick imbalance
    let mut expected_t = initial_estimate;
    let mut prev_b: f64 = 1.0; // first tick direction convention
    let mut prev_price: f64 = 0.0;
    let mut bar_ticks: f64 = 0.0; // ticks in current bar (for EWM span)
    let mut prev_bar_len: f64 = initial_estimate; // previous bar length for EWM span

    for i in 0..n {
        let b = if i == 0 {
            1.0 // b_0 = 1 by convention
        } else {
            tick_direction(prices[i], prev_price, prev_b)
        };

        acc.update(timestamps[i], prices[i], volumes[i]);
        theta += b;
        bar_ticks += 1.0;

        if theta.abs() >= expected_t {
            bars.push(acc.emit());
            // Update expected threshold via EWM of bar lengths
            expected_t = ewm_update(expected_t, bar_ticks, prev_bar_len);
            prev_bar_len = bar_ticks;
            acc.reset();
            theta = 0.0;
            bar_ticks = 0.0;
        }

        prev_b = b;
        prev_price = prices[i];
    }

    if !acc.is_empty() {
        bars.push(acc.emit());
    }

    Ok(bars)
}

/// Volume Imbalance Bars (VIBs) — Chapter 2.3.2.2.
///
/// Same as tick imbalance bars, but weights by volume:
/// theta_t = sum(b_i * v_i) from start of current bar.
///
/// Args:
///     timestamps: tick timestamps
///     prices: tick prices
///     volumes: tick volumes
///     initial_estimate: initial expected |volume imbalance| threshold
#[pyfunction]
pub fn volume_imbalance_bars(
    timestamps: Vec<f64>,
    prices: Vec<f64>,
    volumes: Vec<f64>,
    initial_estimate: f64,
) -> PyResult<Vec<Bar>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    validate_inputs(&timestamps, &prices, &volumes)?;
    let n = timestamps.len();
    if n == 0 || initial_estimate <= 0.0 {
        return Ok(vec![]);
    }

    let mut bars = Vec::with_capacity(n / 100 + 1);
    let mut acc = BarAccumulator::new();
    let mut theta: f64 = 0.0; // cumulative volume imbalance
    let mut expected_t = initial_estimate;
    let mut prev_b: f64 = 1.0;
    let mut prev_price: f64 = 0.0;
    let mut bar_ticks: f64 = 0.0;
    let mut prev_bar_len: f64 = initial_estimate;

    for i in 0..n {
        let b = if i == 0 {
            1.0
        } else {
            tick_direction(prices[i], prev_price, prev_b)
        };

        acc.update(timestamps[i], prices[i], volumes[i]);
        theta += b * volumes[i];
        bar_ticks += 1.0;

        if theta.abs() >= expected_t {
            bars.push(acc.emit());
            expected_t = ewm_update(expected_t, bar_ticks, prev_bar_len);
            prev_bar_len = bar_ticks;
            acc.reset();
            theta = 0.0;
            bar_ticks = 0.0;
        }

        prev_b = b;
        prev_price = prices[i];
    }

    if !acc.is_empty() {
        bars.push(acc.emit());
    }

    Ok(bars)
}

// ===========================================================================
// Information-driven bars: Run bars (Chapter 2.3.3)
// ===========================================================================

/// Tick Run Bars (TRBs) — Chapter 2.3.3.
///
/// Counts runs of buy ticks and sell ticks separately.
/// Samples a new bar when max(buy_run, sell_run) >= expected threshold.
///
/// Args:
///     timestamps: tick timestamps
///     prices: tick prices
///     volumes: tick volumes
///     initial_estimate: initial expected max run length
#[pyfunction]
pub fn tick_run_bars(
    timestamps: Vec<f64>,
    prices: Vec<f64>,
    volumes: Vec<f64>,
    initial_estimate: f64,
) -> PyResult<Vec<Bar>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    validate_inputs(&timestamps, &prices, &volumes)?;
    let n = timestamps.len();
    if n == 0 || initial_estimate <= 0.0 {
        return Ok(vec![]);
    }

    let mut bars = Vec::with_capacity(n / 100 + 1);
    let mut acc = BarAccumulator::new();
    let mut buy_run: f64 = 0.0; // count of buy ticks in current bar
    let mut sell_run: f64 = 0.0; // count of sell ticks in current bar
    let mut expected_t = initial_estimate;
    let mut prev_b: f64 = 1.0;
    let mut prev_price: f64 = 0.0;
    let mut bar_ticks: f64 = 0.0;
    let mut prev_bar_len: f64 = initial_estimate;

    for i in 0..n {
        let b = if i == 0 {
            1.0
        } else {
            tick_direction(prices[i], prev_price, prev_b)
        };

        acc.update(timestamps[i], prices[i], volumes[i]);
        bar_ticks += 1.0;

        if b > 0.0 {
            buy_run += 1.0;
        } else {
            sell_run += 1.0;
        }

        let max_run = buy_run.max(sell_run);
        if max_run >= expected_t {
            bars.push(acc.emit());
            expected_t = ewm_update(expected_t, bar_ticks, prev_bar_len);
            prev_bar_len = bar_ticks;
            acc.reset();
            buy_run = 0.0;
            sell_run = 0.0;
            bar_ticks = 0.0;
        }

        prev_b = b;
        prev_price = prices[i];
    }

    if !acc.is_empty() {
        bars.push(acc.emit());
    }

    Ok(bars)
}

/// Volume Run Bars (VRBs) — Chapter 2.3.3.
///
/// Same as tick run bars, but accumulates volume in buy/sell runs
/// rather than counting ticks.
///
/// Args:
///     timestamps: tick timestamps
///     prices: tick prices
///     volumes: tick volumes
///     initial_estimate: initial expected max volume run
#[pyfunction]
pub fn volume_run_bars(
    timestamps: Vec<f64>,
    prices: Vec<f64>,
    volumes: Vec<f64>,
    initial_estimate: f64,
) -> PyResult<Vec<Bar>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    validate_inputs(&timestamps, &prices, &volumes)?;
    let n = timestamps.len();
    if n == 0 || initial_estimate <= 0.0 {
        return Ok(vec![]);
    }

    let mut bars = Vec::with_capacity(n / 100 + 1);
    let mut acc = BarAccumulator::new();
    let mut buy_vol_run: f64 = 0.0; // cumulative buy volume in current bar
    let mut sell_vol_run: f64 = 0.0; // cumulative sell volume in current bar
    let mut expected_t = initial_estimate;
    let mut prev_b: f64 = 1.0;
    let mut prev_price: f64 = 0.0;
    let mut bar_ticks: f64 = 0.0;
    let mut prev_bar_len: f64 = initial_estimate;

    for i in 0..n {
        let b = if i == 0 {
            1.0
        } else {
            tick_direction(prices[i], prev_price, prev_b)
        };

        acc.update(timestamps[i], prices[i], volumes[i]);
        bar_ticks += 1.0;

        if b > 0.0 {
            buy_vol_run += volumes[i];
        } else {
            sell_vol_run += volumes[i];
        }

        let max_run = buy_vol_run.max(sell_vol_run);
        if max_run >= expected_t {
            bars.push(acc.emit());
            expected_t = ewm_update(expected_t, bar_ticks, prev_bar_len);
            prev_bar_len = bar_ticks;
            acc.reset();
            buy_vol_run = 0.0;
            sell_vol_run = 0.0;
            bar_ticks = 0.0;
        }

        prev_b = b;
        prev_price = prices[i];
    }

    if !acc.is_empty() {
        bars.push(acc.emit());
    }

    Ok(bars)
}

// ===========================================================================
// Registration
// ===========================================================================

/// Register all bar-sampling functions and types in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<Bar>()?;

    // Standard bars
    m.add_function(wrap_pyfunction!(tick_bars, m)?)?;
    m.add_function(wrap_pyfunction!(volume_bars, m)?)?;
    m.add_function(wrap_pyfunction!(dollar_bars, m)?)?;

    // Information-driven: imbalance bars
    m.add_function(wrap_pyfunction!(tick_imbalance_bars, m)?)?;
    m.add_function(wrap_pyfunction!(volume_imbalance_bars, m)?)?;

    // Information-driven: run bars
    m.add_function(wrap_pyfunction!(tick_run_bars, m)?)?;
    m.add_function(wrap_pyfunction!(volume_run_bars, m)?)?;

    Ok(())
}

// ===========================================================================
// Tests
// ===========================================================================

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

    // -- Helpers for test data --

    fn make_timestamps(n: usize) -> Vec<f64> {
        (0..n).map(|i| 1700000000.0 + i as f64).collect()
    }

    fn make_prices_up(n: usize) -> Vec<f64> {
        (0..n).map(|i| 100.0 + i as f64 * 0.01).collect()
    }

    fn make_volumes(n: usize, v: f64) -> Vec<f64> {
        vec![v; n]
    }

    fn make_alternating_prices(n: usize) -> Vec<f64> {
        (0..n)
            .map(|i| if i % 2 == 0 { 100.0 } else { 101.0 })
            .collect()
    }

    // -- Validation --

    #[test]
    fn test_validate_length_mismatch() {
        let result = tick_bars(vec![1.0], vec![1.0, 2.0], vec![1.0], 1);
        assert!(result.is_err());
    }

    #[test]
    fn test_empty_input() {
        let bars = tick_bars(vec![], vec![], vec![], 10).unwrap();
        assert!(bars.is_empty());
    }

    #[test]
    fn test_zero_threshold() {
        let bars = tick_bars(vec![1.0], vec![100.0], vec![10.0], 0).unwrap();
        assert!(bars.is_empty());
    }

    // -- Bar __repr__ --

    #[test]
    fn test_bar_repr() {
        let bar = Bar {
            timestamp: 1700000000.0,
            open: 100.0,
            high: 101.0,
            low: 99.0,
            close: 100.5,
            volume: 1000.0,
            vwap: 100.25,
            n_ticks: 50,
        };
        let r = bar.__repr__();
        assert!(r.contains("Bar("));
        assert!(r.contains("ticks=50"));
    }

    // -- Tick Bars --

    #[test]
    fn test_tick_bars_basic() {
        let n = 100;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = make_volumes(n, 10.0);
        let bars = tick_bars(ts, px, vol, 10).unwrap();
        assert_eq!(bars.len(), 10);
        for bar in &bars {
            assert_eq!(bar.n_ticks, 10);
            assert_eq!(bar.volume, 100.0); // 10 * 10.0
            assert!(bar.high >= bar.low);
            assert!(bar.high >= bar.open);
            assert!(bar.high >= bar.close);
            assert!(bar.low <= bar.open);
            assert!(bar.low <= bar.close);
        }
    }

    #[test]
    fn test_tick_bars_partial_last_bar() {
        let n = 15;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = make_volumes(n, 1.0);
        let bars = tick_bars(ts, px, vol, 10).unwrap();
        assert_eq!(bars.len(), 2);
        assert_eq!(bars[0].n_ticks, 10);
        assert_eq!(bars[1].n_ticks, 5);
    }

    #[test]
    fn test_tick_bars_single_tick() {
        let bars = tick_bars(vec![1.0], vec![100.0], vec![5.0], 1).unwrap();
        assert_eq!(bars.len(), 1);
        assert_eq!(bars[0].open, 100.0);
        assert_eq!(bars[0].close, 100.0);
        assert_eq!(bars[0].high, 100.0);
        assert_eq!(bars[0].low, 100.0);
        assert_eq!(bars[0].volume, 5.0);
        assert_eq!(bars[0].vwap, 100.0);
        assert_eq!(bars[0].n_ticks, 1);
    }

    #[test]
    fn test_tick_bars_vwap() {
        // 2 ticks: price=100 vol=10, price=200 vol=30
        // VWAP = (100*10 + 200*30) / (10+30) = 7000/40 = 175.0
        let bars = tick_bars(
            vec![1.0, 2.0],
            vec![100.0, 200.0],
            vec![10.0, 30.0],
            2,
        )
        .unwrap();
        assert_eq!(bars.len(), 1);
        assert!((bars[0].vwap - 175.0).abs() < 1e-10);
    }

    // -- Volume Bars --

    #[test]
    fn test_volume_bars_basic() {
        let n = 100;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = make_volumes(n, 10.0); // total = 1000
        let bars = volume_bars(ts, px, vol, 100.0).unwrap(); // 100/10 = 10 ticks per bar
        assert_eq!(bars.len(), 10);
        for bar in &bars {
            assert!(bar.volume >= 100.0);
        }
    }

    #[test]
    fn test_volume_bars_zero_threshold() {
        let bars = volume_bars(vec![1.0], vec![100.0], vec![10.0], 0.0).unwrap();
        assert!(bars.is_empty());
    }

    #[test]
    fn test_volume_bars_negative_threshold() {
        let bars = volume_bars(vec![1.0], vec![100.0], vec![10.0], -5.0).unwrap();
        assert!(bars.is_empty());
    }

    #[test]
    fn test_volume_bars_large_threshold() {
        // Threshold larger than total volume => one partial bar
        let bars = volume_bars(
            vec![1.0, 2.0],
            vec![100.0, 101.0],
            vec![5.0, 5.0],
            1000.0,
        )
        .unwrap();
        assert_eq!(bars.len(), 1); // partial bar emitted
        assert_eq!(bars[0].n_ticks, 2);
    }

    // -- Dollar Bars --

    #[test]
    fn test_dollar_bars_basic() {
        // price=100, vol=10 => dollar=1000 per tick
        let n = 50;
        let ts = make_timestamps(n);
        let px = vec![100.0; n];
        let vol = make_volumes(n, 10.0);
        let bars = dollar_bars(ts, px, vol, 5000.0).unwrap(); // 5 ticks per bar
        assert_eq!(bars.len(), 10);
    }

    #[test]
    fn test_dollar_bars_zero_volume() {
        // Zero volume ticks should not trigger bar
        let bars = dollar_bars(
            vec![1.0, 2.0, 3.0],
            vec![100.0, 100.0, 100.0],
            vec![0.0, 0.0, 0.0],
            1000.0,
        )
        .unwrap();
        // No dollar threshold reached; one partial bar
        assert_eq!(bars.len(), 1);
        assert_eq!(bars[0].n_ticks, 3);
        assert_eq!(bars[0].volume, 0.0);
    }

    // -- Tick Imbalance Bars --

    #[test]
    fn test_tick_imbalance_bars_all_up() {
        // Monotonically increasing: all b_t = 1, theta grows linearly
        let n = 100;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = make_volumes(n, 1.0);
        let bars = tick_imbalance_bars(ts, px, vol, 10.0).unwrap();
        assert!(!bars.is_empty());
        // First bar should have ~10 ticks
        assert!(bars[0].n_ticks >= 10);
    }

    #[test]
    fn test_tick_imbalance_bars_alternating() {
        // Alternating prices: b = +1, -1, +1, -1 ...
        // theta oscillates near 0, may take a long time to trigger
        let n = 1000;
        let ts = make_timestamps(n);
        let px = make_alternating_prices(n);
        let vol = make_volumes(n, 1.0);
        let bars = tick_imbalance_bars(ts, px, vol, 5.0).unwrap();
        // With alternating ticks, imbalance stays low so fewer bars expected
        assert!(!bars.is_empty());
    }

    #[test]
    fn test_tick_imbalance_bars_empty() {
        let bars = tick_imbalance_bars(vec![], vec![], vec![], 10.0).unwrap();
        assert!(bars.is_empty());
    }

    #[test]
    fn test_tick_imbalance_bars_constant_price() {
        // Constant prices: b carries forward (b_0=1), so theta grows by 1 each tick
        let n = 50;
        let ts = make_timestamps(n);
        let px = vec![100.0; n];
        let vol = make_volumes(n, 1.0);
        let bars = tick_imbalance_bars(ts, px, vol, 10.0).unwrap();
        assert!(!bars.is_empty());
        // theta = +1 each tick (carry forward), so first bar at ~10 ticks
        assert!(bars[0].n_ticks >= 10);
    }

    // -- Volume Imbalance Bars --

    #[test]
    fn test_volume_imbalance_bars_all_up() {
        let n = 100;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = make_volumes(n, 10.0);
        // theta = sum(b_i * v_i) = 10.0 per tick (all up)
        let bars = volume_imbalance_bars(ts, px, vol, 50.0).unwrap();
        assert!(!bars.is_empty());
        // First bar: 5 ticks * 10 vol * 1.0 direction = 50.0
        assert!(bars[0].n_ticks >= 5);
    }

    #[test]
    fn test_volume_imbalance_bars_zero_estimate() {
        let bars = volume_imbalance_bars(vec![1.0], vec![100.0], vec![10.0], 0.0).unwrap();
        assert!(bars.is_empty());
    }

    // -- Tick Run Bars --

    #[test]
    fn test_tick_run_bars_all_up() {
        // All up: buy_run grows every tick
        let n = 100;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = make_volumes(n, 1.0);
        let bars = tick_run_bars(ts, px, vol, 10.0).unwrap();
        assert!(!bars.is_empty());
        assert!(bars[0].n_ticks >= 10);
    }

    #[test]
    fn test_tick_run_bars_alternating() {
        // Alternating: buy/sell runs each grow by 1 every 2 ticks
        // max(buy, sell) increments by 1 every 2 ticks
        let n = 200;
        let ts = make_timestamps(n);
        let px = make_alternating_prices(n);
        let vol = make_volumes(n, 1.0);
        let bars = tick_run_bars(ts, px, vol, 10.0).unwrap();
        assert!(!bars.is_empty());
    }

    #[test]
    fn test_tick_run_bars_empty() {
        let bars = tick_run_bars(vec![], vec![], vec![], 10.0).unwrap();
        assert!(bars.is_empty());
    }

    // -- Volume Run Bars --

    #[test]
    fn test_volume_run_bars_all_up() {
        let n = 100;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = make_volumes(n, 5.0);
        // buy_vol_run grows by 5.0 per tick
        let bars = volume_run_bars(ts, px, vol, 25.0).unwrap();
        assert!(!bars.is_empty());
        // ~5 ticks per bar
        assert!(bars[0].n_ticks >= 5);
    }

    #[test]
    fn test_volume_run_bars_mixed() {
        // prices: up, up, down, up, up, down, ...
        let prices = vec![100.0, 101.0, 102.0, 101.0, 102.0, 103.0, 102.0, 103.0, 104.0, 103.0];
        let n = prices.len();
        let ts = make_timestamps(n);
        let vol = make_volumes(n, 10.0);
        let bars = volume_run_bars(ts, prices, vol, 20.0).unwrap();
        assert!(!bars.is_empty());
        for bar in &bars {
            assert!(bar.high >= bar.low);
            assert!(bar.volume > 0.0);
        }
    }

    #[test]
    fn test_volume_run_bars_zero_volume() {
        // Zero volume shouldn't accumulate any run volume
        let n = 10;
        let ts = make_timestamps(n);
        let px = make_prices_up(n);
        let vol = vec![0.0; n];
        let bars = volume_run_bars(ts, px, vol, 10.0).unwrap();
        // No volume accumulated => single partial bar
        assert_eq!(bars.len(), 1);
        assert_eq!(bars[0].n_ticks, n as u64);
    }

    // -- OHLCV correctness --

    #[test]
    fn test_ohlcv_correctness() {
        let ts = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let px = vec![100.0, 105.0, 98.0, 102.0, 101.0];
        let vol = vec![10.0, 20.0, 15.0, 25.0, 30.0];
        let bars = tick_bars(ts, px, vol, 5).unwrap();
        assert_eq!(bars.len(), 1);
        let bar = &bars[0];
        assert_eq!(bar.open, 100.0);
        assert_eq!(bar.close, 101.0);
        assert_eq!(bar.high, 105.0);
        assert_eq!(bar.low, 98.0);
        assert_eq!(bar.volume, 100.0);
        assert_eq!(bar.n_ticks, 5);
        assert_eq!(bar.timestamp, 1.0);
        // VWAP = (100*10 + 105*20 + 98*15 + 102*25 + 101*30) / 100
        // = (1000 + 2100 + 1470 + 2550 + 3030) / 100 = 10150/100 = 101.5
        assert!((bar.vwap - 101.5).abs() < 1e-10);
    }

    // -- EWM helper --

    #[test]
    fn test_ewm_update_basic() {
        let v = ewm_update(10.0, 20.0, 10.0);
        // alpha = 2/(10+1) = 2/11 ≈ 0.1818
        // result = 0.1818*20 + 0.8182*10 = 3.636 + 8.182 = 11.818
        let expected = 2.0 / 11.0 * 20.0 + 9.0 / 11.0 * 10.0;
        assert!((v - expected).abs() < 1e-10);
    }

    #[test]
    fn test_ewm_update_zero_span() {
        assert_eq!(ewm_update(10.0, 42.0, 0.0), 42.0);
    }

    // -- Tick direction --

    #[test]
    fn test_tick_direction_up() {
        assert_eq!(tick_direction(101.0, 100.0, -1.0), 1.0);
    }

    #[test]
    fn test_tick_direction_down() {
        assert_eq!(tick_direction(99.0, 100.0, 1.0), -1.0);
    }

    #[test]
    fn test_tick_direction_unchanged() {
        assert_eq!(tick_direction(100.0, 100.0, 1.0), 1.0);
        assert_eq!(tick_direction(100.0, 100.0, -1.0), -1.0);
    }

    // -- VWAP edge case: zero volume --

    #[test]
    fn test_vwap_zero_volume() {
        let bars = tick_bars(
            vec![1.0, 2.0],
            vec![100.0, 200.0],
            vec![0.0, 0.0],
            2,
        )
        .unwrap();
        assert_eq!(bars.len(), 1);
        // With zero total volume, VWAP falls back to close price
        assert_eq!(bars[0].vwap, bars[0].close);
    }
}
