//! Granger Causality analysis for effective connectivity.
//!
//! Granger causality measures the directed influence of one time series on another
//! by testing whether past values of X help predict Y beyond past values of Y alone.
//!
//! ## Time Domain Granger Causality
//!
//! The Granger causality from X to Y is:
//! GC(X→Y) = ln(var(Y|Y_past) / var(Y|Y_past,X_past))
//!
//! ## Spectral Granger Causality
//!
//! Decomposes Granger causality across frequencies using the
//! Geweke spectral measure.
//!
//! ## References
//!
//! - Granger, C. W. (1969). Investigating causal relations by econometric models.
//! - Geweke, J. (1982). Measurement of linear dependence and feedback between
//!   multiple time series.

use crate::{ConnectivityError, ConnectivityResult};
use nalgebra::{DMatrix, DVector};
use num_complex::Complex64;
use std::f64::consts::PI;

/// Configuration for Granger causality analysis
#[derive(Debug, Clone)]
pub struct GrangerConfig {
    /// Model order (number of lags)
    pub order: usize,
    /// Criterion for automatic order selection
    pub order_criterion: OrderCriterion,
    /// Maximum order for automatic selection
    pub max_order: usize,
    /// Include instantaneous effects
    pub include_instantaneous: bool,
}

impl Default for GrangerConfig {
    fn default() -> Self {
        Self {
            order: 10,
            order_criterion: OrderCriterion::Bic,
            max_order: 20,
            include_instantaneous: false,
        }
    }
}

impl GrangerConfig {
    /// Create config with specific model order
    pub fn with_order(order: usize) -> Self {
        Self {
            order,
            ..Self::default()
        }
    }

    /// Create config with automatic order selection
    pub fn auto_order(max_order: usize, criterion: OrderCriterion) -> Self {
        Self {
            order: 0, // Will be determined automatically
            order_criterion: criterion,
            max_order,
            ..Self::default()
        }
    }
}

/// Criterion for selecting VAR model order
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrderCriterion {
    /// Akaike Information Criterion
    Aic,
    /// Bayesian Information Criterion
    Bic,
    /// Hannan-Quinn Criterion
    Hqc,
    /// Use fixed order
    Fixed,
}

/// Result of Granger causality test
#[derive(Debug, Clone)]
pub struct GrangerResult {
    /// Granger causality value (log ratio of variances)
    pub gc_value: f64,
    /// F-statistic for causality test
    pub f_statistic: f64,
    /// P-value (if computed)
    pub p_value: Option<f64>,
    /// Degrees of freedom (numerator, denominator)
    pub df: (usize, usize),
    /// Model order used
    pub order: usize,
    /// Residual variance (restricted model)
    pub var_restricted: f64,
    /// Residual variance (unrestricted model)
    pub var_unrestricted: f64,
}

/// Result of spectral Granger causality
#[derive(Debug, Clone)]
pub struct SpectralGrangerResult {
    /// Frequencies
    pub frequencies: Vec<f64>,
    /// Granger causality X → Y at each frequency
    pub gc_x_to_y: Vec<f64>,
    /// Granger causality Y → X at each frequency
    pub gc_y_to_x: Vec<f64>,
    /// Instantaneous causality at each frequency
    pub gc_instantaneous: Vec<f64>,
    /// Total coherence-like measure
    pub gc_total: Vec<f64>,
    /// Model order used
    pub order: usize,
}

/// Vector Autoregressive (VAR) model for Granger causality
#[derive(Debug)]
pub struct VarModel {
    /// Coefficient matrices A_1, A_2, ..., A_p for each lag
    coefficients: Vec<DMatrix<f64>>,
    /// Residual covariance matrix
    residual_cov: DMatrix<f64>,
    /// Model order
    order: usize,
    /// Number of variables
    n_vars: usize,
}

impl VarModel {
    /// Fit a VAR model to multivariate time series
    ///
    /// # Arguments
    /// * `data` - Multivariate time series [n_vars][n_samples]
    /// * `order` - Model order (number of lags)
    pub fn fit(data: &[Vec<f64>], order: usize) -> ConnectivityResult<Self> {
        let n_vars = data.len();
        let n_samples = data[0].len();

        if n_samples <= order * n_vars + 1 {
            return Err(ConnectivityError::InsufficientData(format!(
                "Need more samples for order {} model (have {}, need > {})",
                order,
                n_samples,
                order * n_vars + 1
            )));
        }

        // Build regression matrices
        // Y = [y_p+1, ..., y_T]^T   (T-p x n_vars)
        // X = [Y_lagged]            (T-p x p*n_vars)

        let t = n_samples - order;

        // Dependent variable matrix Y
        let mut y_mat = DMatrix::zeros(t, n_vars);
        for j in 0..n_vars {
            for i in 0..t {
                y_mat[(i, j)] = data[j][order + i];
            }
        }

        // Lagged predictor matrix X
        let n_pred = order * n_vars;
        let mut x_mat = DMatrix::zeros(t, n_pred);
        for i in 0..t {
            for lag in 0..order {
                for j in 0..n_vars {
                    x_mat[(i, lag * n_vars + j)] = data[j][order + i - lag - 1];
                }
            }
        }

        // OLS estimation: B = (X'X)^(-1) X'Y
        let xtx = x_mat.transpose() * &x_mat;
        let xty = x_mat.transpose() * &y_mat;

        let xtx_inv = Self::pseudoinverse(&xtx)?;
        let b = &xtx_inv * &xty;

        // Extract coefficient matrices
        let mut coefficients = Vec::with_capacity(order);
        for lag in 0..order {
            let mut a_lag = DMatrix::zeros(n_vars, n_vars);
            for i in 0..n_vars {
                for j in 0..n_vars {
                    a_lag[(i, j)] = b[(lag * n_vars + j, i)];
                }
            }
            coefficients.push(a_lag);
        }

        // Compute residuals and covariance
        let y_pred = &x_mat * &b;
        let residuals = &y_mat - &y_pred;
        let residual_cov = (residuals.transpose() * &residuals) / (t - n_pred) as f64;

        Ok(Self {
            coefficients,
            residual_cov,
            order,
            n_vars,
        })
    }

    /// Get coefficient matrix for a specific lag
    pub fn get_coefficients(&self, lag: usize) -> Option<&DMatrix<f64>> {
        self.coefficients.get(lag)
    }

    /// Get residual covariance matrix
    pub fn residual_covariance(&self) -> &DMatrix<f64> {
        &self.residual_cov
    }

    /// Compute AIC for model selection
    pub fn aic(&self, n_samples: usize) -> f64 {
        let n = n_samples - self.order;
        let det = self.residual_cov.determinant().max(1e-30);
        let k = (self.order * self.n_vars * self.n_vars) as f64;

        n as f64 * det.ln() + 2.0 * k
    }

    /// Compute BIC for model selection
    pub fn bic(&self, n_samples: usize) -> f64 {
        let n = n_samples - self.order;
        let det = self.residual_cov.determinant().max(1e-30);
        let k = (self.order * self.n_vars * self.n_vars) as f64;

        n as f64 * det.ln() + k * (n as f64).ln()
    }

    /// Compute transfer function H(f) at a frequency
    pub fn transfer_function(&self, freq: f64, sfreq: f64) -> DMatrix<Complex64> {
        let n = self.n_vars;
        let mut a_f = DMatrix::from_element(n, n, Complex64::new(0.0, 0.0));

        for (lag, a) in self.coefficients.iter().enumerate() {
            let exp_factor = Complex64::new(0.0, -2.0 * PI * freq * (lag + 1) as f64 / sfreq);
            let exp_val = exp_factor.exp();

            for i in 0..n {
                for j in 0..n {
                    a_f[(i, j)] += Complex64::new(a[(i, j)], 0.0) * exp_val;
                }
            }
        }

        // I - A(f)
        let mut result = DMatrix::from_element(n, n, Complex64::new(0.0, 0.0));
        for i in 0..n {
            result[(i, i)] = Complex64::new(1.0, 0.0);
        }

        for i in 0..n {
            for j in 0..n {
                result[(i, j)] -= a_f[(i, j)];
            }
        }

        // H(f) = (I - A(f))^(-1)
        Self::complex_inverse(&result).unwrap_or(result)
    }

    /// Pseudoinverse using SVD
    fn pseudoinverse(m: &DMatrix<f64>) -> ConnectivityResult<DMatrix<f64>> {
        let svd = m.clone().svd(true, true);

        let u = svd
            .u
            .ok_or_else(|| ConnectivityError::ComputationError("SVD failed".to_string()))?;

        let vt = svd
            .v_t
            .ok_or_else(|| ConnectivityError::ComputationError("SVD failed".to_string()))?;

        let s = svd.singular_values;
        let tol = 1e-10 * s[0];

        let n = s.len().min(m.ncols());
        let mut s_inv = DMatrix::zeros(m.ncols(), m.nrows());

        for i in 0..n {
            if s[i] > tol {
                s_inv[(i, i)] = 1.0 / s[i];
            }
        }

        Ok(vt.transpose() * &s_inv * u.transpose())
    }

    /// Complex matrix inverse
    fn complex_inverse(m: &DMatrix<Complex64>) -> ConnectivityResult<DMatrix<Complex64>> {
        let n = m.nrows();
        if n != m.ncols() {
            return Err(ConnectivityError::ComputationError(
                "Matrix not square".to_string(),
            ));
        }

        // Simple Gaussian elimination for small matrices
        let mut aug = DMatrix::from_element(n, 2 * n, Complex64::new(0.0, 0.0));

        // Copy m to left half
        for i in 0..n {
            for j in 0..n {
                aug[(i, j)] = m[(i, j)];
            }
            aug[(i, n + i)] = Complex64::new(1.0, 0.0);
        }

        // Forward elimination
        for i in 0..n {
            // Find pivot
            let mut max_row = i;
            let mut max_val = aug[(i, i)].norm();
            for k in (i + 1)..n {
                if aug[(k, i)].norm() > max_val {
                    max_val = aug[(k, i)].norm();
                    max_row = k;
                }
            }

            // Swap rows
            for j in 0..(2 * n) {
                let tmp = aug[(i, j)];
                aug[(i, j)] = aug[(max_row, j)];
                aug[(max_row, j)] = tmp;
            }

            let pivot = aug[(i, i)];
            if pivot.norm() < 1e-15 {
                return Err(ConnectivityError::ComputationError(
                    "Singular matrix".to_string(),
                ));
            }

            // Scale row
            for j in 0..(2 * n) {
                aug[(i, j)] /= pivot;
            }

            // Eliminate column
            for k in 0..n {
                if k != i {
                    let factor = aug[(k, i)];
                    for j in 0..(2 * n) {
                        let val = aug[(i, j)];
                        aug[(k, j)] -= factor * val;
                    }
                }
            }
        }

        // Extract inverse from right half
        let mut inv = DMatrix::from_element(n, n, Complex64::new(0.0, 0.0));
        for i in 0..n {
            for j in 0..n {
                inv[(i, j)] = aug[(i, n + j)];
            }
        }

        Ok(inv)
    }
}

/// Compute time-domain Granger causality from X to Y
///
/// # Arguments
/// * `x` - Source time series
/// * `y` - Target time series
/// * `order` - Model order (number of lags)
///
/// # Returns
/// Granger causality result including GC value, F-statistic, and p-value
pub fn granger_causality(x: &[f64], y: &[f64], order: usize) -> ConnectivityResult<GrangerResult> {
    if x.len() != y.len() {
        return Err(ConnectivityError::DimensionMismatch(
            "Time series must have same length".to_string(),
        ));
    }

    let n = x.len();
    if n <= 2 * order + 1 {
        return Err(ConnectivityError::InsufficientData(
            "Not enough samples for given order".to_string(),
        ));
    }

    // Restricted model: Y ~ Y_past only
    let var_restricted = fit_ar_variance(y, order)?;

    // Unrestricted model: Y ~ Y_past + X_past (VAR model)
    let data = vec![y.to_vec(), x.to_vec()];
    let var_model = VarModel::fit(&data, order)?;
    let var_unrestricted = var_model.residual_cov[(0, 0)];

    // Granger causality: ln(var_restricted / var_unrestricted)
    let gc_value = (var_restricted / var_unrestricted).ln();

    // F-statistic for testing significance
    // F = ((RSS_r - RSS_u) / p) / (RSS_u / (T - 2p - 1))
    let t = n - order;
    let df_num = order;
    let df_den = t - 2 * order - 1;

    let f_statistic = if df_den > 0 && var_unrestricted > 1e-30 {
        ((var_restricted - var_unrestricted) * df_den as f64) / (var_unrestricted * df_num as f64)
    } else {
        0.0
    };

    // P-value using F-distribution approximation (simplified)
    let p_value = if f_statistic > 0.0 && df_den > 0 {
        Some(f_distribution_sf(f_statistic, df_num, df_den))
    } else {
        None
    };

    Ok(GrangerResult {
        gc_value: gc_value.max(0.0),
        f_statistic: f_statistic.max(0.0),
        p_value,
        df: (df_num, df_den),
        order,
        var_restricted,
        var_unrestricted,
    })
}

/// Compute spectral Granger causality (Geweke's measure)
///
/// # Arguments
/// * `x` - First time series
/// * `y` - Second time series
/// * `sfreq` - Sampling frequency
/// * `order` - Model order
/// * `n_freqs` - Number of frequency points
pub fn spectral_granger(
    x: &[f64],
    y: &[f64],
    sfreq: f64,
    order: usize,
    n_freqs: Option<usize>,
) -> ConnectivityResult<SpectralGrangerResult> {
    if x.len() != y.len() {
        return Err(ConnectivityError::DimensionMismatch(
            "Time series must have same length".to_string(),
        ));
    }

    // Fit bivariate VAR model
    let data = vec![x.to_vec(), y.to_vec()];
    let var_model = VarModel::fit(&data, order)?;

    let sigma = var_model.residual_covariance();

    let n_freqs = n_freqs.unwrap_or(128);
    let df = sfreq / (2.0 * n_freqs as f64);
    let frequencies: Vec<f64> = (0..n_freqs).map(|i| i as f64 * df).collect();

    let mut gc_x_to_y = Vec::with_capacity(n_freqs);
    let mut gc_y_to_x = Vec::with_capacity(n_freqs);
    let mut gc_instantaneous = Vec::with_capacity(n_freqs);
    let mut gc_total = Vec::with_capacity(n_freqs);

    for &freq in &frequencies {
        // Compute transfer function H(f)
        let h = var_model.transfer_function(freq, sfreq);

        // Spectral matrix S(f) = H(f) * Sigma * H(f)^H
        let h_conj = h.transpose().map(|c| c.conj());

        // S(f) computation
        let mut s = DMatrix::from_element(2, 2, Complex64::new(0.0, 0.0));
        for i in 0..2 {
            for j in 0..2 {
                for k in 0..2 {
                    for l in 0..2 {
                        s[(i, j)] +=
                            h[(i, k)] * Complex64::new(sigma[(k, l)], 0.0) * h_conj[(l, j)];
                    }
                }
            }
        }

        // Geweke's spectral Granger causality
        let s_xx = s[(0, 0)].re;
        let s_yy = s[(1, 1)].re;

        // Intrinsic power (what would remain without the other signal)
        let sigma_xx = sigma[(0, 0)];
        let sigma_yy = sigma[(1, 1)];

        let h_xx = h[(0, 0)].norm_sqr();
        let h_yy = h[(1, 1)].norm_sqr();
        let _h_xy = h[(0, 1)].norm_sqr();
        let _h_yx = h[(1, 0)].norm_sqr();

        // GC X->Y: based on how much H_yx contributes
        let gc_xy = if s_yy > 1e-30 {
            let intrinsic_y = h_yy * sigma_yy;
            (s_yy / intrinsic_y.max(1e-30)).ln().max(0.0)
        } else {
            0.0
        };

        // GC Y->X: based on how much H_xy contributes
        let gc_yx = if s_xx > 1e-30 {
            let intrinsic_x = h_xx * sigma_xx;
            (s_xx / intrinsic_x.max(1e-30)).ln().max(0.0)
        } else {
            0.0
        };

        // Instantaneous causality (correlation of innovations)
        let gc_inst = if sigma[(0, 0)] > 1e-30 && sigma[(1, 1)] > 1e-30 {
            let r = sigma[(0, 1)] / (sigma[(0, 0)] * sigma[(1, 1)]).sqrt();
            -(1.0 - r * r).ln().max(0.0)
        } else {
            0.0
        };

        // Total
        let gc_t = gc_xy + gc_yx + gc_inst;

        gc_x_to_y.push(gc_xy);
        gc_y_to_x.push(gc_yx);
        gc_instantaneous.push(gc_inst);
        gc_total.push(gc_t);
    }

    Ok(SpectralGrangerResult {
        frequencies,
        gc_x_to_y,
        gc_y_to_x,
        gc_instantaneous,
        gc_total,
        order,
    })
}

/// Compute Granger causality for all channel pairs
pub fn granger_all_pairs(
    data: &[Vec<f64>], // [n_channels][n_samples]
    order: usize,
) -> ConnectivityResult<Vec<Vec<f64>>> {
    let n_channels = data.len();

    // Generate all directed pairs
    let mut results = Vec::new();

    for i in 0..n_channels {
        for j in 0..n_channels {
            if i != j {
                let gc = granger_causality(&data[i], &data[j], order)?;
                results.push(vec![gc.gc_value]); // Single value per pair
            }
        }
    }

    Ok(results)
}

/// Select optimal model order using information criterion
pub fn select_order(
    data: &[Vec<f64>],
    max_order: usize,
    criterion: OrderCriterion,
) -> ConnectivityResult<usize> {
    let n_samples = data[0].len();

    let mut best_order = 1;
    let mut best_criterion = f64::INFINITY;

    for order in 1..=max_order {
        if let Ok(model) = VarModel::fit(data, order) {
            let crit = match criterion {
                OrderCriterion::Aic => model.aic(n_samples),
                OrderCriterion::Bic => model.bic(n_samples),
                OrderCriterion::Hqc => {
                    // HQ = n * ln(det) + 2k * ln(ln(n))
                    let n = n_samples - order;
                    let det = model.residual_cov.determinant().max(1e-30);
                    let k = (order * data.len() * data.len()) as f64;
                    n as f64 * det.ln() + 2.0 * k * (n as f64).ln().ln()
                }
                OrderCriterion::Fixed => return Ok(max_order),
            };

            if crit < best_criterion {
                best_criterion = crit;
                best_order = order;
            }
        }
    }

    Ok(best_order)
}

// ========== Helper functions ==========

/// Fit AR model and return residual variance
fn fit_ar_variance(y: &[f64], order: usize) -> ConnectivityResult<f64> {
    let n = y.len();
    let t = n - order;

    // Build regression matrix
    let mut x_mat = DMatrix::zeros(t, order);
    let mut y_vec = DVector::zeros(t);

    for i in 0..t {
        y_vec[i] = y[order + i];
        for j in 0..order {
            x_mat[(i, j)] = y[order + i - j - 1];
        }
    }

    // OLS
    let xtx = x_mat.transpose() * &x_mat;
    let xty = x_mat.transpose() * &y_vec;

    let xtx_inv = VarModel::pseudoinverse(&xtx)?;
    let b = &xtx_inv * &xty;

    // Residual variance
    let y_pred = &x_mat * &b;
    let residuals = &y_vec - &y_pred;

    let var = residuals.dot(&residuals) / (t - order) as f64;
    Ok(var)
}

/// Survival function of F-distribution (approximate)
fn f_distribution_sf(f: f64, df1: usize, df2: usize) -> f64 {
    if f <= 0.0 || df1 == 0 || df2 == 0 {
        return 1.0;
    }

    // Use Beta distribution relationship:
    // F(x; d1, d2) = I_{d1*x/(d1*x+d2)}(d1/2, d2/2)
    // where I is the regularized incomplete beta function

    // Simplified approximation for large df
    let d1 = df1 as f64;
    let d2 = df2 as f64;
    let _x = d1 * f / (d1 * f + d2);

    // Wilson-Hilferty approximation
    let _a = d1 / 2.0;
    let _b = d2 / 2.0;

    // Simple normal approximation for p-value
    let mean = d2 / (d2 - 2.0).max(1.0);
    let var = 2.0 * d2 * d2 * (d1 + d2 - 2.0) / (d1 * (d2 - 2.0).powi(2) * (d2 - 4.0).max(1.0));

    let z = (f - mean) / var.sqrt();
    0.5 * (1.0 - erf(z / 2.0_f64.sqrt()))
}

/// Error function approximation
fn erf(x: f64) -> f64 {
    // Horner's method approximation
    let a1 = 0.254829592;
    let a2 = -0.284496736;
    let a3 = 1.421413741;
    let a4 = -1.453152027;
    let a5 = 1.061405429;
    let p = 0.3275911;

    let sign = if x < 0.0 { -1.0 } else { 1.0 };
    let x = x.abs();

    let t = 1.0 / (1.0 + p * x);
    let y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * (-x * x).exp();

    sign * y
}

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

    fn create_coupled_series(n: usize, coupling: f64) -> (Vec<f64>, Vec<f64>) {
        let mut x = vec![0.0; n];
        let mut y = vec![0.0; n];

        // AR(1) process for x
        for i in 1..n {
            x[i] = 0.8 * x[i - 1] + 0.1 * ((i as f64 * 0.1).sin());
        }

        // Y depends on past X (Granger causality from X to Y)
        for i in 1..n {
            y[i] = 0.5 * y[i - 1] + coupling * x[i - 1] + 0.1 * ((i as f64 * 0.2).cos());
        }

        (x, y)
    }

    #[test]
    fn test_granger_causality_coupled() {
        let (x, y) = create_coupled_series(500, 0.5);

        // X should cause Y
        let gc_x_to_y = granger_causality(&x, &y, 5).unwrap();
        assert!(gc_x_to_y.gc_value > 0.0, "GC X→Y should be positive");

        // Y should not cause X (or weakly)
        let gc_y_to_x = granger_causality(&y, &x, 5).unwrap();

        // GC X→Y should be much larger than GC Y→X
        assert!(
            gc_x_to_y.gc_value > gc_y_to_x.gc_value * 0.5,
            "GC X→Y ({}) should be larger than GC Y→X ({})",
            gc_x_to_y.gc_value,
            gc_y_to_x.gc_value
        );
    }

    #[test]
    fn test_granger_uncoupled() {
        // Create independent series
        let x: Vec<f64> = (0..500).map(|i| (i as f64 * 0.1).sin()).collect();
        let y: Vec<f64> = (0..500).map(|i| (i as f64 * 0.17).cos()).collect();

        let gc = granger_causality(&x, &y, 5).unwrap();

        // GC should be small for independent series
        assert!(
            gc.gc_value < 0.5,
            "GC for independent series should be small, got {}",
            gc.gc_value
        );
    }

    #[test]
    fn test_var_model() {
        let (x, y) = create_coupled_series(300, 0.3);
        let data = vec![x, y];

        let model = VarModel::fit(&data, 3).unwrap();

        assert_eq!(model.order, 3);
        assert_eq!(model.n_vars, 2);
        assert_eq!(model.coefficients.len(), 3);
    }

    #[test]
    fn test_spectral_granger() {
        let (x, y) = create_coupled_series(500, 0.5);

        let result = spectral_granger(&x, &y, 100.0, 5, Some(32)).unwrap();

        assert_eq!(result.frequencies.len(), 32);
        assert_eq!(result.gc_x_to_y.len(), 32);
        assert_eq!(result.gc_y_to_x.len(), 32);

        // GC values should be non-negative
        assert!(result.gc_x_to_y.iter().all(|&v| v >= 0.0));
        assert!(result.gc_y_to_x.iter().all(|&v| v >= 0.0));
    }

    #[test]
    fn test_order_selection() {
        let (x, y) = create_coupled_series(500, 0.3);
        let data = vec![x, y];

        let order = select_order(&data, 10, OrderCriterion::Bic).unwrap();

        assert!(order >= 1 && order <= 10);
    }

    #[test]
    fn test_granger_config() {
        let config = GrangerConfig::default();
        assert_eq!(config.order, 10);

        let config2 = GrangerConfig::with_order(5);
        assert_eq!(config2.order, 5);
    }
}