//! Permutation tests for hypothesis testing
//!
//! Provides non-parametric permutation tests that don't assume normal distribution.

use crate::{Result, StatsError, utils};
use rand::prelude::*;
use rayon::prelude::*;

/// Tail of the test (one-sided or two-sided)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Tail {
    /// Two-sided test (H1: mean != popmean)
    TwoSided,
    /// Left-tailed test (H1: mean < popmean)
    Less,
    /// Right-tailed test (H1: mean > popmean)
    Greater,
}

/// Result of a permutation test
#[derive(Debug, Clone)]
pub struct PermutationResult {
    /// The observed test statistic
    pub statistic: f64,
    /// The p-value
    pub pvalue: f64,
    /// The null distribution (permuted statistics)
    pub null_distribution: Vec<f64>,
    /// Number of permutations performed
    pub n_permutations: usize,
}

/// One-sample permutation test
///
/// Tests whether the mean of a sample differs from a population mean.
/// Uses sign-flipping to generate the null distribution.
///
/// # Arguments
/// * `data` - Sample data
/// * `popmean` - Population mean to test against
/// * `n_permutations` - Number of permutations (default: 10000)
/// * `tail` - Type of test (two-sided, less, greater)
/// * `seed` - Optional random seed for reproducibility
///
/// # Returns
/// A `PermutationResult` with the test statistic and p-value
pub fn permutation_test_1samp(
    data: &[f64],
    popmean: f64,
    n_permutations: usize,
    tail: Tail,
    seed: Option<u64>,
) -> Result<PermutationResult> {
    if data.len() < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: data.len(),
        });
    }

    // Center the data
    let centered: Vec<f64> = data.iter().map(|x| x - popmean).collect();

    // Compute observed t-statistic
    let obs_stat = compute_t_statistic_1samp(&centered);

    // Generate null distribution using sign-flipping
    let null_dist: Vec<f64> = (0..n_permutations)
        .into_par_iter()
        .map(|i| {
            let mut rng = match seed {
                Some(s) => StdRng::seed_from_u64(s.wrapping_add(i as u64)),
                None => StdRng::from_entropy(),
            };

            // Randomly flip signs
            let flipped: Vec<f64> = centered
                .iter()
                .map(|&x| if rng.r#gen::<bool>() { x } else { -x })
                .collect();

            compute_t_statistic_1samp(&flipped)
        })
        .collect();

    // Compute p-value
    let pvalue = compute_pvalue(obs_stat, &null_dist, tail);

    Ok(PermutationResult {
        statistic: obs_stat,
        pvalue,
        null_distribution: null_dist,
        n_permutations,
    })
}

/// Paired-sample permutation test
///
/// Tests whether the mean difference between paired samples differs from zero.
///
/// # Arguments
/// * `a` - First sample
/// * `b` - Second sample (paired with a)
/// * `n_permutations` - Number of permutations
/// * `tail` - Type of test
/// * `seed` - Optional random seed
pub fn permutation_test_rel(
    a: &[f64],
    b: &[f64],
    n_permutations: usize,
    tail: Tail,
    seed: Option<u64>,
) -> Result<PermutationResult> {
    if a.len() != b.len() {
        return Err(StatsError::DimensionMismatch(format!(
            "Arrays must have same length: {} vs {}",
            a.len(),
            b.len()
        )));
    }

    if a.len() < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: a.len(),
        });
    }

    // Compute differences
    let diff: Vec<f64> = a.iter().zip(b.iter()).map(|(x, y)| x - y).collect();

    // Use one-sample test on differences
    permutation_test_1samp(&diff, 0.0, n_permutations, tail, seed)
}

/// Independent two-sample permutation test
///
/// Tests whether two independent samples have different means.
///
/// # Arguments
/// * `a` - First sample
/// * `b` - Second sample
/// * `n_permutations` - Number of permutations
/// * `tail` - Type of test
/// * `seed` - Optional random seed
pub fn permutation_test_ind(
    a: &[f64],
    b: &[f64],
    n_permutations: usize,
    tail: Tail,
    seed: Option<u64>,
) -> Result<PermutationResult> {
    if a.len() < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: a.len(),
        });
    }
    if b.len() < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: b.len(),
        });
    }

    // Combine samples
    let mut combined: Vec<f64> = a.to_vec();
    combined.extend_from_slice(b);
    let n_a = a.len();

    // Compute observed t-statistic
    let obs_stat = compute_t_statistic_ind(a, b);

    // Generate null distribution by permuting group labels
    let null_dist: Vec<f64> = (0..n_permutations)
        .into_par_iter()
        .map(|i| {
            let mut rng = match seed {
                Some(s) => StdRng::seed_from_u64(s.wrapping_add(i as u64)),
                None => StdRng::from_entropy(),
            };

            // Shuffle combined data
            let mut shuffled = combined.clone();
            shuffled.shuffle(&mut rng);

            // Split into two groups
            let perm_a = &shuffled[..n_a];
            let perm_b = &shuffled[n_a..];

            compute_t_statistic_ind(perm_a, perm_b)
        })
        .collect();

    // Compute p-value
    let pvalue = compute_pvalue(obs_stat, &null_dist, tail);

    Ok(PermutationResult {
        statistic: obs_stat,
        pvalue,
        null_distribution: null_dist,
        n_permutations,
    })
}

/// Compute one-sample t-statistic
fn compute_t_statistic_1samp(data: &[f64]) -> f64 {
    let _n = data.len() as f64;
    let mean = utils::mean(data);
    let se = utils::sem(data);

    if se < 1e-10 {
        if mean.abs() < 1e-10 {
            0.0
        } else {
            mean.signum() * f64::INFINITY
        }
    } else {
        mean / se
    }
}

/// Compute independent two-sample t-statistic (Welch's t-test)
fn compute_t_statistic_ind(a: &[f64], b: &[f64]) -> f64 {
    let mean_a = utils::mean(a);
    let mean_b = utils::mean(b);
    let var_a = utils::variance(a, 1);
    let var_b = utils::variance(b, 1);
    let n_a = a.len() as f64;
    let n_b = b.len() as f64;

    let se = (var_a / n_a + var_b / n_b).sqrt();

    if se < 1e-10 {
        let diff = mean_a - mean_b;
        if diff.abs() < 1e-10 {
            0.0
        } else {
            diff.signum() * f64::INFINITY
        }
    } else {
        (mean_a - mean_b) / se
    }
}

/// Compute p-value from null distribution
fn compute_pvalue(observed: f64, null_dist: &[f64], tail: Tail) -> f64 {
    let n = null_dist.len() as f64;

    match tail {
        Tail::TwoSided => {
            let abs_obs = observed.abs();
            let count = null_dist.iter().filter(|&&x| x.abs() >= abs_obs).count();
            (count as f64 + 1.0) / (n + 1.0)
        }
        Tail::Greater => {
            let count = null_dist.iter().filter(|&&x| x >= observed).count();
            (count as f64 + 1.0) / (n + 1.0)
        }
        Tail::Less => {
            let count = null_dist.iter().filter(|&&x| x <= observed).count();
            (count as f64 + 1.0) / (n + 1.0)
        }
    }
}

/// Multi-dimensional permutation test for arrays
///
/// Performs element-wise permutation tests on multi-dimensional data.
/// Useful for testing significance at each time point or sensor.
///
/// # Arguments
/// * `data` - 2D array [n_observations x n_features]
/// * `n_permutations` - Number of permutations
/// * `tail` - Type of test
/// * `seed` - Optional random seed
///
/// # Returns
/// Tuple of (t_statistics, p_values) for each feature
pub fn permutation_test_1samp_nd(
    data: &[Vec<f64>],
    popmean: f64,
    n_permutations: usize,
    tail: Tail,
    seed: Option<u64>,
) -> Result<(Vec<f64>, Vec<f64>)> {
    if data.is_empty() {
        return Err(StatsError::InvalidInput("Empty data array".to_string()));
    }

    let n_obs = data.len();
    let n_features = data[0].len();

    if n_obs < 2 {
        return Err(StatsError::InsufficientData {
            needed: 2,
            got: n_obs,
        });
    }

    // Check all rows have same length
    for row in data {
        if row.len() != n_features {
            return Err(StatsError::DimensionMismatch(
                "All rows must have same length".to_string(),
            ));
        }
    }

    // Center data
    let centered: Vec<Vec<f64>> = data
        .iter()
        .map(|row| row.iter().map(|x| x - popmean).collect())
        .collect();

    // Compute observed t-statistics for each feature
    let obs_stats: Vec<f64> = (0..n_features)
        .map(|j| {
            let col: Vec<f64> = centered.iter().map(|row| row[j]).collect();
            compute_t_statistic_1samp(&col)
        })
        .collect();

    // Generate null distributions
    let null_dists: Vec<Vec<f64>> = (0..n_permutations)
        .into_par_iter()
        .map(|i| {
            let mut rng = match seed {
                Some(s) => StdRng::seed_from_u64(s.wrapping_add(i as u64)),
                None => StdRng::from_entropy(),
            };

            // Generate sign flips
            let signs: Vec<f64> = (0..n_obs)
                .map(|_| if rng.r#gen::<bool>() { 1.0 } else { -1.0 })
                .collect();

            // Compute t-statistics for each feature
            (0..n_features)
                .map(|j| {
                    let col: Vec<f64> = centered
                        .iter()
                        .enumerate()
                        .map(|(k, row)| row[j] * signs[k])
                        .collect();
                    compute_t_statistic_1samp(&col)
                })
                .collect()
        })
        .collect();

    // Compute p-values for each feature
    let pvalues: Vec<f64> = (0..n_features)
        .map(|j| {
            let null_dist: Vec<f64> = null_dists.iter().map(|perm| perm[j]).collect();
            compute_pvalue(obs_stats[j], &null_dist, tail)
        })
        .collect();

    Ok((obs_stats, pvalues))
}

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

    #[test]
    fn test_permutation_1samp_significant() {
        // Data clearly above 0
        let data = vec![2.1, 2.3, 1.9, 2.5, 2.2, 2.0, 2.4, 2.1];
        let result = permutation_test_1samp(&data, 0.0, 1000, Tail::TwoSided, Some(42)).unwrap();

        assert!(result.statistic > 0.0);
        assert!(result.pvalue < 0.05);
    }

    #[test]
    fn test_permutation_1samp_not_significant() {
        // Data centered around 0
        let data = vec![0.1, -0.2, 0.15, -0.1, 0.05, -0.05, 0.1, -0.15];
        let result = permutation_test_1samp(&data, 0.0, 1000, Tail::TwoSided, Some(42)).unwrap();

        assert!(result.pvalue > 0.05);
    }

    #[test]
    fn test_permutation_paired() {
        // Paired samples with clear difference
        let a = vec![10.0, 11.0, 12.0, 13.0, 14.0];
        let b = vec![8.0, 9.0, 10.0, 11.0, 12.0];
        let result = permutation_test_rel(&a, &b, 1000, Tail::Greater, Some(42)).unwrap();

        assert!(result.statistic > 0.0);
        assert!(result.pvalue < 0.05);
    }

    #[test]
    fn test_permutation_ind() {
        // Two independent groups with different means
        let a = vec![10.0, 11.0, 12.0, 10.5, 11.5];
        let b = vec![5.0, 6.0, 5.5, 6.5, 5.8];
        let result = permutation_test_ind(&a, &b, 1000, Tail::TwoSided, Some(42)).unwrap();

        assert!(result.statistic > 0.0);
        assert!(result.pvalue < 0.05);
    }

    #[test]
    fn test_permutation_nd() {
        // Multi-dimensional data with more samples for reliable p-values
        let data = vec![
            vec![2.0, 0.1, 3.0],
            vec![2.2, -0.1, 2.8],
            vec![1.9, 0.05, 3.1],
            vec![2.1, -0.05, 2.9],
            vec![2.05, 0.02, 3.05],
            vec![1.95, -0.02, 2.95],
            vec![2.15, 0.08, 3.08],
            vec![2.08, -0.08, 2.92],
        ];
        let (stats, pvals) =
            permutation_test_1samp_nd(&data, 0.0, 1000, Tail::TwoSided, Some(42)).unwrap();

        assert_eq!(stats.len(), 3);
        assert_eq!(pvals.len(), 3);
        // First and third features should have large t-statistics
        assert!(stats[0] > 5.0);
        assert!(stats[2] > 5.0);
        // Second feature should have small t-statistic
        assert!(stats[1].abs() < 2.0);
    }
}