//! Lead-Lag Detection for prediction market trading.
//!
//! Implements:
//! - Hayashi-Yoshida asynchronous covariance estimator for non-synchronous data
//! - Cross-correlation at multiple lags
//! - Granger causality F-test with OLS regression
//! - Lead-lag network construction from pairwise analysis
//!
//! All math from scratch in Rust -- no external crate dependencies beyond PyO3.

use pyo3::prelude::*;

// ===========================================================================
// Frozen result types
// ===========================================================================

/// Result of a lead-lag analysis between two series.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct LeadLagResult {
    /// Peak cross-correlation magnitude.
    #[pyo3(get)]
    pub correlation: f64,
    /// Name of the leading series.
    #[pyo3(get)]
    pub leader: String,
    /// Name of the lagging series.
    #[pyo3(get)]
    pub lagger: String,
    /// Lag (in periods) at which peak correlation occurs.
    #[pyo3(get)]
    pub lag_periods: i32,
    /// Approximate p-value for the peak correlation.
    #[pyo3(get)]
    pub p_value: f64,
}

#[pymethods]
impl LeadLagResult {
    fn __repr__(&self) -> String {
        format!(
            "LeadLagResult(leader='{}', lagger='{}', corr={:.4}, lag={}, p={:.4})",
            self.leader, self.lagger, self.correlation, self.lag_periods, self.p_value
        )
    }
}

/// Result of a Granger causality test.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct GrangerResult {
    /// F-statistic from the Granger causality F-test.
    #[pyo3(get)]
    pub f_statistic: f64,
    /// Approximate p-value of the F-test.
    #[pyo3(get)]
    pub p_value: f64,
    /// Lag order used in the test.
    #[pyo3(get)]
    pub lag: usize,
    /// Whether the test is significant at the 5% level.
    #[pyo3(get)]
    pub is_significant: bool,
}

#[pymethods]
impl GrangerResult {
    fn __repr__(&self) -> String {
        format!(
            "GrangerResult(F={:.4}, p={:.4}, lag={}, significant={})",
            self.f_statistic, self.p_value, self.lag, self.is_significant
        )
    }
}

/// A network of lead-lag relationships.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct LeadLagNetwork {
    /// Directed edges: (leader, lagger, correlation).
    #[pyo3(get)]
    pub edges: Vec<(String, String, f64)>,
    /// Names of series that tend to lead (more out-edges than in-edges).
    #[pyo3(get)]
    pub leaders: Vec<String>,
    /// Names of series that tend to lag (more in-edges than out-edges).
    #[pyo3(get)]
    pub laggers: Vec<String>,
}

#[pymethods]
impl LeadLagNetwork {
    fn __repr__(&self) -> String {
        format!(
            "LeadLagNetwork(edges={}, leaders={:?}, laggers={:?})",
            self.edges.len(),
            self.leaders,
            self.laggers
        )
    }
}

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

/// Clamp a value to finite, returning 0.0 for NaN/Inf.
#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

/// Compute the mean of a slice.
#[inline]
fn mean(data: &[f64]) -> f64 {
    if data.is_empty() {
        return 0.0;
    }
    data.iter().sum::<f64>() / data.len() as f64
}

/// Compute the standard deviation of a slice (population).
#[inline]
fn std_dev(data: &[f64]) -> f64 {
    if data.len() < 2 {
        return 0.0;
    }
    let m = mean(data);
    let var = data.iter().map(|x| (x - m).powi(2)).sum::<f64>() / data.len() as f64;
    var.sqrt()
}

/// Ordinary Least Squares regression.
/// Fits y = X * beta + epsilon.
/// X is (n_obs x n_features), y is (n_obs,).
/// Returns coefficients beta (n_features,) and residual sum of squares.
fn ols(x: &[Vec<f64>], y: &[f64]) -> Option<(Vec<f64>, f64)> {
    let n = y.len();
    if n == 0 {
        return None;
    }
    let p = x[0].len();
    if p == 0 || n <= p {
        return None;
    }

    // X^T X (p x p)
    let mut xtx = vec![vec![0.0f64; p]; p];
    for i in 0..p {
        for j in 0..p {
            let mut s = 0.0;
            for k in 0..n {
                s += x[k][i] * x[k][j];
            }
            xtx[i][j] = s;
        }
    }

    // X^T y (p,)
    let mut xty = vec![0.0f64; p];
    for i in 0..p {
        let mut s = 0.0;
        for k in 0..n {
            s += x[k][i] * y[k];
        }
        xty[i] = s;
    }

    // Solve via Cholesky: X^T X beta = X^T y
    let beta = cholesky_solve(&xtx, &xty)?;

    // Residual sum of squares
    let mut rss = 0.0;
    for k in 0..n {
        let mut pred = 0.0;
        for j in 0..p {
            pred += x[k][j] * beta[j];
        }
        rss += (y[k] - pred).powi(2);
    }

    Some((beta, rss))
}

/// Cholesky decomposition of a symmetric positive-definite matrix.
/// Returns lower-triangular L such that A = L L^T, or None if not PD.
fn cholesky_decompose(a: &[Vec<f64>]) -> Option<Vec<Vec<f64>>> {
    let n = a.len();
    let mut l = vec![vec![0.0f64; n]; n];

    for i in 0..n {
        for j in 0..=i {
            let mut s = 0.0;
            for k in 0..j {
                s += l[i][k] * l[j][k];
            }
            if i == j {
                let diag = a[i][i] - s;
                if diag <= 1e-15 {
                    return None; // Not positive definite
                }
                l[i][j] = diag.sqrt();
            } else {
                if l[j][j].abs() < 1e-30 {
                    return None;
                }
                l[i][j] = (a[i][j] - s) / l[j][j];
            }
        }
    }
    Some(l)
}

/// Solve A x = b given Cholesky factor L (A = L L^T).
/// Forward substitution: L z = b, then backward: L^T x = z.
fn cholesky_solve(a: &[Vec<f64>], b: &[f64]) -> Option<Vec<f64>> {
    let l = cholesky_decompose(a)?;
    let n = b.len();

    // Forward substitution: L z = b
    let mut z = vec![0.0f64; n];
    for i in 0..n {
        let mut s = 0.0;
        for j in 0..i {
            s += l[i][j] * z[j];
        }
        if l[i][i].abs() < 1e-30 {
            return None;
        }
        z[i] = (b[i] - s) / l[i][i];
    }

    // Backward substitution: L^T x = z
    let mut x = vec![0.0f64; n];
    for i in (0..n).rev() {
        let mut s = 0.0;
        for j in (i + 1)..n {
            s += l[j][i] * x[j];
        }
        if l[i][i].abs() < 1e-30 {
            return None;
        }
        x[i] = (z[i] - s) / l[i][i];
    }

    Some(x)
}

/// Approximate the upper-tail p-value for an F-distribution.
/// Uses the Abramowitz-Stegun normal approximation for large df.
/// F(d1, d2) is approximated by transforming to a chi-squared ratio.
fn f_distribution_p_value(f_stat: f64, d1: usize, d2: usize) -> f64 {
    if !f_stat.is_finite() || f_stat <= 0.0 || d1 == 0 || d2 == 0 {
        return 1.0;
    }

    let d1f = d1 as f64;
    let d2f = d2 as f64;

    // Use the normal approximation to the F-distribution
    // Based on the cube-root transformation (Paulson, 1942)
    let a = 2.0 / (9.0 * d1f);
    let b = 2.0 / (9.0 * d2f);
    let f_third = f_stat.powf(1.0 / 3.0);

    let num = (1.0 - b) * f_third - (1.0 - a);
    let den = (b * f_third * f_third + a).sqrt();

    if den.abs() < 1e-30 {
        return if f_stat > 1.0 { 0.0 } else { 1.0 };
    }

    let z = num / den;
    let p = 1.0 - normal_cdf(z);
    finite_or_zero(p.clamp(0.0, 1.0))
}

/// Standard normal CDF (Abramowitz & Stegun approximation).
#[inline]
fn normal_cdf(x: f64) -> f64 {
    const A1: f64 = 0.254829592;
    const A2: f64 = -0.284496736;
    const A3: f64 = 1.421413741;
    const A4: f64 = -1.453152027;
    const A5: f64 = 1.061405429;
    const P: f64 = 0.3275911;

    let sign = if x < 0.0 { -1.0 } else { 1.0 };
    let x_abs = x.abs() / std::f64::consts::SQRT_2;
    let t = 1.0 / (1.0 + P * x_abs);
    let y = 1.0 - (((((A5 * t + A4) * t) + A3) * t + A2) * t + A1) * t * (-x_abs * x_abs).exp();
    0.5 * (1.0 + sign * y)
}

/// Approximate p-value for a correlation coefficient using t-distribution approximation.
/// t = r * sqrt((n-2) / (1 - r^2)), then approximated via normal for large n.
#[allow(dead_code)]
fn correlation_p_value(r: f64, n: usize) -> f64 {
    if n < 3 {
        return 1.0;
    }
    let r_abs = r.abs().min(0.9999);
    let nf = n as f64;
    let t = r_abs * ((nf - 2.0) / (1.0 - r_abs * r_abs)).sqrt();
    // Two-sided p-value via normal approximation (good for n > 30)
    let p = 2.0 * (1.0 - normal_cdf(t.abs()));
    finite_or_zero(p.clamp(0.0, 1.0))
}

// ===========================================================================
// Public functions
// ===========================================================================

/// Hayashi-Yoshida asynchronous covariance estimator.
///
/// Computes the realized covariance between two non-synchronously observed
/// price series. Unlike standard covariance, this estimator correctly handles
/// irregularly spaced data by summing products of overlapping return intervals.
///
/// Reference: Hayashi & Yoshida (2005), "On covariance estimation of
/// non-synchronously observed diffusion processes."
///
/// Args:
///     ts_a: Timestamps for series A (must be sorted ascending).
///     prices_a: Prices for series A (same length as ts_a).
///     ts_b: Timestamps for series B (must be sorted ascending).
///     prices_b: Prices for series B (same length as ts_b).
///
/// Returns:
///     Hayashi-Yoshida covariance estimate.
#[pyfunction]
pub fn hayashi_yoshida(
    ts_a: Vec<f64>,
    prices_a: Vec<f64>,
    ts_b: Vec<f64>,
    prices_b: Vec<f64>,
) -> PyResult<f64> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if ts_a.len() != prices_a.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "ts_a and prices_a must have the same length",
        ));
    }
    if ts_b.len() != prices_b.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "ts_b and prices_b must have the same length",
        ));
    }
    if ts_a.len() < 2 || ts_b.len() < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "each series must have at least 2 observations",
        ));
    }

    // Validate finite values
    for (i, (t, p)) in ts_a.iter().zip(prices_a.iter()).enumerate() {
        if !t.is_finite() || !p.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "ts_a/prices_a[{}] is not finite",
                i
            )));
        }
    }
    for (i, (t, p)) in ts_b.iter().zip(prices_b.iter()).enumerate() {
        if !t.is_finite() || !p.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "ts_b/prices_b[{}] is not finite",
                i
            )));
        }
    }

    // Compute returns for each series
    let returns_a: Vec<f64> = (1..prices_a.len())
        .map(|i| prices_a[i] - prices_a[i - 1])
        .collect();
    let returns_b: Vec<f64> = (1..prices_b.len())
        .map(|i| prices_b[i] - prices_b[i - 1])
        .collect();

    // Hayashi-Yoshida: sum over all pairs of intervals that overlap in time
    let mut cov = 0.0;
    for i in 0..returns_a.len() {
        let a_start = ts_a[i];
        let a_end = ts_a[i + 1];

        for j in 0..returns_b.len() {
            let b_start = ts_b[j];
            let b_end = ts_b[j + 1];

            // Check if intervals overlap: max(start) < min(end)
            let overlap_start = a_start.max(b_start);
            let overlap_end = a_end.min(b_end);

            if overlap_start < overlap_end {
                cov += returns_a[i] * returns_b[j];
            }
        }
    }

    Ok(finite_or_zero(cov))
}

/// Compute cross-correlation between two series at multiple lags.
///
/// For each lag k in [-max_lag, +max_lag], computes the Pearson correlation
/// between x[max_lag..n-max_lag] and y shifted by k positions.
/// Positive lag means y is shifted forward (x leads y).
///
/// Args:
///     x: First time series.
///     y: Second time series (same length as x).
///     max_lag: Maximum lag to compute (positive integer).
///
/// Returns:
///     List of (lag, correlation) tuples sorted by lag.
#[pyfunction]
pub fn cross_correlation_lags(
    x: Vec<f64>,
    y: Vec<f64>,
    max_lag: i32,
) -> PyResult<Vec<(i32, f64)>> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if x.len() != y.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x and y must have the same length",
        ));
    }
    let n = x.len() as i32;
    if max_lag < 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "max_lag must be non-negative",
        ));
    }
    if n < 3 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series must have at least 3 observations",
        ));
    }
    let max_lag = max_lag.min(n - 2); // clamp to valid range

    let mut results = Vec::with_capacity((2 * max_lag + 1) as usize);

    for lag in -max_lag..=max_lag {
        // Determine overlapping range
        let (x_start, y_start) = if lag >= 0 {
            (0i32, lag)
        } else {
            (-lag, 0i32)
        };
        let len = (n - x_start).min(n - y_start);
        if len < 2 {
            results.push((lag, 0.0));
            continue;
        }

        let xs = x_start as usize;
        let ys = y_start as usize;
        let l = len as usize;

        let x_slice = &x[xs..xs + l];
        let y_slice = &y[ys..ys + l];

        let mx = mean(x_slice);
        let my = mean(y_slice);
        let sx = std_dev(x_slice);
        let sy = std_dev(y_slice);

        if sx < 1e-15 || sy < 1e-15 {
            results.push((lag, 0.0));
            continue;
        }

        let mut cov = 0.0;
        for k in 0..l {
            cov += (x_slice[k] - mx) * (y_slice[k] - my);
        }
        cov /= l as f64;

        let corr = cov / (sx * sy);
        results.push((lag, finite_or_zero(corr.clamp(-1.0, 1.0))));
    }

    Ok(results)
}

/// Granger causality F-test.
///
/// Tests whether past values of x help predict y beyond y's own past values.
/// Fits two OLS models:
///   Restricted:   y_t = c + sum(a_i * y_{t-i}) + e_t
///   Unrestricted: y_t = c + sum(a_i * y_{t-i}) + sum(b_i * x_{t-i}) + e_t
/// and computes the F-statistic for the joint significance of the b_i.
///
/// Args:
///     x: Potential causal series (Granger-causes y?).
///     y: Target series.
///     max_lag: Maximum lag order to test.
///
/// Returns:
///     GrangerResult with the best lag order (lowest p-value).
#[pyfunction]
pub fn granger_causality(
    x: Vec<f64>,
    y: Vec<f64>,
    max_lag: usize,
) -> PyResult<GrangerResult> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if x.len() != y.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x and y must have the same length",
        ));
    }
    if max_lag == 0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "max_lag must be at least 1",
        ));
    }
    let n = x.len();
    if n <= max_lag + 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series too short for the given max_lag",
        ));
    }

    // Validate finite values
    for (i, (&xv, &yv)) in x.iter().zip(y.iter()).enumerate() {
        if !xv.is_finite() || !yv.is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "x[{}] or y[{}] is not finite",
                i, i
            )));
        }
    }

    let mut best_f = 0.0f64;
    let mut best_p = 1.0f64;
    let mut best_lag = 1usize;

    for lag in 1..=max_lag {
        let t_start = lag;
        let n_obs = n - t_start;
        if n_obs <= 2 * lag + 2 {
            continue;
        }

        // Build dependent variable: y[t_start..n]
        let y_dep: Vec<f64> = y[t_start..n].to_vec();

        // Restricted model: y_t ~ constant + y_{t-1} + ... + y_{t-lag}
        let mut x_restricted = Vec::with_capacity(n_obs);
        for t in t_start..n {
            let mut row = Vec::with_capacity(lag + 1);
            row.push(1.0); // intercept
            for l in 1..=lag {
                row.push(y[t - l]);
            }
            x_restricted.push(row);
        }

        // Unrestricted model: y_t ~ constant + y_{t-1..lag} + x_{t-1..lag}
        let mut x_unrestricted = Vec::with_capacity(n_obs);
        for t in t_start..n {
            let mut row = Vec::with_capacity(2 * lag + 1);
            row.push(1.0); // intercept
            for l in 1..=lag {
                row.push(y[t - l]);
            }
            for l in 1..=lag {
                row.push(x[t - l]);
            }
            x_unrestricted.push(row);
        }

        // Fit both models
        let restricted = match ols(&x_restricted, &y_dep) {
            Some(r) => r,
            None => continue,
        };
        let unrestricted = match ols(&x_unrestricted, &y_dep) {
            Some(r) => r,
            None => continue,
        };

        let rss_r = restricted.1;
        let rss_u = unrestricted.1;

        // F = ((RSS_r - RSS_u) / q) / (RSS_u / (n - k))
        // where q = number of restrictions (lag), k = parameters in unrestricted model
        let q = lag as f64;
        let k = (2 * lag + 1) as f64;
        let df2 = n_obs as f64 - k;

        if df2 <= 0.0 || rss_u <= 0.0 {
            continue;
        }

        let f_stat = ((rss_r - rss_u) / q) / (rss_u / df2);
        let f_stat = finite_or_zero(f_stat.max(0.0));

        let p = f_distribution_p_value(f_stat, lag, df2 as usize);

        if p < best_p {
            best_p = p;
            best_f = f_stat;
            best_lag = lag;
        }
    }

    Ok(GrangerResult {
        f_statistic: finite_or_zero(best_f),
        p_value: finite_or_zero(best_p),
        lag: best_lag,
        is_significant: best_p < 0.05,
    })
}

/// Build a lead-lag network from pairwise cross-correlation analysis.
///
/// For each pair of series, computes cross-correlations at all lags and
/// identifies the dominant lead-lag relationship. Constructs a directed
/// network where edges point from leader to lagger.
///
/// Args:
///     series: List of time series (all same length).
///     names: List of series names (same length as series).
///     max_lag: Maximum lag to consider.
///
/// Returns:
///     LeadLagNetwork with directed edges and leader/lagger rankings.
#[pyfunction]
pub fn lead_lag_network(
    series: Vec<Vec<f64>>,
    names: Vec<String>,
    max_lag: i32,
) -> PyResult<LeadLagNetwork> {
    crate::auth::require_pro()
        .map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;

    if series.len() != names.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series and names must have the same length",
        ));
    }
    if series.len() < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "need at least 2 series to build a network",
        ));
    }
    if max_lag < 1 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "max_lag must be at least 1",
        ));
    }

    // Verify all series have the same length
    let n = series[0].len();
    for (i, s) in series.iter().enumerate() {
        if s.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "series[{}] has length {}, expected {}",
                i,
                s.len(),
                n
            )));
        }
    }

    let k = series.len();
    let mut edges = Vec::new();

    // Track net lead score: positive = tends to lead, negative = tends to lag
    let mut lead_scores: Vec<i32> = vec![0; k];

    for i in 0..k {
        for j in (i + 1)..k {
            // Compute cross-correlations between series i and j
            let cc = cross_correlation_lags(
                series[i].clone(),
                series[j].clone(),
                max_lag,
            )?;

            // Find the lag with maximum absolute correlation
            let mut best_lag = 0i32;
            let mut best_abs_corr = 0.0f64;
            let mut best_corr = 0.0f64;

            for &(lag, corr) in &cc {
                let ac = corr.abs();
                if ac > best_abs_corr {
                    best_abs_corr = ac;
                    best_lag = lag;
                    best_corr = corr;
                }
            }

            // Skip weak relationships
            if best_abs_corr < 0.05 {
                continue;
            }

            // Determine leader/lagger:
            // Positive lag means x (series i) leads y (series j) when corr > 0
            // at that lag. If the best lag is positive and corr > 0, i leads j.
            let (leader_idx, lagger_idx) = if best_lag > 0 {
                // Series i leads series j
                (i, j)
            } else if best_lag < 0 {
                // Series j leads series i
                (j, i)
            } else {
                // Lag 0: contemporaneous, skip direction
                continue;
            };

            edges.push((
                names[leader_idx].clone(),
                names[lagger_idx].clone(),
                finite_or_zero(best_corr.abs()),
            ));

            lead_scores[leader_idx] += 1;
            lead_scores[lagger_idx] -= 1;
        }
    }

    // Classify leaders (positive score) and laggers (negative score)
    let mut leaders = Vec::new();
    let mut laggers = Vec::new();

    let mut scored: Vec<(i32, &String)> = lead_scores.iter().zip(names.iter()).map(|(&s, n)| (s, n)).collect();
    scored.sort_by(|a, b| b.0.cmp(&a.0));

    for (score, name) in &scored {
        if *score > 0 {
            leaders.push((*name).clone());
        } else if *score < 0 {
            laggers.push((*name).clone());
        }
    }

    Ok(LeadLagNetwork {
        edges,
        leaders,
        laggers,
    })
}

/// Register lead-lag types and functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<LeadLagResult>()?;
    m.add_class::<GrangerResult>()?;
    m.add_class::<LeadLagNetwork>()?;
    m.add_function(wrap_pyfunction!(hayashi_yoshida, m)?)?;
    m.add_function(wrap_pyfunction!(cross_correlation_lags, m)?)?;
    m.add_function(wrap_pyfunction!(granger_causality, m)?)?;
    m.add_function(wrap_pyfunction!(lead_lag_network, m)?)?;
    Ok(())
}

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

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

    #[test]
    fn test_finite_or_zero() {
        assert_eq!(finite_or_zero(1.0), 1.0);
        assert_eq!(finite_or_zero(f64::NAN), 0.0);
        assert_eq!(finite_or_zero(f64::INFINITY), 0.0);
        assert_eq!(finite_or_zero(f64::NEG_INFINITY), 0.0);
    }

    #[test]
    fn test_mean_empty() {
        assert_eq!(mean(&[]), 0.0);
    }

    #[test]
    fn test_mean_values() {
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        assert!((mean(&data) - 3.0).abs() < 1e-10);
    }

    #[test]
    fn test_std_dev() {
        let data = vec![2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0];
        let sd = std_dev(&data);
        assert!(sd > 0.0);
        // Known population std dev of this dataset is 2.0
        assert!((sd - 2.0).abs() < 1e-10);
    }

    #[test]
    fn test_ols_simple() {
        // y = 2 + 3*x
        let x_data: Vec<Vec<f64>> = (0..10)
            .map(|i| vec![1.0, i as f64])
            .collect();
        let y_data: Vec<f64> = (0..10).map(|i| 2.0 + 3.0 * i as f64).collect();
        let (beta, rss) = ols(&x_data, &y_data).unwrap();
        assert!((beta[0] - 2.0).abs() < 1e-8);
        assert!((beta[1] - 3.0).abs() < 1e-8);
        assert!(rss < 1e-10);
    }

    #[test]
    fn test_hayashi_yoshida_synchronized() {
        // When timestamps are identical, HY should approximate standard covariance
        let ts: Vec<f64> = (0..100).map(|i| i as f64).collect();
        let prices_a: Vec<f64> = (0..100).map(|i| 100.0 + 0.1 * i as f64).collect();
        let prices_b: Vec<f64> = (0..100).map(|i| 50.0 + 0.05 * i as f64).collect();
        let cov = hayashi_yoshida(ts.clone(), prices_a, ts, prices_b).unwrap();
        assert!(cov > 0.0);
    }

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

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

    #[test]
    fn test_cross_correlation_self() {
        let x: Vec<f64> = (0..50).map(|i| (i as f64 * 0.1).sin()).collect();
        let cc = cross_correlation_lags(x.clone(), x, 5).unwrap();
        // At lag 0, self-correlation should be 1.0
        let lag_0 = cc.iter().find(|&&(lag, _)| lag == 0).unwrap();
        assert!((lag_0.1 - 1.0).abs() < 1e-10);
    }

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

    #[test]
    fn test_cross_correlation_lagged_signal() {
        // y is a lagged copy of x by 2 positions
        let n = 100;
        let x: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin()).collect();
        let mut y = vec![0.0; n];
        for i in 2..n {
            y[i] = x[i - 2];
        }
        let cc = cross_correlation_lags(x, y, 5).unwrap();
        // Best correlation should be at lag = +2 (x leads y)
        let mut best_lag = 0;
        let mut best_abs = 0.0;
        for &(lag, corr) in &cc {
            if corr.abs() > best_abs {
                best_abs = corr.abs();
                best_lag = lag;
            }
        }
        assert_eq!(best_lag, 2);
    }

    #[test]
    fn test_granger_causality_basic() {
        // Create a simple Granger-causal relationship: y_t = 0.5 * x_{t-1} + noise
        let n = 200;
        let mut rng: u64 = 42;
        let mut x = Vec::with_capacity(n);
        let mut y = Vec::with_capacity(n);

        for _ in 0..n {
            rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
            let u = ((rng >> 33) as f64 / (1u64 << 31) as f64 - 0.5) * 0.2;
            x.push(u);
        }
        y.push(0.0);
        for i in 1..n {
            rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1);
            let noise = ((rng >> 33) as f64 / (1u64 << 31) as f64 - 0.5) * 0.05;
            y.push(0.5 * x[i - 1] + noise);
        }

        let result = granger_causality(x, y, 3).unwrap();
        assert!(result.f_statistic > 0.0);
        assert!(result.lag >= 1);
    }

    #[test]
    fn test_granger_mismatched_lengths() {
        let result = granger_causality(vec![1.0; 50], vec![1.0; 30], 2);
        assert!(result.is_err());
    }

    #[test]
    fn test_granger_zero_lag() {
        let result = granger_causality(vec![1.0; 50], vec![1.0; 50], 0);
        assert!(result.is_err());
    }

    #[test]
    fn test_f_distribution_p_value() {
        // F=1 with equal df should give p ~ 0.5
        let p = f_distribution_p_value(1.0, 10, 10);
        assert!(p > 0.3 && p < 0.7, "p={}", p);

        // Large F should give small p
        let p2 = f_distribution_p_value(100.0, 5, 100);
        assert!(p2 < 0.01, "p2={}", p2);
    }

    #[test]
    fn test_lead_lag_network_basic() {
        let n = 80;
        let s1: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin()).collect();
        let mut s2 = vec![0.0; n];
        for i in 1..n {
            s2[i] = s1[i - 1]; // s2 lags s1 by 1
        }
        let names = vec!["alpha".to_string(), "beta".to_string()];
        let net = lead_lag_network(vec![s1, s2], names, 5).unwrap();
        // alpha should appear as leader
        if !net.edges.is_empty() {
            let first_edge = &net.edges[0];
            assert_eq!(first_edge.0, "alpha");
            assert_eq!(first_edge.1, "beta");
        }
    }

    #[test]
    fn test_lead_lag_network_too_few_series() {
        let result = lead_lag_network(
            vec![vec![1.0, 2.0, 3.0]],
            vec!["only".to_string()],
            1,
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_lead_lag_network_mismatched_names() {
        let result = lead_lag_network(
            vec![vec![1.0; 10], vec![2.0; 10]],
            vec!["a".to_string()],
            1,
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_normal_cdf_symmetric() {
        let p1 = normal_cdf(-1.0);
        let p2 = normal_cdf(1.0);
        assert!((p1 + p2 - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_correlation_p_value_high_corr() {
        let p = correlation_p_value(0.9, 100);
        assert!(p < 0.01);
    }

    #[test]
    fn test_correlation_p_value_zero_corr() {
        let p = correlation_p_value(0.0, 100);
        assert!((p - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_cholesky_solve() {
        // A = [[4, 2], [2, 3]], b = [1, 2]
        // Solution: x = [-0.125, 0.75]
        let a = vec![vec![4.0, 2.0], vec![2.0, 3.0]];
        let b = vec![1.0, 2.0];
        let x = cholesky_solve(&a, &b).unwrap();
        assert!((x[0] - (-0.125)).abs() < 1e-10);
        assert!((x[1] - 0.75).abs() < 1e-10);
    }
}
