//! Cluster-based permutation tests for spatio-temporal data
//!
//! Implements cluster-based permutation tests as described in:
//! Maris & Oostenveld (2007) "Nonparametric statistical testing of EEG- and MEG-data"
//!
//! These tests control for multiple comparisons by forming clusters of
//! adjacent significant points and testing cluster-level statistics.

use crate::{Result, StatsError, permutation::Tail, utils};
use rand::prelude::*;
use rayon::prelude::*;
use std::collections::{HashSet, VecDeque};

/// Sparse adjacency matrix for spatial connectivity
#[derive(Debug, Clone)]
pub struct Adjacency {
    /// Number of nodes (e.g., channels or sources)
    pub n_nodes: usize,
    /// Adjacency list: for each node, list of neighbor indices
    pub neighbors: Vec<Vec<usize>>,
}

impl Adjacency {
    /// Create adjacency from neighbor lists
    pub fn new(neighbors: Vec<Vec<usize>>) -> Self {
        let n_nodes = neighbors.len();
        Self { n_nodes, neighbors }
    }

    /// Create fully connected adjacency (all nodes are neighbors)
    pub fn fully_connected(n_nodes: usize) -> Self {
        let neighbors: Vec<Vec<usize>> = (0..n_nodes)
            .map(|i| (0..n_nodes).filter(|&j| j != i).collect())
            .collect();
        Self { n_nodes, neighbors }
    }

    /// Create grid adjacency (4-connectivity for 2D grid)
    pub fn grid_2d(n_rows: usize, n_cols: usize) -> Self {
        let n_nodes = n_rows * n_cols;
        let mut neighbors = vec![Vec::new(); n_nodes];

        for row in 0..n_rows {
            for col in 0..n_cols {
                let idx = row * n_cols + col;
                // Up
                if row > 0 {
                    neighbors[idx].push((row - 1) * n_cols + col);
                }
                // Down
                if row < n_rows - 1 {
                    neighbors[idx].push((row + 1) * n_cols + col);
                }
                // Left
                if col > 0 {
                    neighbors[idx].push(row * n_cols + col - 1);
                }
                // Right
                if col < n_cols - 1 {
                    neighbors[idx].push(row * n_cols + col + 1);
                }
            }
        }

        Self { n_nodes, neighbors }
    }

    /// Create adjacency from distance matrix (connect nodes within threshold)
    pub fn from_distance_matrix(distances: &[Vec<f64>], threshold: f64) -> Self {
        let n_nodes = distances.len();
        let neighbors: Vec<Vec<usize>> = (0..n_nodes)
            .map(|i| {
                (0..n_nodes)
                    .filter(|&j| i != j && distances[i][j] <= threshold)
                    .collect()
            })
            .collect();
        Self { n_nodes, neighbors }
    }
}

/// A cluster of connected significant points
#[derive(Debug, Clone)]
pub struct Cluster {
    /// Indices of points in the cluster (flattened: node * n_times + time)
    pub indices: Vec<usize>,
    /// Cluster statistic (sum of t-values)
    pub statistic: f64,
    /// Node indices in the cluster
    pub nodes: Vec<usize>,
    /// Time indices in the cluster
    pub times: Vec<usize>,
}

/// Result of cluster-based permutation test
#[derive(Debug, Clone)]
pub struct ClusterTestResult {
    /// Observed t-statistics [n_nodes x n_times]
    pub t_obs: Vec<Vec<f64>>,
    /// Clusters found in observed data
    pub clusters: Vec<Cluster>,
    /// P-values for each cluster
    pub cluster_pvalues: Vec<f64>,
    /// Null distribution of maximum cluster statistics
    pub h0: Vec<f64>,
    /// Number of permutations
    pub n_permutations: usize,
    /// Threshold used for clustering
    pub threshold: f64,
}

/// Find clusters in a 2D array of statistics
///
/// # Arguments
/// * `stats` - 2D array of test statistics [n_nodes x n_times]
/// * `threshold` - Cluster-forming threshold (absolute value)
/// * `adjacency` - Spatial adjacency matrix (None = no spatial clustering)
/// * `tail` - Which tail to consider for clustering
///
/// # Returns
/// Vector of clusters found
pub fn find_clusters(
    stats: &[Vec<f64>],
    threshold: f64,
    adjacency: Option<&Adjacency>,
    tail: Tail,
) -> Vec<Cluster> {
    if stats.is_empty() || stats[0].is_empty() {
        return Vec::new();
    }

    let n_nodes = stats.len();
    let n_times = stats[0].len();

    // Create mask of significant points
    let mask: Vec<Vec<bool>> = stats
        .iter()
        .map(|row| {
            row.iter()
                .map(|&v| match tail {
                    Tail::TwoSided => v.abs() >= threshold,
                    Tail::Greater => v >= threshold,
                    Tail::Less => v <= -threshold,
                })
                .collect()
        })
        .collect();

    // Find connected components
    let mut visited = vec![vec![false; n_times]; n_nodes];
    let mut clusters = Vec::new();

    for node in 0..n_nodes {
        for time in 0..n_times {
            if mask[node][time] && !visited[node][time] {
                // BFS to find cluster
                let cluster = bfs_cluster(
                    &mask,
                    &mut visited,
                    node,
                    time,
                    n_nodes,
                    n_times,
                    adjacency,
                    stats,
                );
                if !cluster.indices.is_empty() {
                    clusters.push(cluster);
                }
            }
        }
    }

    // Sort clusters by absolute statistic (largest first)
    clusters.sort_by(|a, b| {
        b.statistic
            .abs()
            .partial_cmp(&a.statistic.abs())
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    clusters
}

/// BFS to find connected cluster
fn bfs_cluster(
    mask: &[Vec<bool>],
    visited: &mut [Vec<bool>],
    start_node: usize,
    start_time: usize,
    n_nodes: usize,
    n_times: usize,
    adjacency: Option<&Adjacency>,
    stats: &[Vec<f64>],
) -> Cluster {
    let mut indices = Vec::new();
    let mut nodes_set = HashSet::new();
    let mut times_set = HashSet::new();
    let mut statistic = 0.0;

    let mut queue = VecDeque::new();
    queue.push_back((start_node, start_time));
    visited[start_node][start_time] = true;

    while let Some((node, time)) = queue.pop_front() {
        indices.push(node * n_times + time);
        nodes_set.insert(node);
        times_set.insert(time);
        statistic += stats[node][time];

        // Check temporal neighbors
        if time > 0 && mask[node][time - 1] && !visited[node][time - 1] {
            visited[node][time - 1] = true;
            queue.push_back((node, time - 1));
        }
        if time < n_times - 1 && mask[node][time + 1] && !visited[node][time + 1] {
            visited[node][time + 1] = true;
            queue.push_back((node, time + 1));
        }

        // Check spatial neighbors
        if let Some(adj) = adjacency {
            for &neighbor in &adj.neighbors[node] {
                if neighbor < n_nodes && mask[neighbor][time] && !visited[neighbor][time] {
                    visited[neighbor][time] = true;
                    queue.push_back((neighbor, time));
                }
            }
        }
    }

    let mut nodes: Vec<usize> = nodes_set.into_iter().collect();
    let mut times: Vec<usize> = times_set.into_iter().collect();
    nodes.sort_unstable();
    times.sort_unstable();

    Cluster {
        indices,
        statistic,
        nodes,
        times,
    }
}

/// Spatio-temporal cluster-based permutation test
///
/// # Arguments
/// * `x` - Condition A data [n_observations x n_nodes x n_times]
/// * `y` - Condition B data (optional, for two-sample test)
/// * `threshold` - Cluster-forming threshold (t-value)
/// * `n_permutations` - Number of permutations
/// * `adjacency` - Spatial adjacency (None = temporal clustering only)
/// * `tail` - Test tail
/// * `seed` - Optional random seed
///
/// # Returns
/// ClusterTestResult with clusters and p-values
pub fn spatio_temporal_cluster_test(
    x: &[Vec<Vec<f64>>],
    y: Option<&[Vec<Vec<f64>>]>,
    threshold: f64,
    n_permutations: usize,
    adjacency: Option<&Adjacency>,
    tail: Tail,
    seed: Option<u64>,
) -> Result<ClusterTestResult> {
    if x.is_empty() {
        return Err(StatsError::InvalidInput("Empty data".to_string()));
    }

    let n_obs_x = x.len();
    let n_nodes = x[0].len();
    let n_times = x[0][0].len();

    // Validate dimensions
    for obs in x {
        if obs.len() != n_nodes || obs.iter().any(|row| row.len() != n_times) {
            return Err(StatsError::DimensionMismatch(
                "All observations must have same dimensions".to_string(),
            ));
        }
    }

    // Determine test type
    let (t_obs, _compute_stat): (
        Vec<Vec<f64>>,
        Box<dyn Fn(&[Vec<Vec<f64>>], Option<&[Vec<Vec<f64>>]>) -> Vec<Vec<f64>> + Send + Sync>,
    ) = if let Some(y_data) = y {
        // Two-sample test
        let _n_obs_y = y_data.len();
        for obs in y_data {
            if obs.len() != n_nodes || obs.iter().any(|row| row.len() != n_times) {
                return Err(StatsError::DimensionMismatch(
                    "Y dimensions must match X".to_string(),
                ));
            }
        }

        let t = compute_t_stats_ind(x, y_data, n_nodes, n_times);
        let func = Box::new(move |a: &[Vec<Vec<f64>>], b: Option<&[Vec<Vec<f64>>]>| {
            compute_t_stats_ind(a, b.unwrap(), n_nodes, n_times)
        });
        (t, func)
    } else {
        // One-sample test
        let t = compute_t_stats_1samp(x, n_nodes, n_times);
        let func = Box::new(move |a: &[Vec<Vec<f64>>], _: Option<&[Vec<Vec<f64>>]>| {
            compute_t_stats_1samp(a, n_nodes, n_times)
        });
        (t, func)
    };

    // Find clusters in observed data
    let clusters = find_clusters(&t_obs, threshold, adjacency, tail);

    // Generate null distribution
    let h0: Vec<f64> = if y.is_some() {
        // Two-sample: permute group labels
        let _n_obs_y = y.unwrap().len();
        let mut combined: Vec<Vec<Vec<f64>>> = x.to_vec();
        combined.extend(y.unwrap().iter().cloned());

        (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(),
                };

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

                let perm_x = &shuffled[..n_obs_x];
                let perm_y = &shuffled[n_obs_x..];

                let t_perm = compute_t_stats_ind(perm_x, perm_y, n_nodes, n_times);
                let perm_clusters = find_clusters(&t_perm, threshold, adjacency, tail);

                perm_clusters.first().map_or(0.0, |c| c.statistic.abs())
            })
            .collect()
    } else {
        // One-sample: sign flipping
        (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_x)
                    .map(|_| if rng.r#gen::<bool>() { 1.0 } else { -1.0 })
                    .collect();

                // Apply signs
                let flipped: Vec<Vec<Vec<f64>>> = x
                    .iter()
                    .enumerate()
                    .map(|(k, obs)| {
                        obs.iter()
                            .map(|row| row.iter().map(|&v| v * signs[k]).collect())
                            .collect()
                    })
                    .collect();

                let t_perm = compute_t_stats_1samp(&flipped, n_nodes, n_times);
                let perm_clusters = find_clusters(&t_perm, threshold, adjacency, tail);

                perm_clusters.first().map_or(0.0, |c| c.statistic.abs())
            })
            .collect()
    };

    // Compute cluster p-values
    let cluster_pvalues: Vec<f64> = clusters
        .iter()
        .map(|cluster| {
            let count = h0.iter().filter(|&&v| v >= cluster.statistic.abs()).count();
            (count as f64 + 1.0) / (n_permutations as f64 + 1.0)
        })
        .collect();

    Ok(ClusterTestResult {
        t_obs,
        clusters,
        cluster_pvalues,
        h0,
        n_permutations,
        threshold,
    })
}

/// Compute one-sample t-statistics for each node and time
fn compute_t_stats_1samp(data: &[Vec<Vec<f64>>], n_nodes: usize, n_times: usize) -> Vec<Vec<f64>> {
    let _n_obs = data.len() as f64;

    (0..n_nodes)
        .map(|node| {
            (0..n_times)
                .map(|time| {
                    let values: Vec<f64> = data.iter().map(|obs| obs[node][time]).collect();
                    let mean = utils::mean(&values);
                    let se = utils::sem(&values);
                    if se < 1e-10 { 0.0 } else { mean / se }
                })
                .collect()
        })
        .collect()
}

/// Compute independent two-sample t-statistics for each node and time
fn compute_t_stats_ind(
    x: &[Vec<Vec<f64>>],
    y: &[Vec<Vec<f64>>],
    n_nodes: usize,
    n_times: usize,
) -> Vec<Vec<f64>> {
    let n_x = x.len() as f64;
    let n_y = y.len() as f64;

    (0..n_nodes)
        .map(|node| {
            (0..n_times)
                .map(|time| {
                    let vals_x: Vec<f64> = x.iter().map(|obs| obs[node][time]).collect();
                    let vals_y: Vec<f64> = y.iter().map(|obs| obs[node][time]).collect();

                    let mean_x = utils::mean(&vals_x);
                    let mean_y = utils::mean(&vals_y);
                    let var_x = utils::variance(&vals_x, 1);
                    let var_y = utils::variance(&vals_y, 1);

                    let se = (var_x / n_x + var_y / n_y).sqrt();
                    if se < 1e-10 {
                        0.0
                    } else {
                        (mean_x - mean_y) / se
                    }
                })
                .collect()
        })
        .collect()
}

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

    #[test]
    fn test_adjacency_grid() {
        let adj = Adjacency::grid_2d(3, 3);
        assert_eq!(adj.n_nodes, 9);
        // Corner has 2 neighbors
        assert_eq!(adj.neighbors[0].len(), 2);
        // Center has 4 neighbors
        assert_eq!(adj.neighbors[4].len(), 4);
        // Edge has 3 neighbors
        assert_eq!(adj.neighbors[1].len(), 3);
    }

    #[test]
    fn test_find_clusters() {
        // Simple 2x5 array with clusters
        // Without spatial adjacency, each row forms separate temporal clusters
        let stats = vec![vec![0.5, 2.5, 2.8, 0.3, 0.1], vec![0.2, 2.6, 2.4, 0.4, 3.0]];

        let clusters = find_clusters(&stats, 2.0, None, Tail::Greater);

        // Should find clusters (row0 cols 1-2, row1 cols 1-2, row1 col 4)
        assert!(!clusters.is_empty());
        // With spatial adjacency, the two middle clusters merge
        let adj = Adjacency::fully_connected(2);
        let clusters_spatial = find_clusters(&stats, 2.0, Some(&adj), Tail::Greater);
        assert!(clusters_spatial.len() <= clusters.len());
    }

    #[test]
    fn test_cluster_test_1samp() {
        // Create data with a significant cluster
        let mut rng = StdRng::seed_from_u64(42);

        let n_obs = 20;
        let n_nodes = 4;
        let n_times = 10;

        // Generate data with effect in nodes 1-2, times 3-6
        let data: Vec<Vec<Vec<f64>>> = (0..n_obs)
            .map(|_| {
                (0..n_nodes)
                    .map(|node| {
                        (0..n_times)
                            .map(|time| {
                                let noise: f64 = rng.r#gen::<f64>() * 0.5 - 0.25;
                                if (1..=2).contains(&node) && (3..=6).contains(&time) {
                                    1.5 + noise // Effect
                                } else {
                                    noise // No effect
                                }
                            })
                            .collect()
                    })
                    .collect()
            })
            .collect();

        let adj = Adjacency::fully_connected(n_nodes);
        let result = spatio_temporal_cluster_test(
            &data,
            None,
            2.0,
            500,
            Some(&adj),
            Tail::Greater,
            Some(42),
        )
        .unwrap();

        // Should find at least one significant cluster
        assert!(!result.clusters.is_empty());
        // First cluster should be significant
        assert!(result.cluster_pvalues[0] < 0.05);
    }
}