//! Triple Barrier Labeling and CUSUM event sampling.
//!
//! Implements Chapters 2-3 of Marcos Lopez de Prado's
//! "Advances in Financial Machine Learning":
//! - Exponentially weighted daily volatility (get_daily_vol)
//! - Symmetric CUSUM filter for event-driven sampling (cusum_filter)
//! - Triple barrier labeling (triple_barrier_labels)
//! - Meta-labeling for bet sizing (meta_labels)
//! - Rare label dropping (drop_labels)
//!
//! All functions are pure Rust + PyO3. No external crate dependencies.

use pyo3::prelude::*;

// ===========================================================================
// BarrierLabel
// ===========================================================================

/// Result of triple barrier labeling for a single event.
///
/// Fields:
///   event_idx: index in the original price series where the event started
///   label: -1 (stop-loss), 0 (vertical barrier / below min_ret), +1 (profit-taking)
///   ret: realized return from entry to barrier touch
///   barrier: which barrier was touched ("pt", "sl", or "vb")
///   touch_idx: index in the price series where the barrier was touched
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct BarrierLabel {
    #[pyo3(get)]
    pub event_idx: usize,
    #[pyo3(get)]
    pub label: i32,
    #[pyo3(get)]
    pub ret: f64,
    #[pyo3(get)]
    pub barrier: String,
    #[pyo3(get)]
    pub touch_idx: usize,
}

#[pymethods]
impl BarrierLabel {
    fn __repr__(&self) -> String {
        format!(
            "BarrierLabel(event_idx={}, label={}, ret={:.6}, barrier=\"{}\", touch_idx={})",
            self.event_idx, self.label, self.ret, self.barrier, self.touch_idx
        )
    }
}

// ===========================================================================
// Private helpers
// ===========================================================================

/// Compute log returns from a price series.
/// Returns a Vec of length `prices.len() - 1`.
/// ret[i] = ln(prices[i+1] / prices[i]).
#[inline]
fn log_returns(prices: &[f64]) -> Vec<f64> {
    let n = prices.len();
    if n < 2 {
        return Vec::new();
    }
    let mut rets = Vec::with_capacity(n - 1);
    for i in 0..n - 1 {
        if prices[i] > 0.0 && prices[i + 1] > 0.0 {
            rets.push((prices[i + 1] / prices[i]).ln());
        } else {
            rets.push(0.0);
        }
    }
    rets
}

/// Exponentially weighted moving standard deviation.
///
/// Uses the EWM alpha = 2 / (span + 1) convention.
/// For the first `span` values, uses an expanding window to avoid NaN.
/// Returns a Vec of length `values.len()` where each element is the
/// EWM standard deviation up to and including that index.
fn ewm_std(values: &[f64], span: usize) -> Vec<f64> {
    let n = values.len();
    if n == 0 {
        return Vec::new();
    }
    let alpha = 2.0 / (span as f64 + 1.0);
    let mut result = vec![0.0; n];

    // We track the EWM mean and EWM variance (biased).
    // ewm_var = alpha * (x - ewm_mean_prev)^2 + (1 - alpha) * ewm_var_prev
    let mut ewm_mean = values[0];
    let mut ewm_var = 0.0;

    // For the expanding window phase, also track simple stats
    let mut sum = values[0];
    let mut sum_sq = values[0] * values[0];

    result[0] = 0.0; // single point has no std

    for i in 1..n {
        let x = values[i];
        let count = i + 1;

        if count <= span {
            // Expanding window: use simple std for stability
            sum += x;
            sum_sq += x * x;
            let mean = sum / count as f64;
            // Sample variance
            let var = if count > 1 {
                (sum_sq - count as f64 * mean * mean) / (count as f64 - 1.0)
            } else {
                0.0
            };
            result[i] = if var > 0.0 { var.sqrt() } else { 0.0 };
        } else {
            // EWM phase
            let diff = x - ewm_mean;
            ewm_mean += alpha * diff;
            ewm_var = (1.0 - alpha) * (ewm_var + alpha * diff * diff);
            result[i] = if ewm_var > 0.0 { ewm_var.sqrt() } else { 0.0 };
        }
    }

    // Transition: seed EWM state from expanding window at the boundary.
    // Recompute the EWM from scratch so the transition is smooth.
    if n > span && span > 0 {
        ewm_mean = values[0];
        ewm_var = 0.0;
        for i in 1..n {
            let diff = values[i] - ewm_mean;
            ewm_mean += alpha * diff;
            ewm_var = (1.0 - alpha) * (ewm_var + alpha * diff * diff);
            if i >= span {
                result[i] = if ewm_var > 0.0 { ewm_var.sqrt() } else { 0.0 };
            }
        }
    }

    result
}

// ===========================================================================
// get_daily_vol
// ===========================================================================

/// Compute exponentially weighted daily volatility from a price series.
///
/// Computes log returns, then applies EWM standard deviation with the given span.
/// First `span` entries use an expanding window for stability.
///
/// Returns a Vec of length `prices.len()`. The first element is always 0.0
/// (no return can be computed from a single price). The i-th element is the
/// EWM std of log returns up to time i.
///
/// Args:
///     prices: raw price series (must have at least 2 elements)
///     span: lookback for EWM (e.g., 20 for ~20-day half-life)
#[pyfunction]
pub fn get_daily_vol(prices: Vec<f64>, span: usize) -> PyResult<Vec<f64>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    if prices.len() < 2 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "prices must have at least 2 elements",
        ));
    }
    if span == 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "span must be > 0",
        ));
    }

    let rets = log_returns(&prices);
    let vol = ewm_std(&rets, span);

    // Align output to prices: prepend 0.0 for the first price (no return available)
    let mut result = Vec::with_capacity(prices.len());
    result.push(0.0);
    result.extend_from_slice(&vol);
    Ok(result)
}

// ===========================================================================
// cusum_filter
// ===========================================================================

/// Symmetric CUSUM filter for event-driven sampling (Chapter 2.5.2.1).
///
/// Tracks cumulative positive and negative sums of price changes.
/// When either sum exceeds the threshold, records the index and resets both.
///
/// This produces a structurally meaningful subsample of the time series,
/// filtering out noise and focusing on significant moves.
///
/// Args:
///     prices: raw price series
///     threshold: CUSUM trigger threshold (in price units)
///
/// Returns:
///     indices where structural breaks occurred
#[pyfunction]
pub fn cusum_filter(prices: Vec<f64>, threshold: f64) -> PyResult<Vec<usize>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    if prices.len() < 2 {
        return Ok(Vec::new());
    }
    if threshold <= 0.0 || !threshold.is_finite() {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "threshold must be a positive finite number",
        ));
    }

    let mut events = Vec::new();
    let mut s_pos = 0.0_f64;
    let mut s_neg = 0.0_f64;

    for i in 1..prices.len() {
        let diff = prices[i] - prices[i - 1];
        if !diff.is_finite() {
            continue;
        }

        // Cumulative sums: positive deviations and negative deviations
        s_pos = (s_pos + diff).max(0.0);
        s_neg = (s_neg - diff).max(0.0); // note: subtracting diff to track negative moves

        if s_pos >= threshold || s_neg >= threshold {
            events.push(i);
            s_pos = 0.0;
            s_neg = 0.0;
        }
    }

    Ok(events)
}

// ===========================================================================
// triple_barrier_labels
// ===========================================================================

/// Triple barrier labeling (Chapter 3.2-3.4).
///
/// For each event index, sets three barriers:
/// - Upper (profit-taking): entry_price * (1 + daily_vol * pt_sl[0])
/// - Lower (stop-loss):     entry_price * (1 - daily_vol * pt_sl[1])
/// - Vertical:              max_holding bars forward from entry
///
/// Scans forward from each event to find which barrier is touched first.
///
/// Args:
///     prices: raw price series
///     timestamps: timestamp for each price (same length as prices; used for alignment)
///     events: indices from cusum_filter or user-provided event indices
///     pt_sl: [profit_taking_multiplier, stop_loss_multiplier]; 0.0 = disabled
///     min_ret: minimum absolute return to assign +1/-1 (below this, label = 0)
///     max_holding: vertical barrier in bars (0 = no vertical barrier)
///     vol_span: lookback span for daily vol estimation
///
/// Returns:
///     Vec of BarrierLabel, one per event
#[pyfunction]
pub fn triple_barrier_labels(
    prices: Vec<f64>,
    timestamps: Vec<f64>,
    events: Vec<usize>,
    pt_sl: [f64; 2],
    min_ret: f64,
    max_holding: usize,
    vol_span: usize,
) -> PyResult<Vec<BarrierLabel>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    let n = prices.len();

    // --- Input validation ---
    if n < 2 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "prices must have at least 2 elements",
        ));
    }
    if timestamps.len() != n {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
            "timestamps length ({}) must match prices length ({})",
            timestamps.len(),
            n
        )));
    }
    if vol_span == 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "vol_span must be > 0",
        ));
    }

    // Validate events are within bounds
    for &idx in &events {
        if idx >= n {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                "event index {} is out of bounds (prices length {})",
                idx, n
            )));
        }
    }

    // Precompute daily vol for the entire series
    let rets = log_returns(&prices);
    let vol_series = ewm_std(&rets, vol_span);
    // vol_series has length n-1 (aligned to returns).
    // vol_series[i] corresponds to the vol estimate at prices[i+1].
    // To get vol at prices[j], use vol_series[j-1] (when j >= 1).

    let pt_mult = pt_sl[0];
    let sl_mult = pt_sl[1];

    let mut labels = Vec::with_capacity(events.len());

    for &event_idx in &events {
        let entry_price = prices[event_idx];
        if entry_price <= 0.0 || !entry_price.is_finite() {
            // Skip invalid entry prices
            labels.push(BarrierLabel {
                event_idx,
                label: 0,
                ret: 0.0,
                barrier: "vb".to_string(),
                touch_idx: event_idx,
            });
            continue;
        }

        // Get daily vol at event time
        let daily_vol = if event_idx == 0 {
            // No return data before the first price; use a small default
            if !vol_series.is_empty() {
                vol_series[0]
            } else {
                0.0
            }
        } else {
            // vol_series[i] is the vol of returns up to index i (return i = prices[i+1]/prices[i])
            // For event at prices[event_idx], the latest return is at event_idx-1
            vol_series[(event_idx - 1).min(vol_series.len() - 1)]
        };

        // Set barrier levels (in price space)
        let pt_barrier = if pt_mult > 0.0 && daily_vol > 0.0 {
            Some(entry_price * (1.0 + daily_vol * pt_mult))
        } else {
            None
        };
        let sl_barrier = if sl_mult > 0.0 && daily_vol > 0.0 {
            Some(entry_price * (1.0 - daily_vol * sl_mult))
        } else {
            None
        };

        // Vertical barrier: index of the last bar to consider
        let vb_idx = if max_holding > 0 {
            (event_idx + max_holding).min(n - 1)
        } else {
            n - 1 // no vertical barrier: scan to end of series
        };

        // Scan forward for barrier touches
        let mut touch_idx = vb_idx;
        let mut barrier_type = "vb";
        let mut first_touch = false;

        for j in (event_idx + 1)..=vb_idx {
            let p = prices[j];
            if !p.is_finite() {
                continue;
            }

            // Check profit-taking (upper) barrier
            if let Some(pt) = pt_barrier {
                if p >= pt {
                    touch_idx = j;
                    barrier_type = "pt";
                    first_touch = true;
                    break;
                }
            }

            // Check stop-loss (lower) barrier
            if let Some(sl) = sl_barrier {
                if p <= sl {
                    touch_idx = j;
                    barrier_type = "sl";
                    first_touch = true;
                    break;
                }
            }
        }

        // Compute return
        let exit_price = prices[touch_idx];
        let ret = if entry_price > 0.0 && exit_price > 0.0 {
            (exit_price / entry_price).ln()
        } else {
            0.0
        };

        // Assign label
        let label = if !first_touch {
            // Vertical barrier hit (or end of series)
            if ret.abs() < min_ret {
                0
            } else if ret > 0.0 {
                1
            } else {
                -1
            }
        } else {
            match barrier_type {
                "pt" => {
                    if ret.abs() < min_ret {
                        0
                    } else {
                        1
                    }
                }
                "sl" => {
                    if ret.abs() < min_ret {
                        0
                    } else {
                        -1
                    }
                }
                _ => 0,
            }
        };

        labels.push(BarrierLabel {
            event_idx,
            label,
            ret,
            barrier: barrier_type.to_string(),
            touch_idx,
        });
    }

    Ok(labels)
}

// ===========================================================================
// meta_labels
// ===========================================================================

/// Meta-labeling (Chapter 3.6).
///
/// Given a primary model's directional signals, meta-labeling determines
/// whether acting on each signal is profitable (label=1) or not (label=0).
///
/// The primary model provides the side (+1 long, -1 short). The meta-label
/// indicates if the bet was correct:
/// - Primary says +1 (long) and price goes up: meta-label = 1 (correct)
/// - Primary says +1 (long) and price goes down: meta-label = 0 (wrong)
/// - Primary says -1 (short) and price goes down: meta-label = 1 (correct)
/// - Primary says -1 (short) and price goes up: meta-label = 0 (wrong)
///
/// Args:
///     prices: raw price series
///     timestamps: timestamp for each price (same length as prices)
///     primary_signals: Vec of (event_idx, side) where side is +1 or -1
///     pt_sl: [profit_taking_multiplier, stop_loss_multiplier]; 0.0 = disabled
///     max_holding: vertical barrier in bars (0 = no vertical barrier)
///     vol_span: lookback span for daily vol estimation
///
/// Returns:
///     Vec of BarrierLabel with label = 0 or 1 (binary classification)
#[pyfunction]
pub fn meta_labels(
    prices: Vec<f64>,
    timestamps: Vec<f64>,
    primary_signals: Vec<(usize, i32)>,
    pt_sl: [f64; 2],
    max_holding: usize,
    vol_span: usize,
) -> PyResult<Vec<BarrierLabel>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    let n = prices.len();

    // --- Input validation ---
    if n < 2 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "prices must have at least 2 elements",
        ));
    }
    if timestamps.len() != n {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
            "timestamps length ({}) must match prices length ({})",
            timestamps.len(),
            n
        )));
    }
    if vol_span == 0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "vol_span must be > 0",
        ));
    }

    // Validate signals
    for &(idx, side) in &primary_signals {
        if idx >= n {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                "signal index {} is out of bounds (prices length {})",
                idx, n
            )));
        }
        if side != 1 && side != -1 {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                "signal side must be +1 or -1, got {}",
                side
            )));
        }
    }

    // Precompute daily vol
    let rets = log_returns(&prices);
    let vol_series = ewm_std(&rets, vol_span);

    let pt_mult = pt_sl[0];
    let sl_mult = pt_sl[1];

    let mut labels = Vec::with_capacity(primary_signals.len());

    for &(event_idx, side) in &primary_signals {
        let entry_price = prices[event_idx];
        if entry_price <= 0.0 || !entry_price.is_finite() {
            labels.push(BarrierLabel {
                event_idx,
                label: 0,
                ret: 0.0,
                barrier: "vb".to_string(),
                touch_idx: event_idx,
            });
            continue;
        }

        // Get daily vol at event time
        let daily_vol = if event_idx == 0 {
            if !vol_series.is_empty() {
                vol_series[0]
            } else {
                0.0
            }
        } else {
            vol_series[(event_idx - 1).min(vol_series.len() - 1)]
        };

        // For meta-labeling, the barriers are adjusted by the side of the primary signal.
        // If side = +1 (long): upper = profit-taking (at +pt_mult*vol), lower = stop-loss (at -sl_mult*vol).
        // If side = -1 (short): upper = stop-loss (at +sl_mult*vol), lower = profit-taking (at -pt_mult*vol).
        // We multiply by side_f to flip the barrier directions for shorts.
        let side_f = side as f64;

        // Barrier levels in price space
        // For long: upper = entry*(1 + vol*pt_mult), lower = entry*(1 - vol*sl_mult)
        // For short: upper = entry*(1 + vol*sl_mult), lower = entry*(1 - vol*pt_mult)
        let (upper_mult, lower_mult) = if side == 1 {
            (pt_mult, sl_mult)
        } else {
            (sl_mult, pt_mult)
        };

        let upper_barrier = if upper_mult > 0.0 && daily_vol > 0.0 {
            Some(entry_price * (1.0 + daily_vol * upper_mult))
        } else {
            None
        };
        let lower_barrier = if lower_mult > 0.0 && daily_vol > 0.0 {
            Some(entry_price * (1.0 - daily_vol * lower_mult))
        } else {
            None
        };

        // Vertical barrier
        let vb_idx = if max_holding > 0 {
            (event_idx + max_holding).min(n - 1)
        } else {
            n - 1
        };

        // Scan forward
        let mut touch_idx = vb_idx;
        let mut barrier_type = "vb";

        for j in (event_idx + 1)..=vb_idx {
            let p = prices[j];
            if !p.is_finite() {
                continue;
            }

            // Check upper barrier
            if let Some(ub) = upper_barrier {
                if p >= ub {
                    touch_idx = j;
                    barrier_type = if side == 1 { "pt" } else { "sl" };
                    break;
                }
            }

            // Check lower barrier
            if let Some(lb) = lower_barrier {
                if p <= lb {
                    touch_idx = j;
                    barrier_type = if side == 1 { "sl" } else { "pt" };
                    break;
                }
            }
        }

        // Compute directional return (adjusted for side)
        let exit_price = prices[touch_idx];
        let raw_ret = if entry_price > 0.0 && exit_price > 0.0 {
            (exit_price / entry_price).ln()
        } else {
            0.0
        };
        let directional_ret = raw_ret * side_f;

        // Meta-label: 1 if the primary signal was correct, 0 otherwise
        let label = if directional_ret > 0.0 { 1 } else { 0 };

        labels.push(BarrierLabel {
            event_idx,
            label,
            ret: raw_ret,
            barrier: barrier_type.to_string(),
            touch_idx,
        });
    }

    Ok(labels)
}

// ===========================================================================
// drop_labels
// ===========================================================================

/// Drop rare labels that don't meet a minimum percentage threshold (Chapter 3.7).
///
/// Counts occurrences of each label class (-1, 0, +1). Removes events whose
/// label class has fewer than `min_pct` fraction of the total.
///
/// This prevents classifier training on heavily imbalanced datasets where
/// rare classes can degrade model performance.
///
/// Args:
///     labels: Vec of BarrierLabel from triple_barrier_labels or meta_labels
///     min_pct: minimum fraction of total (e.g., 0.05 = 5%). Must be in [0, 1].
///
/// Returns:
///     Filtered Vec of BarrierLabel
#[pyfunction]
pub fn drop_labels(labels: Vec<BarrierLabel>, min_pct: f64) -> PyResult<Vec<BarrierLabel>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    if !min_pct.is_finite() || min_pct < 0.0 || min_pct > 1.0 {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "min_pct must be in [0.0, 1.0]",
        ));
    }

    if labels.is_empty() {
        return Ok(Vec::new());
    }

    let total = labels.len() as f64;
    let min_count = (total * min_pct).ceil() as usize;

    // Count occurrences of each label
    let mut count_neg = 0usize;
    let mut count_zero = 0usize;
    let mut count_pos = 0usize;

    for label in &labels {
        match label.label {
            -1 => count_neg += 1,
            0 => count_zero += 1,
            1 => count_pos += 1,
            _ => {} // ignore unexpected labels
        }
    }

    // Determine which classes to keep
    let keep_neg = count_neg >= min_count;
    let keep_zero = count_zero >= min_count;
    let keep_pos = count_pos >= min_count;

    let result: Vec<BarrierLabel> = labels
        .into_iter()
        .filter(|l| match l.label {
            -1 => keep_neg,
            0 => keep_zero,
            1 => keep_pos,
            _ => false,
        })
        .collect();

    Ok(result)
}

// ===========================================================================
// Module registration
// ===========================================================================

/// Register all labeling classes and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<BarrierLabel>()?;
    m.add_function(wrap_pyfunction!(get_daily_vol, m)?)?;
    m.add_function(wrap_pyfunction!(cusum_filter, m)?)?;
    m.add_function(wrap_pyfunction!(triple_barrier_labels, m)?)?;
    m.add_function(wrap_pyfunction!(meta_labels, m)?)?;
    m.add_function(wrap_pyfunction!(drop_labels, m)?)?;
    Ok(())
}

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

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

    // ---- log_returns ----

    #[test]
    fn log_returns_basic() {
        let prices = vec![100.0, 110.0, 105.0];
        let rets = log_returns(&prices);
        assert_eq!(rets.len(), 2);
        assert!((rets[0] - (110.0_f64 / 100.0).ln()).abs() < 1e-12);
        assert!((rets[1] - (105.0_f64 / 110.0).ln()).abs() < 1e-12);
    }

    #[test]
    fn log_returns_single_price() {
        let rets = log_returns(&[100.0]);
        assert!(rets.is_empty());
    }

    #[test]
    fn log_returns_zero_price_handled() {
        let rets = log_returns(&[100.0, 0.0, 50.0]);
        assert_eq!(rets.len(), 2);
        assert_eq!(rets[0], 0.0); // 0.0 price -> 0.0 return
        assert_eq!(rets[1], 0.0); // from 0.0 -> also 0.0
    }

    // ---- ewm_std ----

    #[test]
    fn ewm_std_constant_series() {
        let vals = vec![1.0; 50];
        let std = ewm_std(&vals, 10);
        assert_eq!(std.len(), 50);
        for v in &std {
            assert!(*v < 1e-12, "constant series should have zero vol");
        }
    }

    #[test]
    fn ewm_std_increasing() {
        // Alternating up/down should produce nonzero vol
        let vals: Vec<f64> = (0..100).map(|i| if i % 2 == 0 { 1.0 } else { -1.0 }).collect();
        let std = ewm_std(&vals, 20);
        assert!(std[99] > 0.0);
    }

    #[test]
    fn ewm_std_empty() {
        let std = ewm_std(&[], 10);
        assert!(std.is_empty());
    }

    // ---- cusum_filter ----

    #[test]
    fn cusum_filter_no_events() {
        // Tiny moves, high threshold
        let prices: Vec<f64> = (0..100).map(|i| 100.0 + 0.001 * i as f64).collect();
        let events = cusum_filter(prices, 10.0).unwrap();
        assert!(events.is_empty());
    }

    #[test]
    fn cusum_filter_detects_jump() {
        let mut prices = vec![100.0; 10];
        prices.push(110.0); // +10 jump
        prices.extend(vec![110.0; 10]);
        let events = cusum_filter(prices, 5.0).unwrap();
        assert!(!events.is_empty());
        assert_eq!(events[0], 10); // jump happens at index 10
    }

    #[test]
    fn cusum_filter_detects_drop() {
        let mut prices = vec![100.0; 10];
        prices.push(90.0); // -10 drop
        prices.extend(vec![90.0; 10]);
        let events = cusum_filter(prices, 5.0).unwrap();
        assert!(!events.is_empty());
        assert_eq!(events[0], 10);
    }

    #[test]
    fn cusum_filter_resets() {
        // Two jumps separated by flat
        let mut prices = vec![100.0; 5];
        prices.push(120.0); // +20
        prices.extend(vec![120.0; 10]);
        prices.push(140.0); // +20
        prices.extend(vec![140.0; 5]);
        let events = cusum_filter(prices, 10.0).unwrap();
        assert_eq!(events.len(), 2);
    }

    #[test]
    fn cusum_filter_invalid_threshold() {
        assert!(cusum_filter(vec![1.0, 2.0], 0.0).is_err());
        assert!(cusum_filter(vec![1.0, 2.0], -1.0).is_err());
        assert!(cusum_filter(vec![1.0, 2.0], f64::NAN).is_err());
    }

    // ---- triple_barrier_labels ----

    #[test]
    fn triple_barrier_profit_taking() {
        // Price goes up sharply -> should hit PT barrier
        let n = 100;
        let mut prices = Vec::with_capacity(n);
        for i in 0..n {
            prices.push(100.0 + i as f64 * 0.5); // steady uptrend
        }
        let timestamps: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let events = vec![10]; // event at index 10

        let labels = triple_barrier_labels(
            prices,
            timestamps,
            events,
            [1.0, 1.0], // PT and SL at 1x vol
            0.0,        // no min_ret
            50,         // max_holding 50 bars
            20,         // vol_span
        )
        .unwrap();

        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].event_idx, 10);
        assert_eq!(labels[0].label, 1);
        assert_eq!(labels[0].barrier, "pt");
        assert!(labels[0].ret > 0.0);
    }

    #[test]
    fn triple_barrier_stop_loss() {
        // Price goes down sharply -> should hit SL barrier
        let n = 100;
        let mut prices = Vec::with_capacity(n);
        for i in 0..n {
            prices.push(100.0 - i as f64 * 0.5); // steady downtrend
        }
        let timestamps: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let events = vec![10];

        let labels = triple_barrier_labels(
            prices,
            timestamps,
            events,
            [1.0, 1.0],
            0.0,
            50,
            20,
        )
        .unwrap();

        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].event_idx, 10);
        assert_eq!(labels[0].label, -1);
        assert_eq!(labels[0].barrier, "sl");
        assert!(labels[0].ret < 0.0);
    }

    #[test]
    fn triple_barrier_vertical() {
        // Flat prices -> no horizontal barriers hit -> vertical barrier
        let n = 100;
        let prices = vec![100.0; n];
        let timestamps: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let events = vec![10];

        let labels = triple_barrier_labels(
            prices,
            timestamps,
            events,
            [1.0, 1.0],
            0.0,
            10,
            5,
        )
        .unwrap();

        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].barrier, "vb");
        assert_eq!(labels[0].label, 0);
    }

    #[test]
    fn triple_barrier_min_ret() {
        // Small move, but below min_ret -> label 0
        let n = 100;
        let mut prices = Vec::with_capacity(n);
        for i in 0..n {
            prices.push(100.0 + i as f64 * 0.001); // tiny uptrend
        }
        let timestamps: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let events = vec![10];

        let labels = triple_barrier_labels(
            prices,
            timestamps,
            events,
            [0.0, 0.0], // no horizontal barriers
            0.5,         // high min_ret
            20,
            5,
        )
        .unwrap();

        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].label, 0); // below min_ret
    }

    #[test]
    fn triple_barrier_out_of_bounds_event() {
        let prices = vec![100.0; 10];
        let timestamps: Vec<f64> = (0..10).map(|i| i as f64).collect();
        let result = triple_barrier_labels(
            prices, timestamps, vec![100], // out of bounds
            [1.0, 1.0], 0.0, 5, 3,
        );
        assert!(result.is_err());
    }

    #[test]
    fn triple_barrier_mismatched_lengths() {
        let result = triple_barrier_labels(
            vec![100.0; 10],
            vec![1.0; 5], // wrong length
            vec![0],
            [1.0, 1.0],
            0.0,
            5,
            3,
        );
        assert!(result.is_err());
    }

    // ---- meta_labels ----

    #[test]
    fn meta_label_correct_long() {
        // Primary says long (+1), price goes up -> meta = 1
        let n = 100;
        let mut prices = Vec::with_capacity(n);
        for i in 0..n {
            prices.push(100.0 + i as f64 * 0.5);
        }
        let timestamps: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let signals = vec![(10, 1)]; // long at index 10

        let labels = meta_labels(
            prices,
            timestamps,
            signals,
            [1.0, 1.0],
            50,
            20,
        )
        .unwrap();

        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].label, 1); // correct prediction
    }

    #[test]
    fn meta_label_wrong_long() {
        // Primary says long (+1), price goes down -> meta = 0
        let n = 100;
        let mut prices = Vec::with_capacity(n);
        for i in 0..n {
            prices.push(100.0 - i as f64 * 0.5);
        }
        let timestamps: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let signals = vec![(10, 1)];

        let labels = meta_labels(
            prices,
            timestamps,
            signals,
            [1.0, 1.0],
            50,
            20,
        )
        .unwrap();

        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].label, 0); // wrong prediction
    }

    #[test]
    fn meta_label_correct_short() {
        // Primary says short (-1), price goes down -> meta = 1
        let n = 100;
        let mut prices = Vec::with_capacity(n);
        for i in 0..n {
            prices.push(100.0 - i as f64 * 0.5);
        }
        let timestamps: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let signals = vec![(10, -1)];

        let labels = meta_labels(
            prices,
            timestamps,
            signals,
            [1.0, 1.0],
            50,
            20,
        )
        .unwrap();

        assert_eq!(labels.len(), 1);
        assert_eq!(labels[0].label, 1); // correct (price went down, we were short)
    }

    #[test]
    fn meta_label_invalid_side() {
        let result = meta_labels(
            vec![100.0; 10],
            (0..10).map(|i| i as f64).collect(),
            vec![(5, 2)], // invalid side
            [1.0, 1.0],
            5,
            3,
        );
        assert!(result.is_err());
    }

    // ---- drop_labels ----

    #[test]
    fn drop_labels_removes_rare() {
        // 90 labels of +1, 10 labels of -1
        let mut labels = Vec::new();
        for i in 0..90 {
            labels.push(BarrierLabel {
                event_idx: i,
                label: 1,
                ret: 0.01,
                barrier: "pt".to_string(),
                touch_idx: i + 1,
            });
        }
        for i in 90..100 {
            labels.push(BarrierLabel {
                event_idx: i,
                label: -1,
                ret: -0.01,
                barrier: "sl".to_string(),
                touch_idx: i + 1,
            });
        }

        // min_pct = 0.15 (15%) -> -1 class has only 10% -> should be dropped
        let filtered = drop_labels(labels, 0.15).unwrap();
        assert_eq!(filtered.len(), 90);
        assert!(filtered.iter().all(|l| l.label == 1));
    }

    #[test]
    fn drop_labels_keeps_all() {
        let mut labels = Vec::new();
        for i in 0..50 {
            labels.push(BarrierLabel {
                event_idx: i,
                label: 1,
                ret: 0.01,
                barrier: "pt".to_string(),
                touch_idx: i + 1,
            });
        }
        for i in 50..100 {
            labels.push(BarrierLabel {
                event_idx: i,
                label: -1,
                ret: -0.01,
                barrier: "sl".to_string(),
                touch_idx: i + 1,
            });
        }

        // Both classes at 50% -> nothing dropped at 0.10 threshold
        let filtered = drop_labels(labels, 0.10).unwrap();
        assert_eq!(filtered.len(), 100);
    }

    #[test]
    fn drop_labels_empty() {
        let filtered = drop_labels(Vec::new(), 0.05).unwrap();
        assert!(filtered.is_empty());
    }

    #[test]
    fn drop_labels_invalid_pct() {
        assert!(drop_labels(Vec::new(), -0.1).is_err());
        assert!(drop_labels(Vec::new(), 1.5).is_err());
        assert!(drop_labels(Vec::new(), f64::NAN).is_err());
    }

    // ---- get_daily_vol ----

    #[test]
    fn get_daily_vol_basic() {
        let prices: Vec<f64> = (1..=100).map(|i| 100.0 + i as f64 * 0.1).collect();
        let vol = get_daily_vol(prices.clone(), 20).unwrap();
        assert_eq!(vol.len(), prices.len());
        assert_eq!(vol[0], 0.0); // first element always 0
        // Vol should be small but nonzero for a steady uptrend
        assert!(vol[50] > 0.0);
    }

    #[test]
    fn get_daily_vol_too_short() {
        assert!(get_daily_vol(vec![100.0], 10).is_err());
    }

    #[test]
    fn get_daily_vol_zero_span() {
        assert!(get_daily_vol(vec![100.0, 101.0], 0).is_err());
    }

    // ---- Integration: cusum -> triple_barrier ----

    #[test]
    fn integration_cusum_to_labels() {
        // Generate a price series with a clear regime change
        let mut prices = Vec::with_capacity(200);
        for i in 0..100 {
            prices.push(100.0 + (i as f64 * 0.01).sin() * 0.5); // oscillating near 100
        }
        for i in 100..200 {
            prices.push(100.0 + (i - 100) as f64 * 0.3); // strong uptrend
        }
        let timestamps: Vec<f64> = (0..200).map(|i| i as f64).collect();

        // CUSUM should detect the regime change
        let events = cusum_filter(prices.clone(), 1.0).unwrap();
        assert!(!events.is_empty(), "CUSUM should detect the structural break");

        // Label the events
        let labels = triple_barrier_labels(
            prices,
            timestamps,
            events,
            [2.0, 2.0],
            0.0,
            20,
            10,
        )
        .unwrap();

        assert!(!labels.is_empty());
        // At least some labels in the uptrend should be +1
        let pos_count = labels.iter().filter(|l| l.label == 1).count();
        assert!(pos_count > 0, "should have some positive labels in uptrend");
    }
}
