//! Fractional Differentiation from Lopez de Prado's AFML Chapter 5.
//!
//! Provides fixed-width window (FFD), expanding window, and the minimum
//! differentiation order (min_d) search to make a series stationary while
//! preserving maximum memory.
//!
//! All functions are `#[pyfunction]` with `f64` at the PyO3 boundary.

use pyo3::prelude::*;

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

/// Finite-or-zero guard.
#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

// ===========================================================================
// Fractional differentiation weights
// ===========================================================================

/// Internal implementation of fractional differentiation weights (no tier check).
fn frac_diff_weights_impl(d: f64, threshold: f64) -> Vec<f64> {
    let threshold = threshold.abs().max(1e-15);
    let mut weights = vec![1.0_f64];
    let mut k = 1usize;
    loop {
        let w_prev = weights[k - 1];
        let w_k = -w_prev * (d - k as f64 + 1.0) / k as f64;
        if !w_k.is_finite() || w_k.abs() < threshold {
            break;
        }
        weights.push(w_k);
        k += 1;
        // Safety cap to avoid infinite loops for extreme d values
        if k > 100_000 {
            break;
        }
    }
    weights
}

/// Compute the fractional differentiation weights for order `d`.
///
/// Weights follow the recursion: w_k = -w_{k-1} * (d - k + 1) / k,
/// starting with w_0 = 1. Generation stops when |w_k| < `threshold`.
///
/// This is the core building block for both FFD and expanding-window methods.
#[pyfunction]
#[pyo3(signature = (d, threshold=1e-5))]
pub fn frac_diff_weights(d: f64, threshold: f64) -> PyResult<Vec<f64>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    Ok(frac_diff_weights_impl(d, threshold))
}

// ===========================================================================
// Fixed-Width Window Fractional Differentiation (Chapter 5.4)
// ===========================================================================

/// Fixed-Width Window Fractional Differentiation (FFD).
///
/// Computes weights via `frac_diff_weights(d, threshold)` and applies them as
/// a convolution over the series.  The output is shorter than the input by
/// `len(weights) - 1` because the first entries lack enough history.
///
/// This is the recommended method from AFML Chapter 5.4: it fixes the window
/// width (determined by the threshold) so every output point uses the same
/// number of lags, making the resulting series suitable for modelling.
#[pyfunction]
#[pyo3(signature = (series, d, threshold=1e-5))]
pub fn frac_diff_ffd(series: Vec<f64>, d: f64, threshold: f64) -> PyResult<Vec<f64>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    // --- input validation ---
    if series.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series must not be empty",
        ));
    }
    if !d.is_finite() || d < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "d must be a non-negative finite number",
        ));
    }
    if !threshold.is_finite() || threshold <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "threshold must be a positive finite number",
        ));
    }

    // Compute weights
    let weights = frac_diff_weights_impl(d, threshold);
    let w_len = weights.len();

    if w_len > series.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series too short for the given d and threshold (need more history)",
        ));
    }

    // Convolve: output[t] = sum_{k=0}^{w_len-1} w[k] * series[t - k]
    let n = series.len();
    let out_len = n - w_len + 1;
    let mut output = Vec::with_capacity(out_len);

    for t in (w_len - 1)..n {
        let mut val = 0.0_f64;
        for (k, &w) in weights.iter().enumerate() {
            val += w * series[t - k];
        }
        output.push(finite_or_zero(val));
    }

    Ok(output)
}

// ===========================================================================
// Expanding Window Fractional Differentiation (Chapter 5.3)
// ===========================================================================

/// Expanding-window (full-memory) fractional differentiation.
///
/// At each point `t`, uses all weights from lag 0 to lag `t`.  This preserves
/// the full information content of the original series but produces a
/// non-stationary weight structure (earlier points use fewer lags).
///
/// Slower than FFD (O(n^2) vs O(n * w_len)) but useful for analysis.
#[pyfunction]
pub fn frac_diff_expanding(series: Vec<f64>, d: f64) -> PyResult<Vec<f64>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    if series.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series must not be empty",
        ));
    }
    if !d.is_finite() || d < 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "d must be a non-negative finite number",
        ));
    }

    let n = series.len();

    // Pre-compute all weights up to length n (one per lag).
    // w_0 = 1, w_k = -w_{k-1} * (d - k + 1) / k
    let mut weights = Vec::with_capacity(n);
    weights.push(1.0_f64);
    for k in 1..n {
        let w_prev = weights[k - 1];
        let w_k = -w_prev * (d - k as f64 + 1.0) / k as f64;
        if !w_k.is_finite() {
            weights.push(0.0);
        } else {
            weights.push(w_k);
        }
    }

    // output[t] = sum_{k=0}^{t} w[k] * series[t - k]
    let mut output = Vec::with_capacity(n);
    for t in 0..n {
        let mut val = 0.0_f64;
        for k in 0..=t {
            val += weights[k] * series[t - k];
        }
        output.push(finite_or_zero(val));
    }

    Ok(output)
}

// ===========================================================================
// Augmented Dickey-Fuller test statistic (simplified, no lags)
// ===========================================================================

/// Simplified Augmented Dickey-Fuller test statistic (no augmenting lags).
///
/// Fits the regression: delta_y[t] = alpha + beta * y[t-1] + epsilon[t]
/// and returns ADF stat = beta / SE(beta).
///
/// More negative values indicate stronger stationarity evidence.
/// Critical values (approximate, n > 100):
///   1%: -3.43,  5%: -2.862,  10%: -2.567
#[pyfunction]
pub fn adf_statistic(series: Vec<f64>) -> PyResult<f64> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    let n = series.len();
    if n < 3 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series must have at least 3 observations for ADF test",
        ));
    }

    // Check for all-NaN or all-identical
    let all_finite = series.iter().all(|x| x.is_finite());
    if !all_finite {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series contains non-finite values (NaN or Inf)",
        ));
    }

    // delta_y = y[t] - y[t-1],  x = y[t-1]
    let m = n - 1; // number of observations in the regression
    let mut delta_y = Vec::with_capacity(m);
    let mut y_lag = Vec::with_capacity(m);
    for t in 1..n {
        delta_y.push(series[t] - series[t - 1]);
        y_lag.push(series[t - 1]);
    }

    // OLS regression: delta_y = alpha + beta * y_lag
    // Using the standard two-variable OLS formulas:
    //   beta = (sum(x*y) - n*mean_x*mean_y) / (sum(x^2) - n*mean_x^2)
    //   alpha = mean_y - beta * mean_x
    let mf = m as f64;
    let mean_dy = delta_y.iter().sum::<f64>() / mf;
    let mean_yl = y_lag.iter().sum::<f64>() / mf;

    let mut ss_xy = 0.0_f64; // sum of (x - mean_x)(y - mean_y)
    let mut ss_xx = 0.0_f64; // sum of (x - mean_x)^2
    for i in 0..m {
        let dx = y_lag[i] - mean_yl;
        let dy = delta_y[i] - mean_dy;
        ss_xy += dx * dy;
        ss_xx += dx * dx;
    }

    if ss_xx < 1e-30 {
        // y_lag is constant — series has zero variance, trivially stationary
        return Ok(-f64::INFINITY);
    }

    let beta = ss_xy / ss_xx;

    // Residuals and standard error of beta
    let alpha = mean_dy - beta * mean_yl;
    let mut sse = 0.0_f64; // sum of squared residuals
    for i in 0..m {
        let residual = delta_y[i] - alpha - beta * y_lag[i];
        sse += residual * residual;
    }

    // Degrees of freedom: m - 2 (two parameters estimated)
    let df = (m as i64 - 2).max(1) as f64;
    let mse = sse / df;
    let se_beta = (mse / ss_xx).sqrt();

    if se_beta < 1e-30 || !se_beta.is_finite() {
        // Perfect fit or degenerate — treat as highly stationary
        return Ok(-f64::INFINITY);
    }

    let adf = beta / se_beta;
    Ok(finite_or_zero(adf))
}

// ===========================================================================
// Minimum fractional differentiation order (Chapter 5.5)
// ===========================================================================

/// Find the minimum `d` that makes the series stationary (Chapter 5.5).
///
/// Searches `d` from 0 to `max_d` in `n_steps` equal increments.  For each `d`,
/// applies `frac_diff_ffd`, then computes the ADF test statistic.  Returns the
/// smallest `d` whose ADF stat is below the 5% critical value (-2.862).
///
/// Returns `(optimal_d, Vec<(d, adf_stat)>)`.  If no `d` in the range achieves
/// stationarity, `optimal_d` is set to `max_d`.
///
/// # Arguments
/// * `series` - The original price (or log-price) series.
/// * `p_threshold` - Unused (reserved for future p-value based stopping); kept
///   for API compatibility with the AFML interface.
/// * `max_d` - Upper bound on d to search (typically 1.0 or 2.0).
/// * `n_steps` - Number of grid points between 0 and max_d.
/// * `weight_threshold` - Threshold for weight truncation in FFD.
#[pyfunction]
#[pyo3(signature = (series, p_threshold=0.05, max_d=1.0, n_steps=20, weight_threshold=1e-5))]
pub fn min_frac_diff(
    series: Vec<f64>,
    #[allow(unused_variables)] p_threshold: f64,
    max_d: f64,
    n_steps: usize,
    weight_threshold: f64,
) -> PyResult<(f64, Vec<(f64, f64)>)> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    if series.len() < 10 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "series must have at least 10 observations for min_frac_diff",
        ));
    }
    if !max_d.is_finite() || max_d <= 0.0 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "max_d must be a positive finite number",
        ));
    }
    let n_steps = n_steps.max(2);
    let weight_threshold = weight_threshold.abs().max(1e-15);

    // ADF 5% critical value (asymptotic, constant term in regression)
    const ADF_CRITICAL_5PCT: f64 = -2.862;

    let step = max_d / n_steps as f64;
    let mut results: Vec<(f64, f64)> = Vec::with_capacity(n_steps + 1);
    let mut optimal_d = max_d;
    let mut found = false;

    for i in 0..=n_steps {
        let d = step * i as f64;

        // d = 0 means no differencing; compute ADF on original series
        let diffed = if d < 1e-15 {
            series.clone()
        } else {
            match frac_diff_ffd(series.clone(), d, weight_threshold) {
                Ok(v) if v.len() >= 3 => v,
                _ => {
                    // Not enough output points — skip this d
                    continue;
                }
            }
        };

        match adf_statistic(diffed) {
            Ok(adf) => {
                results.push((d, adf));
                if !found && adf < ADF_CRITICAL_5PCT {
                    optimal_d = d;
                    found = true;
                }
            }
            Err(_) => {
                // Skip if ADF computation fails (e.g., degenerate series)
                continue;
            }
        }
    }

    Ok((optimal_d, results))
}

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

/// Register all fractional differentiation functions in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(frac_diff_weights, m)?)?;
    m.add_function(wrap_pyfunction!(frac_diff_ffd, m)?)?;
    m.add_function(wrap_pyfunction!(frac_diff_expanding, m)?)?;
    m.add_function(wrap_pyfunction!(adf_statistic, m)?)?;
    m.add_function(wrap_pyfunction!(min_frac_diff, m)?)?;
    Ok(())
}

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

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

    // --- frac_diff_weights ---

    #[test]
    fn weights_d_zero() {
        let w = frac_diff_weights_impl(0.0, 1e-5);
        assert_eq!(w.len(), 1);
        assert!((w[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn weights_d_one() {
        // d=1 → w = [1, -1, 0, 0, ...] (first difference)
        let w = frac_diff_weights_impl(1.0, 1e-5);
        assert!(w.len() >= 2);
        assert!((w[0] - 1.0).abs() < 1e-10);
        assert!((w[1] - (-1.0)).abs() < 1e-10);
        // w_2 = -(-1)*(1-2+1)/2 = 0 → stops
        if w.len() > 2 {
            assert!(w[2].abs() < 1e-5);
        }
    }

    #[test]
    fn weights_d_half() {
        let w = frac_diff_weights_impl(0.5, 1e-5);
        // w_0 = 1, w_1 = -1*(0.5-0)/1 = -0.5, w_2 = -(-0.5)*(0.5-1)/2 = -0.125
        assert!(w.len() >= 3);
        assert!((w[0] - 1.0).abs() < 1e-10);
        assert!((w[1] - (-0.5)).abs() < 1e-10);
        assert!((w[2] - (-0.125)).abs() < 1e-10);
    }

    #[test]
    fn weights_all_diminish() {
        let w = frac_diff_weights_impl(0.3, 1e-8);
        // Weights should generally decrease in magnitude
        for k in 1..w.len() {
            assert!(
                w[k].abs() <= w[0].abs() + 1e-10,
                "w[{}]={} exceeds w[0]={}",
                k,
                w[k],
                w[0]
            );
        }
    }

    // --- frac_diff_ffd ---

    #[test]
    fn ffd_d_one_is_first_diff() {
        // d=1.0, large threshold so only [1, -1] weights
        let series: Vec<f64> = vec![100.0, 102.0, 101.0, 105.0, 103.0];
        let result = frac_diff_ffd(series, 1.0, 0.5).unwrap();
        // weights = [1, -1], output[t] = series[t] - series[t-1]
        assert_eq!(result.len(), 4);
        assert!((result[0] - 2.0).abs() < 1e-10);
        assert!((result[1] - (-1.0)).abs() < 1e-10);
        assert!((result[2] - 4.0).abs() < 1e-10);
        assert!((result[3] - (-2.0)).abs() < 1e-10);
    }

    #[test]
    fn ffd_d_zero_is_identity() {
        let series: Vec<f64> = vec![10.0, 20.0, 30.0];
        let result = frac_diff_ffd(series.clone(), 0.0, 1e-5).unwrap();
        // d=0 → weights = [1] → output = series itself
        assert_eq!(result.len(), 3);
        for (a, b) in result.iter().zip(series.iter()) {
            assert!((a - b).abs() < 1e-10);
        }
    }

    #[test]
    fn ffd_empty_series() {
        assert!(frac_diff_ffd(vec![], 0.5, 1e-5).is_err());
    }

    #[test]
    fn ffd_negative_d() {
        assert!(frac_diff_ffd(vec![1.0, 2.0, 3.0], -0.5, 1e-5).is_err());
    }

    #[test]
    fn ffd_invalid_threshold() {
        assert!(frac_diff_ffd(vec![1.0, 2.0, 3.0], 0.5, 0.0).is_err());
        assert!(frac_diff_ffd(vec![1.0, 2.0, 3.0], 0.5, -1.0).is_err());
    }

    #[test]
    fn ffd_output_length() {
        let series: Vec<f64> = (0..100).map(|i| i as f64).collect();
        let result = frac_diff_ffd(series.clone(), 0.5, 1e-3).unwrap();
        let weights = frac_diff_weights_impl(0.5, 1e-3);
        assert_eq!(result.len(), series.len() - weights.len() + 1);
    }

    // --- frac_diff_expanding ---

    #[test]
    fn expanding_d_zero() {
        let series = vec![5.0, 10.0, 15.0];
        let result = frac_diff_expanding(series.clone(), 0.0).unwrap();
        assert_eq!(result.len(), 3);
        for (a, b) in result.iter().zip(series.iter()) {
            assert!((a - b).abs() < 1e-10);
        }
    }

    #[test]
    fn expanding_d_one() {
        let series = vec![100.0, 102.0, 105.0];
        let result = frac_diff_expanding(series, 1.0).unwrap();
        assert_eq!(result.len(), 3);
        // t=0: w=[1]*100 = 100
        assert!((result[0] - 100.0).abs() < 1e-10);
        // t=1: w=[1,-1], 1*102 + (-1)*100 = 2
        assert!((result[1] - 2.0).abs() < 1e-10);
        // t=2: w=[1,-1,0], 1*105 + (-1)*102 + 0*100 = 3
        assert!((result[2] - 3.0).abs() < 1e-10);
    }

    #[test]
    fn expanding_preserves_length() {
        let series: Vec<f64> = (0..50).map(|i| i as f64).collect();
        let result = frac_diff_expanding(series.clone(), 0.4).unwrap();
        assert_eq!(result.len(), series.len());
    }

    #[test]
    fn expanding_empty() {
        assert!(frac_diff_expanding(vec![], 0.5).is_err());
    }

    #[test]
    fn expanding_negative_d() {
        assert!(frac_diff_expanding(vec![1.0, 2.0], -0.1).is_err());
    }

    // --- adf_statistic ---

    #[test]
    fn adf_random_walk_not_stationary() {
        // Random walk: y[t] = y[t-1] + noise → ADF should NOT be very negative
        let mut series = vec![0.0_f64];
        let mut state = 0.0_f64;
        for i in 1..200 {
            state += (i as f64 * 0.37).sin() * 0.01;
            series.push(state);
        }
        let adf = adf_statistic(series).unwrap();
        // A random walk should have an ADF stat close to 0 (not rejecting unit root)
        assert!(adf > -5.0, "adf={} unexpectedly stationary", adf);
    }

    #[test]
    fn adf_stationary_series() {
        // Mean-reverting series: y[t] = 0.5 * y[t-1] + noise
        let mut series = vec![0.0_f64];
        for i in 1..300 {
            let prev = series[i - 1];
            let noise = (i as f64 * 0.73).sin() * 0.01;
            series.push(0.3 * prev + noise);
        }
        let adf = adf_statistic(series).unwrap();
        // Should be very negative (stationary)
        assert!(adf < -2.0, "adf={} not stationary enough", adf);
    }

    #[test]
    fn adf_too_short() {
        assert!(adf_statistic(vec![1.0, 2.0]).is_err());
    }

    #[test]
    fn adf_constant_series() {
        let series = vec![5.0; 100];
        let adf = adf_statistic(series).unwrap();
        // Constant series is trivially stationary
        assert!(adf == f64::NEG_INFINITY || adf < -10.0);
    }

    #[test]
    fn adf_nan_input() {
        assert!(adf_statistic(vec![1.0, f64::NAN, 3.0]).is_err());
    }

    // --- min_frac_diff ---

    #[test]
    fn min_frac_diff_trending_series() {
        // Trending series: needs some differentiation
        let series: Vec<f64> = (0..200).map(|i| i as f64 * 0.1 + (i as f64 * 0.05).sin()).collect();
        let (d, results) = min_frac_diff(series, 0.05, 1.0, 10, 1e-3).unwrap();
        assert!(d >= 0.0 && d <= 1.0, "d={}", d);
        assert!(!results.is_empty());
        // Results should be sorted by d
        for window in results.windows(2) {
            assert!(window[0].0 <= window[1].0);
        }
    }

    #[test]
    fn min_frac_diff_stationary_series() {
        // Already stationary: d should be 0 or very small
        let series: Vec<f64> = (0..200)
            .map(|i| (i as f64 * 0.1).sin() * 0.01)
            .collect();
        let (d, results) = min_frac_diff(series, 0.05, 1.0, 20, 1e-3).unwrap();
        assert!(d <= 0.2, "d={} too high for stationary series", d);
        assert!(!results.is_empty());
    }

    #[test]
    fn min_frac_diff_too_short() {
        assert!(min_frac_diff(vec![1.0; 5], 0.05, 1.0, 10, 1e-3).is_err());
    }

    #[test]
    fn min_frac_diff_invalid_max_d() {
        let series: Vec<f64> = (0..50).map(|i| i as f64).collect();
        assert!(min_frac_diff(series.clone(), 0.05, 0.0, 10, 1e-3).is_err());
        assert!(min_frac_diff(series, 0.05, f64::NAN, 10, 1e-3).is_err());
    }

    #[test]
    fn min_frac_diff_returns_tuple() {
        let series: Vec<f64> = (0..100).map(|i| (i as f64).ln_1p()).collect();
        let (d, scan) = min_frac_diff(series, 0.05, 2.0, 10, 1e-3).unwrap();
        assert!(d.is_finite());
        for &(di, adf_i) in &scan {
            assert!(di.is_finite());
            assert!(adf_i.is_finite() || adf_i == f64::NEG_INFINITY);
        }
    }
}
