manifolds-rs 0.3.3

//! Module containing helper functions to generate the graphs used by UMAP,
//! tSNE and PHATE.

use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};

use thousands::*;

use crate::data::structures::*;
use crate::prelude::*;

//////////
// UMAP //
//////////

/// UMAP algorithm parameters
///
/// Controls the fuzzy simplicial set construction and graph symmetrisation.
#[derive(Clone, Debug)]
pub struct UmapGraphParams<T> {
    /// Convergence tolerance for smooth kNN distance binary search (typically
    /// 1e-5). Controls how precisely sigma values are computed.
    pub bandwidth: T,
    /// Number of nearest neighbours assumed to be at distance zero (typically
    /// 1.0). Allows for local manifold structure by treating the nearest
    /// neighbour(s) as having maximal membership strength.
    pub local_connectivity: T,
    /// Balance between fuzzy union and directed graph during symmetrisation
    /// (typically 1.0).
    pub mix_weight: T,
}

impl<T> Default for UmapGraphParams<T>
where
    T: ManifoldsFloat,
{
    /// Returns sensible defaults for UMAP
    ///
    /// ### Returns
    ///
    /// * `bandwidth = 1e-5` - Tight convergence for sigma computation
    /// * `local_connectivity = 1.0` - Treat nearest neighbour as connected
    /// * `mix_weight = 1.0` - Standard symmetric fuzzy union
    fn default() -> Self {
        Self {
            local_connectivity: T::from(1.0).unwrap(),
            bandwidth: T::from(1e-5).unwrap(),
            mix_weight: T::from(1.0).unwrap(),
        }
    }
}

/// Smooth kNN distances via binary search to find sigma for each point
///
/// For each point, finds the bandwidth (sigma) such that the sum of
/// similarities to its k nearest neighbours approximates log(k). Uses binary
/// search for efficiency.
///
/// ### Params
///
/// * `dist` - kNN distance matrix where each row contains distances to k
///   nearest neighbours
/// * `k` - Number of nearest neighbours (used to compute target = ln(k))
/// * `local_connectivity` - Number of nearest neighbours to assume are at
///   distance zero (typically 1.0). Allows for local manifold structure.
/// * `bandwidth` - Convergence tolerance for binary search (typically 1e-5)
/// * `n_iter` - Maximum number of binary search iterations (typically 64)
///
/// ### Returns
///
/// * `sigmas` - Smoothing bandwidth for each point
/// * `rhos` - Distance to the `local_connectivity`-th nearest neighbour for
///   each point
///
/// ### Notes
///
/// Used for UMAP
pub fn smooth_knn_dist<T>(
    dist: &[Vec<T>],
    k: usize,
    local_connectivity: T,
    bandwidth: T,
    n_iter: usize,
) -> (Vec<T>, Vec<T>)
where
    T: ManifoldsFloat,
{
    dist.par_iter()
        .map(|dists| {
            let target = (k as f64).ln();

            let rho = if local_connectivity > T::zero() {
                let idx = (local_connectivity - T::one())
                    .max(T::zero())
                    .floor()
                    .to_usize()
                    .unwrap()
                    .min(dists.len() - 1);

                let fraction = (local_connectivity - T::one()).max(T::zero())
                    - (local_connectivity - T::one()).max(T::zero()).floor();

                if fraction > T::zero() && idx + 1 < dists.len() {
                    dists[idx] * (T::one() - fraction) + dists[idx + 1] * fraction
                } else {
                    dists[idx]
                }
            } else {
                T::zero()
            };

            // Binary search for sigma (rest unchanged)
            let mut lo = T::zero();
            let mut hi = T::max_value();
            let mut mid = T::one();

            for _ in 0..n_iter {
                let mut val = T::zero();
                for &d in dists.iter() {
                    let adjusted = (d - rho).max(T::zero());
                    val += (-(adjusted / mid)).exp();
                }

                if (val.to_f64().unwrap() - target).abs() < bandwidth.to_f64().unwrap() {
                    break;
                }

                if val.to_f64().unwrap() > target {
                    hi = mid;
                    mid = (lo + hi) / (T::one() + T::one());
                } else {
                    lo = mid;
                    if hi == T::max_value() {
                        mid *= T::one() + T::one();
                    } else {
                        mid = (lo + hi) / (T::one() + T::one());
                    }
                }
            }

            (mid, rho)
        })
        .unzip()
}

/// Convert kNN graph to sparse COO (Coordinate) format with membership
/// strengths
///
/// Computes fuzzy simplicial set membership strengths based on distances,
/// local connectivity (rho), and smoothed bandwidths (sigma).
///
/// ### Params
///
/// * `knn_indices` - Indices of k nearest neighbours for each point
/// * `knn_dists` - Distances to k nearest neighbours for each point
/// * `sigmas` - Smoothing bandwidth for each point (from `smooth_knn_dist`)
/// * `rhos` - Local connectivity distance for each point (from
///   `smooth_knn_dist`)
///
/// ### Returns
///
/// Sparse graph in COO format where weights represent membership strengths
/// computed as exp(-(max(0, dist - rho) / sigma))
pub fn knn_to_coo<T>(
    knn_indices: &[Vec<usize>],
    knn_dists: &[Vec<T>],
    sigmas: &[T],
    rhos: &[T],
) -> CoordinateList<T>
where
    T: ManifoldsFloat,
{
    let n = knn_indices.len();
    let capacity: usize = knn_indices.iter().map(|v| v.len()).sum();

    let mut row_indices = Vec::with_capacity(capacity);
    let mut col_indices = Vec::with_capacity(capacity);
    let mut values = Vec::with_capacity(capacity);

    for (i, (neighbours, dists)) in knn_indices.iter().zip(knn_dists.iter()).enumerate() {
        let sigma = sigmas[i];
        let rho = rhos[i];

        for (&j, &dist) in neighbours.iter().zip(dists.iter()) {
            if i == j {
                continue;
            }

            let adjusted = (dist - rho).max(T::zero());
            let weight = if sigma > T::zero() {
                (-(adjusted / sigma)).exp()
            } else if adjusted > T::zero() {
                T::zero()
            } else {
                T::one()
            };

            row_indices.push(i);
            col_indices.push(j);
            values.push(weight);
        }
    }

    CoordinateList {
        row_indices,
        col_indices,
        values,
        n_samples: n,
    }
}

/// Symmetrise graph using probabilistic t-conorm (fuzzy set union)
///
/// Creates symmetric graph by combining directed edges using fuzzy union:
/// w_sym = w_ij + w_ji - w_ij * w_ji, weighted by `mix_weight`.
///
/// ### Params
///
/// * `graph` - Input directed graph in COO format
/// * `mix_weight` - Balance between fuzzy union (0.5) and directed graph (1.0).
///   Controls how much to weight the union operation.
///
/// ### Returns
///
/// Symmetrised graph in COO format
///
/// ### Notes
///
/// * `mix_weight = 1.0`: Full fuzzy union (standard UMAP, symmetric)
/// * `mix_weight = 0.5`: Weighted average of union and directed)
/// * `mix_weight = 0.0`: Use only outgoing edges (directed)
pub fn symmetrise_graph<T>(graph: CoordinateList<T>, mix_weight: T) -> CoordinateList<T>
where
    T: ManifoldsFloat,
{
    let n = graph.n_samples;

    // Build adjacency maps for fast lookups
    let mut forward: Vec<FxHashMap<usize, T>> = vec![FxHashMap::default(); n];
    let mut backward: Vec<FxHashMap<usize, T>> = vec![FxHashMap::default(); n];

    for ((&i, &j), &w) in graph
        .row_indices
        .iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
    {
        forward[i].insert(j, w);
        backward[j].insert(i, w);
    }

    // Parallel symmetrisation
    let edges: Vec<Vec<(usize, T)>> = (0..n)
        .into_par_iter()
        .map(|i| {
            let mut combined = FxHashMap::default();

            // Process all unique neighbours
            for &j in forward[i].keys().chain(backward[i].keys()) {
                let w_ij = forward[i].get(&j).copied().unwrap_or(T::zero());
                let w_ji = backward[i].get(&j).copied().unwrap_or(T::zero());

                // Fuzzy union: a + b - a*b, weighted by mix_weight
                let union = w_ij + w_ji - w_ij * w_ji;
                let w_sym = mix_weight * union + (T::one() - mix_weight) * w_ij;

                if w_sym > T::zero() {
                    combined.insert(j, w_sym);
                }
            }

            let mut result: Vec<(usize, T)> = combined.into_iter().collect();
            result.sort_unstable_by_key(|&(idx, _)| idx);
            result
        })
        .collect();

    // Flatten back to COO
    let capacity: usize = edges.iter().map(|v| v.len()).sum();
    let mut row_indices = Vec::with_capacity(capacity);
    let mut col_indices = Vec::with_capacity(capacity);
    let mut values = Vec::with_capacity(capacity);

    for (i, neighbours) in edges.into_iter().enumerate() {
        for (j, w) in neighbours {
            row_indices.push(i);
            col_indices.push(j);
            values.push(w);
        }
    }

    CoordinateList {
        row_indices,
        col_indices,
        values,
        n_samples: n,
    }
}

/// Convert COO sparse graph to adjacency list representation
///
/// More efficient for SGD optimisation where we need to iterate over neighbours
/// of each vertex.
///
/// ### Params
///
/// * `graph` - Sparse graph in COO format
///
/// ### Returns
///
/// Adjacency list where `result[i]` contains `(neighbour_index, edge_weight)`
/// pairs for vertex `i`
pub fn coo_to_adjacency_list<T>(graph: &CoordinateList<T>) -> Vec<Vec<(usize, T)>>
where
    T: ManifoldsFloat,
{
    let mut adj = vec![Vec::new(); graph.n_samples];

    for ((&i, &j), &w) in graph
        .row_indices
        .iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
    {
        adj[i].push((j, w));
    }

    adj
}

/// Filter out edges that are too weak to be sampled during optimization
///
/// Removes edges where weight < max_weight / n_epochs, matching uwot's
/// preprocessing step. These weak edges would never be sampled during
/// optimization and can cause fragmentation.
///
/// ### Params
///
/// * `graph` - Input graph in COO format
/// * `n_epochs` - Optimization parameters (uses n_epochs for threshold)
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// Filtered graph with weak edges removed
pub fn filter_weak_edges<T>(
    graph: CoordinateList<T>,
    n_epochs: usize,
    verbose: usize,
) -> CoordinateList<T>
where
    T: ManifoldsFloat,
{
    let verbosity = parse_verbosity_level(verbose);

    let max_weight = graph
        .values
        .iter()
        .copied()
        .fold(T::zero(), |acc, w| if w > acc { w } else { acc });

    let original_edge_no = graph.col_indices.len();

    let threshold = max_weight / T::from(n_epochs).unwrap();

    let mut filtered_rows = Vec::new();
    let mut filtered_cols = Vec::new();
    let mut filtered_vals = Vec::new();

    for ((&i, &j), &w) in graph
        .row_indices
        .iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
    {
        if w >= threshold {
            filtered_rows.push(i);
            filtered_cols.push(j);
            filtered_vals.push(w);
        }
    }

    let filtered_edge_no = filtered_cols.len();

    if verbosity.detailed_verbosity() {
        println!(
            " Filtered out {} weak edges.",
            (original_edge_no - filtered_edge_no).separate_with_underscores(),
        );
    }

    CoordinateList {
        row_indices: filtered_rows,
        col_indices: filtered_cols,
        values: filtered_vals,
        n_samples: graph.n_samples,
    }
}

//////////
// tSNE //
//////////

/// Compute Gaussian affinities from k-nearest neighbours using perplexity-based
/// calibration
///
/// For each point i, computes conditional probabilities p_{j|i} using a
/// Gaussian kernel with bandwidth calibrated via binary search to achieve a
/// target perplexity. The result is a sparse graph where edge (i,j) has weight
/// p_{j|i}.
///
/// ### Params
///
/// * `knn_indices` - For each point, indices of its k nearest neighbours
/// * `knn_dists` - For each point, distances to its k nearest neighbours
///   (same order as indices!)
/// * `perplexity` - Target perplexity (effective number of neighbours). Typical
///   values: 5-50
/// * `tol` - Convergence tolerance for entropy (typical: 1e-5)
/// * `max_iter` - Maximum iterations for binary search (typical: 50-200)
/// * `distances_squared` - If true, distances are already squared (e.g.,
///   squared Euclidean). If false, distances will be squared before computing
///   the kernel.
///
/// ### Returns
///
/// A `CoordinateList` containing the asymmetric conditional probabilities p_{j|i}
///
/// ### Notes
///
/// Used for tSNE
/// Compute Gaussian affinities from k-nearest neighbours using perplexity-based
/// calibration
///
/// For each point i, computes conditional probabilities p_{j|i} using a
/// Gaussian kernel with bandwidth calibrated via binary search to achieve a
/// target perplexity. The result is a sparse graph where edge (i,j) has weight
/// p_{j|i}.
///
/// ### Params
///
/// * `knn_indices` - For each point, indices of its k nearest neighbours
/// * `knn_dists` - For each point, distances to its k nearest neighbours
///   (same order as indices!)
/// * `perplexity` - Target perplexity (effective number of neighbours). Typical
///   values: 5-50
/// * `tol` - Convergence tolerance for entropy (typical: 1e-5)
/// * `max_iter` - Maximum iterations for binary search (typical: 50-200)
/// * `distances_squared` - If true, distances are already squared (e.g.,
///   squared Euclidean). If false, distances will be squared before computing
///   the kernel.
///
/// ### Returns
///
/// A `CoordinateList` containing the asymmetric conditional probabilities p_{j|i}
///
/// ### Errors
///
/// Returns an error if `perplexity >= k` for any point, since the maximum
/// achievable entropy with k neighbours is log2(k) and the binary search
/// cannot converge to a target entropy of log2(perplexity).
///
/// ### Notes
///
/// Used for tSNE. Emits a warning when perplexity exceeds (k - 1) / 3, as the
/// binary search still converges but the bandwidth becomes wide enough that
/// tail neighbours dominate the affinities, degrading locality.
pub fn gaussian_knn_affinities<T>(
    knn_indices: &[Vec<usize>],
    knn_dists: &[Vec<T>],
    perplexity: T,
    tol: T,
    max_iter: usize,
    distances_squared: bool,
) -> Result<CoordinateList<T>, ManifoldsError>
where
    T: ManifoldsFloat,
{
    let n = knn_indices.len();

    // perplexity vs kNN size validation
    let min_k = knn_indices.iter().map(|idx| idx.len()).min().unwrap_or(0);

    let perp_f64 = perplexity.to_f64().unwrap_or(f64::NAN);

    // hard error: perplexity must be strictly less than k
    if perp_f64 >= min_k as f64 {
        return Err(ManifoldsError::PerplexityTooLarge {
            perplexity: perp_f64,
            k: min_k,
        });
    }

    let target_entropy = perplexity.log2();
    let machine_epsilon = T::epsilon();

    let results: Vec<Vec<T>> = knn_indices
        .par_iter()
        .zip(knn_dists.par_iter())
        .map(|(_, dists)| {
            // binary search for precision (beta = 1 / (2*sigma^2))
            let mut beta = T::one();
            let mut min_beta = T::neg_infinity();
            let mut max_beta = T::infinity();
            let mut current_probs = vec![T::zero(); dists.len()];

            for _ in 0..max_iter {
                // compute P_i with current beta: p_{j|i} = exp(-beta * d²)
                let mut sum_p = T::zero();
                for (j, &d) in dists.iter().enumerate() {
                    if d < T::epsilon() {
                        continue;
                    }
                    let d_sq = if distances_squared { d } else { d * d };
                    let p = (-beta * d_sq).exp();
                    current_probs[j] = p;
                    sum_p += p;
                }

                // check for numerical stability
                if sum_p.abs() < machine_epsilon {
                    sum_p = machine_epsilon;
                }

                // normalise to get probabilities and compute entropy H
                let mut entropy = T::zero();
                for p in current_probs.iter_mut() {
                    *p /= sum_p;
                    if *p > machine_epsilon {
                        entropy -= *p * p.log2();
                    }
                }

                // check convergence
                let entropy_diff = entropy - target_entropy;
                if entropy_diff.abs() < tol {
                    break;
                }

                // adjust beta
                if entropy_diff > T::zero() {
                    // entropy too high → distribution too flat → increase beta (narrow curve)
                    min_beta = beta;
                    if max_beta.is_infinite() {
                        beta *= T::one() + T::one();
                    } else {
                        beta = (beta + max_beta) / (T::one() + T::one());
                    }
                } else {
                    // entropy too low → distribution too peaked → decrease beta (widen curve)
                    max_beta = beta;
                    if min_beta.is_infinite() {
                        beta /= T::one() + T::one();
                    } else {
                        beta = (beta + min_beta) / (T::one() + T::one());
                    }
                }
            }

            current_probs
        })
        .collect();

    // build sparse graph
    let capacity: usize = results.iter().map(|p| p.len()).sum();
    let mut row_indices = Vec::with_capacity(capacity);
    let mut col_indices = Vec::with_capacity(capacity);
    let mut values = Vec::with_capacity(capacity);

    for (i, probs) in results.into_iter().enumerate() {
        for (&j, p) in knn_indices[i].iter().zip(probs) {
            if p > machine_epsilon && j != i {
                row_indices.push(i);
                col_indices.push(j);
                values.push(p);
            }
        }
    }

    Ok(CoordinateList {
        row_indices,
        col_indices,
        values,
        n_samples: n,
    })
}

/// Symmetrise graph for t-SNE: P_sym = (P + P^T) / 2N
///
/// Converts conditional probabilities P(j|i) to symmetric joint probabilities
/// P_ij. This ensures P_ij = P_ji and Σ_ij P_ij = 1.
///
/// ### Params
///
/// * `graph` - Directed sparse graph containing conditional probabilities P(j|i)
///
/// ### Returns
///
/// Symmetric `CoordinateList` where:
/// - Each edge (i,j) has weight P_ij = (P(j|i) + P(i|j)) / 2N
/// - P_ij = P_ji (symmetric)
/// - All weights sum to 1.0
///
/// ### Algorithm
///
/// 1. Build adjacency map for fast lookup of P(j|i)
/// 2. Collect all unique unordered pairs {i,j} from input edges
/// 3. For each pair: compute P_ij = (P(j|i) + P(i|j)) / 2N
/// 4. Add both directions (i,j) and (j,i) to output with weight P_ij
pub fn symmetrise_affinities_tsne<T>(graph: CoordinateList<T>) -> CoordinateList<T>
where
    T: ManifoldsFloat,
{
    let n = graph.n_samples;
    let n_float = T::from_usize(n).unwrap();
    let two = T::from_f64(2.0).unwrap();
    let normalization = two * n_float;

    // Build adjacency map for O(1) lookup
    let mut adj: Vec<FxHashMap<usize, T>> = vec![FxHashMap::default(); n];
    for ((&i, &j), &w) in graph
        .row_indices
        .iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
    {
        adj[i].insert(j, w);
    }

    // Collect all unique unordered pairs {i,j} from edges
    let mut pairs_set: FxHashSet<(usize, usize)> = FxHashSet::default();
    for ((&i, &j), _) in graph
        .row_indices
        .iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
    {
        let pair = if i < j {
            (i, j)
        } else if i > j {
            (j, i)
        } else {
            (i, i)
        };
        pairs_set.insert(pair);
    }

    let pairs: Vec<(usize, usize)> = pairs_set.into_iter().collect();

    // Process each unique pair in parallel
    let edges: Vec<Vec<(usize, usize, T)>> = pairs
        .par_iter()
        .map(|&(i, j)| {
            // Get both directions (may not exist)
            let w_ij = adj[i].get(&j).copied().unwrap_or_else(T::zero);
            let w_ji = adj[j].get(&i).copied().unwrap_or_else(T::zero);
            let p_sym = (w_ij + w_ji) / normalization;

            let mut local_edges = Vec::new();
            if p_sym > T::epsilon() {
                local_edges.push((i, j, p_sym));
                if i != j {
                    local_edges.push((j, i, p_sym));
                }
            }
            local_edges
        })
        .collect();

    // Flatten into final arrays
    let total_capacity: usize = edges.iter().map(|v| v.len()).sum();
    let mut rows = Vec::with_capacity(total_capacity);
    let mut cols = Vec::with_capacity(total_capacity);
    let mut vals = Vec::with_capacity(total_capacity);

    for edge_vec in edges {
        for (i, j, w) in edge_vec {
            rows.push(i);
            cols.push(j);
            vals.push(w);
        }
    }

    CoordinateList {
        row_indices: rows,
        col_indices: cols,
        values: vals,
        n_samples: n,
    }
}

///////////
// PHATE //
///////////

///////////
// Enums //
///////////

/// Which symmetrisation to use for the PHATE graph.
#[derive(Default)]
pub enum PhateGraphSymmetrisation {
    /// Additive symmetrisation - used in PHATE
    #[default]
    Additive,
    /// Multiplicative symmetrisation
    Multiplicative,
    /// Min-max symmetrisation
    Mnn,
    /// No symmetrisation
    None,
}

/// Parse a string into a PhateGraphSymmetrisation enum
///
/// ### Params
///
/// * `s` - String to parse
///
/// ### Returns
///
/// `Some(PhateGraphSymmetrisation)` pending on parsing.
pub fn parse_phate_symmetrisation(s: &str) -> Option<PhateGraphSymmetrisation> {
    match s.to_lowercase().as_str() {
        "additive" | "add" => Some(PhateGraphSymmetrisation::Additive),
        "multiplicative" | "mult" | "multiply" => Some(PhateGraphSymmetrisation::Multiplicative),
        "mnn" => Some(PhateGraphSymmetrisation::Mnn),
        "none" => Some(PhateGraphSymmetrisation::None),
        _ => None,
    }
}

/////////////
// Helpers //
/////////////

/// Additive symmetrisation
///
/// K = (K + K^T) / 2
///
/// ### Params
///
/// * `graph` - Reference to the graph to symmetrise
fn symmetrise_additive<T>(graph: &mut CoordinateList<T>)
where
    T: ManifoldsFloat,
{
    let two = T::one() + T::one();

    // do this in parallel
    // maybe slower for smaller graphs, but for sure much faster for larger ones
    let edge_map: FxHashMap<(usize, usize), T> = graph
        .row_indices
        .par_iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
        .fold(
            FxHashMap::default, // Changed here
            |mut local_map, ((&i, &j), &v)| {
                *local_map.entry((i, j)).or_insert(T::zero()) += v;
                *local_map.entry((j, i)).or_insert(T::zero()) += v;
                local_map
            },
        )
        .reduce(
            FxHashMap::default, // And here
            |mut map1, map2| {
                for (key, val) in map2 {
                    *map1.entry(key).or_insert(T::zero()) += val;
                }
                map1
            },
        );

    // Rebuild graph with (K + K^T) / 2
    graph.row_indices.clear();
    graph.col_indices.clear();
    graph.values.clear();

    graph.row_indices.reserve(edge_map.len());
    graph.col_indices.reserve(edge_map.len());
    graph.values.reserve(edge_map.len());

    for ((i, j), v) in edge_map {
        graph.row_indices.push(i);
        graph.col_indices.push(j);
        graph.values.push(v / two);
    }
}

/// Multiplicative symmetrisation
///
/// K = K ⊙ K^T (element-wise product)
///
/// ### Params
///
/// * `graph` - Reference to the graph to symmetrise
fn symmetrise_multiplicative<T>(graph: &mut CoordinateList<T>)
where
    T: ManifoldsFloat,
{
    // forward map
    let forward_map: FxHashMap<(usize, usize), T> = graph
        .row_indices
        .par_iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
        .fold(FxHashMap::default, |mut map, ((&i, &j), &v)| {
            map.insert((i, j), v);
            map
        })
        .reduce(FxHashMap::default, |mut map1, map2| {
            map1.extend(map2);
            map1
        });

    // backward map
    let backward_map: FxHashMap<(usize, usize), T> = graph
        .row_indices
        .par_iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
        .fold(FxHashMap::default, |mut map, ((&i, &j), &v)| {
            map.insert((j, i), v);
            map
        })
        .reduce(FxHashMap::default, |mut map1, map2| {
            map1.extend(map2);
            map1
        });

    // compute element-wise product (parallel)
    let products: Vec<(usize, usize, T)> = forward_map
        .par_iter()
        .filter_map(|(&(i, j), &v_ij)| backward_map.get(&(i, j)).map(|&v_ji| (i, j, v_ij * v_ji)))
        .collect();

    // rebuild graph
    graph.row_indices.clear();
    graph.col_indices.clear();
    graph.values.clear();

    graph.row_indices.reserve(products.len());
    graph.col_indices.reserve(products.len());
    graph.values.reserve(products.len());

    for (i, j, v) in products {
        graph.row_indices.push(i);
        graph.col_indices.push(j);
        graph.values.push(v);
    }
}

/// MNN symmetrisation
///
/// K = θ * min(K, K^T) + (1-θ) * max(K, K^T)
///
/// ### Params
///
/// * `graph` - Reference to the graph to symmetrise
fn symmetrise_mnn<T>(graph: &mut CoordinateList<T>, theta: T)
where
    T: ManifoldsFloat,
{
    let edge_map: FxHashMap<(usize, usize), (Option<T>, Option<T>)> = graph
        .row_indices
        .par_iter()
        .zip(&graph.col_indices)
        .zip(&graph.values)
        .fold(FxHashMap::default, |mut map, ((&i, &j), &v)| {
            map.entry((i, j)).or_insert((None, None)).0 = Some(v);
            map.entry((j, i)).or_insert((None, None)).1 = Some(v);
            map
        })
        .reduce(FxHashMap::default, |mut map1, map2| {
            for (key, (v_ij, v_ji)) in map2 {
                let entry = map1.entry(key).or_insert((None, None));
                if v_ij.is_some() {
                    entry.0 = v_ij;
                }
                if v_ji.is_some() {
                    entry.1 = v_ji;
                }
            }
            map1
        });

    let one_minus_theta = T::one() - theta;

    // compute weighted min-max (parallel)
    let results: Vec<(usize, usize, T)> = edge_map
        .par_iter()
        .filter_map(|(&(i, j), &(v_ij, v_ji))| {
            let v_ij = v_ij.unwrap_or(T::zero());
            let v_ji = v_ji.unwrap_or(T::zero());
            let min_val = v_ij.min(v_ji);
            let max_val = v_ij.max(v_ji);
            let combined = theta * min_val + one_minus_theta * max_val;

            if combined > T::epsilon() {
                Some((i, j, combined))
            } else {
                None
            }
        })
        .collect();

    // Rebuild graph
    graph.row_indices.clear();
    graph.col_indices.clear();
    graph.values.clear();

    graph.row_indices.reserve(results.len());
    graph.col_indices.reserve(results.len());
    graph.values.reserve(results.len());

    for (i, j, v) in results {
        graph.row_indices.push(i);
        graph.col_indices.push(j);
        graph.values.push(v);
    }
}

/// Binary connectivity
///
/// Used for decay = None case
///
/// ### Params
///
/// * `knn_indices`: The indices of the k-nearest neighbors for each vertex.
/// * `knn`: The number of nearest neighbors to consider.
/// * `symmetrise`: The symmetrisation method to use.
///
/// ### Returns
///
/// The CoordinateList representing the binary connectivity.
fn binary_knn_connectivity<T>(
    knn_indices: &[Vec<usize>],
    knn: usize,
    symmetrise: PhateGraphSymmetrisation,
) -> CoordinateList<T>
where
    T: ManifoldsFloat,
{
    let n = knn_indices.len();
    let k_actual = knn.min(knn_indices[0].len());

    let mut row_indices = Vec::new();
    let mut col_indices = Vec::new();
    let mut values = Vec::new();

    for (i, indices) in knn_indices.iter().enumerate() {
        for &j in indices.iter().take(k_actual) {
            if j != i {
                row_indices.push(i);
                col_indices.push(j);
                values.push(T::one());
            }
        }
    }

    let mut graph = CoordinateList {
        row_indices,
        col_indices,
        values,
        n_samples: n,
    };

    if !matches!(symmetrise, PhateGraphSymmetrisation::None) {
        symmetrise_additive(&mut graph);
    }

    graph
}

////////////////////////
// Alpha decay kernel //
////////////////////////

/// Compute alpha-decay affinities from k-nearest neighbours for PHATE
///
/// For each point i, computes affinities using an adaptive Gaussian kernel:
///
/// `K(i,j) = exp(-(d(i,j) / σ_i)^α)`
///
/// where σ_i is the distance to the kth nearest neighbour.
///
/// ### Params
///
/// * `knn_indices` - kNN indices (including self)
/// * `knn_dists` - kNN distances (including self)
/// * `knn` - Which neighbour to use for bandwidth (e.g., 5 means use 5th
///   nearest neighbour distance)
/// * `decay` - Decay exponent alpha (typical: 40). If None, returns binary
///   connectivity
/// * `bandwidth_scale` - Multiplicative factor for bandwidth (default: 1.0)
/// * `thresh` - Threshold below which affinities are set to 0 (default: 1e-4,
///   for sparsity)
/// * `distances_squared` - If true, distances are already squared (squared
///   Euclidean). If false, use as-is (cosine, etc.)
/// * `symmetrise` - symmetrisation method: "add" for (K+K^T)/2, "multiply" for
///   K*K^T, "none" for asymmetric.
///
/// ### Returns
///
/// A `CoordinateList` containing the (optionally symmetrized) affinities
#[allow(clippy::too_many_arguments)]
pub fn phate_alpha_decay_affinities<T>(
    knn_indices: &[Vec<usize>],
    knn_dists: &[Vec<T>],
    knn: usize,
    decay: Option<T>,
    bandwidth_scale: T,
    thresh: T,
    symmetrise: &str,
    distances_squared: bool,
) -> CoordinateList<T>
where
    T: ManifoldsFloat,
{
    let n = knn_indices.len();
    let machine_epsilon = T::epsilon();

    let symmetrise = parse_phate_symmetrisation(symmetrise).unwrap_or_default();

    // handle binary connectivity case (decay = None)
    if decay.is_none() {
        return binary_knn_connectivity(knn_indices, knn, symmetrise);
    }

    let decay_val = decay.unwrap();

    // parallel computations of affinities
    let results: Vec<(Vec<usize>, Vec<T>)> = knn_indices
        .par_iter()
        .zip(knn_dists.par_iter())
        .enumerate()
        .map(|(i, (indices, dists))| {
            // bandwidth: dist to kth nearest neighbour
            // note: indices[0] is self, so indices[knn-1] is the kth neighbour
            // (excluding self)
            let bandwidth_dist = if knn > 0 && knn <= dists.len() {
                // this is needed as the ANN liibraries return squared distances
                // for speed
                if distances_squared {
                    dists[knn - 1].sqrt() // convert squared distance to distance
                } else {
                    dists[knn - 1] // already a distance
                }
            } else {
                // fallback: use last neighbour
                if distances_squared {
                    dists[dists.len() - 1].sqrt()
                } else {
                    dists[dists.len() - 1]
                }
            };

            let bandwidth = bandwidth_dist * bandwidth_scale;

            // handle edge case of zero bandwidth
            let bandwidth = bandwidth.max(machine_epsilon);

            // pre-allocate
            let mut neighbor_indices = Vec::with_capacity(indices.len());
            let mut neighbor_values = Vec::with_capacity(indices.len());

            // compute affinities for each neighbour
            for (&j, &dist_val) in indices.iter().zip(dists.iter()) {
                // skip self-loops
                if j == i {
                    continue;
                }

                // handle zero distances
                if dist_val < machine_epsilon {
                    neighbor_indices.push(j);
                    neighbor_values.push(T::one());
                    continue;
                }

                // convert to actual distance if needed
                let d = if distances_squared {
                    dist_val.sqrt() // convert squared distance to distance
                } else {
                    dist_val // already a distance
                };

                // compute affinity: exp(-(d / σ)^α)
                let scaled = d / bandwidth;
                let powered = scaled.powf(decay_val);
                let affinity = (-powered).exp();

                // apply threshold for sparsity
                if affinity >= thresh {
                    neighbor_indices.push(j);
                    neighbor_values.push(affinity);
                }
            }

            (neighbor_indices, neighbor_values)
        })
        .collect();

    // Build asymmetric sparse graph
    let capacity: usize = results.iter().map(|(idx, _)| idx.len()).sum();
    let mut row_indices = Vec::with_capacity(capacity);
    let mut col_indices = Vec::with_capacity(capacity);
    let mut values = Vec::with_capacity(capacity);

    for (i, (indices, vals)) in results.into_iter().enumerate() {
        for (&j, v) in indices.iter().zip(vals) {
            row_indices.push(i);
            col_indices.push(j);
            values.push(v);
        }
    }

    let mut graph = CoordinateList {
        row_indices,
        col_indices,
        values,
        n_samples: n,
    };

    match symmetrise {
        PhateGraphSymmetrisation::Additive => symmetrise_additive(&mut graph),
        PhateGraphSymmetrisation::Multiplicative => symmetrise_multiplicative(&mut graph),
        PhateGraphSymmetrisation::Mnn => symmetrise_mnn(&mut graph, T::one()),
        PhateGraphSymmetrisation::None => {}
    };

    graph
}

///////////
// Tests //
///////////

#[cfg(test)]
mod test_data_gen {
    use super::*;
    use approx::assert_relative_eq;
    use num_traits::Float;

    ////////////////
    // Umap stuff //
    ////////////////

    #[test]
    fn test_smooth_knn_dist_basic() {
        // Simple test with 3 points, k=2
        let dist = vec![vec![1.0, 2.0], vec![1.5, 3.0], vec![0.5, 1.5]];

        let (sigmas, rhos) = smooth_knn_dist(&dist, 2, 1.0, 1e-5, 64);

        assert_eq!(sigmas.len(), 3);
        assert_eq!(rhos.len(), 3);

        // Rhos should be approximately the distance to the first neighbour
        assert_relative_eq!(rhos[0], 1.0, epsilon = 1e-4);
        assert_relative_eq!(rhos[1], 1.5, epsilon = 1e-4);
        assert_relative_eq!(rhos[2], 0.5, epsilon = 1e-4);

        // Sigmas should be positive
        for sigma in sigmas.iter() {
            assert!(*sigma > 0.0);
        }
    }

    #[test]
    fn test_smooth_knn_dist_zero_local_connectivity() {
        let dist = vec![vec![0.0, 1.0, 2.0], vec![0.0, 1.0, 2.0]];

        let (sigmas, rhos) = smooth_knn_dist(&dist, 2, 0.0, 1e-5, 64);

        // With zero local connectivity, rhos should all be zero
        assert!(rhos.iter().all(|&r| r == 0.0));
        assert_eq!(sigmas.len(), 2);
    }

    #[test]
    fn test_knn_to_coo_basic() {
        let knn_indices = vec![vec![1, 2], vec![0, 2], vec![0, 1]];
        let knn_dists = vec![vec![1.0, 2.0], vec![1.0, 1.5], vec![2.0, 1.5]];
        let sigmas = vec![1.0, 1.0, 1.0];
        let rhos = vec![0.0, 0.0, 0.0];

        let graph = knn_to_coo(&knn_indices, &knn_dists, &sigmas, &rhos);

        assert_eq!(graph.n_samples, 3);
        assert_eq!(graph.row_indices.len(), 6); // 3 points × 2 neighbours
        assert_eq!(graph.col_indices.len(), 6);
        assert_eq!(graph.values.len(), 6);

        // All weights should be between 0 and 1
        for &w in &graph.values {
            assert!((0.0..=1.0).contains(&w));
        }
    }

    #[test]
    fn test_knn_to_coo_self_loop_excluded() {
        // Include self in neighbours
        let knn_indices = vec![vec![0, 1], vec![1, 0]];
        let knn_dists = vec![vec![0.0, 1.0], vec![0.0, 1.0]];
        let sigmas = vec![1.0, 1.0];
        let rhos = vec![0.0, 0.0];

        let graph = knn_to_coo(&knn_indices, &knn_dists, &sigmas, &rhos);

        // Self-loops should be excluded
        assert_eq!(graph.values.len(), 2); // Only 2 edges, not 4
        assert!(graph
            .row_indices
            .iter()
            .zip(&graph.col_indices)
            .all(|(&i, &j)| i != j));
    }

    #[test]
    fn test_symmetrise_graph_full_union() {
        let graph = CoordinateList {
            row_indices: vec![0, 1],
            col_indices: vec![1, 0],
            values: vec![0.8, 0.6],
            n_samples: 2,
        };

        let sym_graph = symmetrise_graph(graph, 0.5);

        assert_eq!(sym_graph.n_samples, 2);

        // With mix_weight = 0.5:
        // union = 0.8 + 0.6 - 0.8*0.6 = 0.92
        // w_sym = 0.5 * union + 0.5 * w_ij
        // For 0->1: 0.5 * 0.92 + 0.5 * 0.8 = 0.86
        // For 1->0: 0.5 * 0.92 + 0.5 * 0.6 = 0.76

        let edges = sym_graph.to_edge_list();
        assert_eq!(edges.len(), 2);

        let edge_01 = edges.iter().find(|&&(i, j, _)| i == 0 && j == 1).unwrap();
        let edge_10 = edges.iter().find(|&&(i, j, _)| i == 1 && j == 0).unwrap();

        assert_relative_eq!(edge_01.2, 0.86, epsilon = 1e-6);
        assert_relative_eq!(edge_10.2, 0.76, epsilon = 1e-6);
    }

    #[test]
    fn test_symmetrise_graph_directed() {
        let graph = CoordinateList {
            row_indices: vec![0, 1],
            col_indices: vec![1, 0],
            values: vec![0.8, 0.6],
            n_samples: 2,
        };

        // With mix_weight = 1.0, we get full fuzzy union
        // union = 0.8 + 0.6 - 0.8*0.6 = 0.92
        // w_sym = 1.0 * union + 0.0 * w_ij = 0.92 for both edges
        let sym_graph = symmetrise_graph(graph.clone(), 1.0);

        let edges = sym_graph.to_edge_list();
        assert_eq!(edges.len(), 2);

        let union = 0.8 + 0.6 - 0.8 * 0.6;

        for (_, _, w) in edges {
            assert_relative_eq!(w, union, epsilon = 1e-6);
        }
    }

    #[test]
    fn test_coo_to_adjacency_list() {
        let graph = CoordinateList {
            row_indices: vec![0, 0, 1, 2],
            col_indices: vec![1, 2, 2, 0],
            values: vec![0.5, 0.3, 0.8, 0.9],
            n_samples: 3,
        };

        let adj = coo_to_adjacency_list(&graph);

        assert_eq!(adj.len(), 3);
        assert_eq!(adj[0].len(), 2); // vertex 0 has 2 neighbours
        assert_eq!(adj[1].len(), 1); // vertex 1 has 1 neighbour
        assert_eq!(adj[2].len(), 1); // vertex 2 has 1 neighbour

        assert!(adj[0].contains(&(1, 0.5)));
        assert!(adj[0].contains(&(2, 0.3)));
        assert!(adj[1].contains(&(2, 0.8)));
        assert!(adj[2].contains(&(0, 0.9)));
    }

    #[test]
    fn test_coo_to_adjacency_list_empty() {
        let graph: CoordinateList<f64> = CoordinateList {
            row_indices: vec![],
            col_indices: vec![],
            values: vec![],
            n_samples: 3,
        };

        let adj = coo_to_adjacency_list(&graph);

        assert_eq!(adj.len(), 3);
        assert!(adj[0].is_empty());
        assert!(adj[1].is_empty());
        assert!(adj[2].is_empty());
    }

    ////////////////
    // tSNE stuff //
    ////////////////

    /// Helper: build adjacency map from sparse graph for easier testing
    fn graph_to_adj<T: Float + Copy>(graph: &CoordinateList<T>) -> Vec<Vec<(usize, T)>> {
        let mut adj = vec![Vec::new(); graph.n_samples];
        for ((&i, &j), &w) in graph
            .row_indices
            .iter()
            .zip(&graph.col_indices)
            .zip(&graph.values)
        {
            adj[i].push((j, w));
        }
        adj
    }

    /// Helper: compute entropy of a probability distribution
    fn entropy(probs: &[f64]) -> f64 {
        probs
            .iter()
            .filter(|&&p| p > 1e-12)
            .map(|&p| -p * p.log2())
            .sum()
    }

    #[test]
    fn test_row_probabilities_sum_to_one() {
        // 5 points, each has 4 neighbours (excluding self)
        let knn_indices = vec![
            vec![1, 2, 3, 4],
            vec![0, 2, 3, 4],
            vec![0, 1, 3, 4],
            vec![0, 1, 2, 4],
            vec![0, 1, 2, 3],
        ];
        // Squared Euclidean distances
        let knn_dists = vec![
            vec![1.0, 4.0, 9.0, 16.0],
            vec![1.0, 1.0, 4.0, 9.0],
            vec![4.0, 1.0, 1.0, 4.0],
            vec![9.0, 4.0, 1.0, 1.0],
            vec![16.0, 9.0, 4.0, 1.0],
        ];

        let perplexity = 2.0;
        let graph =
            gaussian_knn_affinities(&knn_indices, &knn_dists, perplexity, 1e-5, 200, true).unwrap();
        let adj = graph_to_adj(&graph);

        for (i, neighbours) in adj.iter().enumerate() {
            let sum: f64 = neighbours.iter().map(|(_, w)| *w).sum();
            assert_relative_eq!(sum, 1.0, epsilon = 1e-4, max_relative = 1e-4);
            println!("Row {}: sum = {:.6}", i, sum);
        }
    }

    #[test]
    fn test_entropy_matches_target_perplexity() {
        // Create data where we can verify entropy
        let knn_indices = vec![
            vec![1, 2, 3, 4, 5, 6, 7],
            vec![0, 2, 3, 4, 5, 6, 7],
            vec![0, 1, 3, 4, 5, 6, 7],
            vec![0, 1, 2, 4, 5, 6, 7],
            vec![0, 1, 2, 3, 5, 6, 7],
            vec![0, 1, 2, 3, 4, 6, 7],
            vec![0, 1, 2, 3, 4, 5, 7],
            vec![0, 1, 2, 3, 4, 5, 6],
        ];
        // Squared distances with some variation
        let knn_dists: Vec<Vec<f64>> = (0..8)
            .map(|i| {
                (0..7)
                    .map(|j| ((j + 1) as f64) * (1.0 + 0.1 * (i as f64)))
                    .collect()
            })
            .collect();

        let perplexity = 3.0;
        let target_entropy = perplexity.log2();

        let graph =
            gaussian_knn_affinities(&knn_indices, &knn_dists, perplexity, 1e-5, 200, true).unwrap();
        let adj = graph_to_adj(&graph);

        for (i, neighbours) in adj.iter().enumerate() {
            let probs: Vec<f64> = neighbours.iter().map(|(_, w)| *w).collect();
            let h = entropy(&probs);
            println!(
                "Row {}: entropy = {:.4}, target = {:.4}, diff = {:.6}",
                i,
                h,
                target_entropy,
                (h - target_entropy).abs()
            );
            assert_relative_eq!(h, target_entropy, epsilon = 1e-3, max_relative = 1e-3);
        }
    }

    #[test]
    fn test_squared_vs_unsquared_equivalence() {
        // Same underlying distances, but one is squared, one is not
        let knn_indices = vec![
            vec![1, 2, 3, 4],
            vec![0, 2, 3, 4],
            vec![0, 1, 3, 4],
            vec![0, 1, 2, 4],
            vec![0, 1, 2, 3],
        ];

        // Unsquared Euclidean distances
        let unsquared: Vec<Vec<f64>> = vec![
            vec![1.0, 2.0, 3.0, 4.0],
            vec![1.0, 1.0, 2.0, 3.0],
            vec![2.0, 1.0, 1.0, 2.0],
            vec![3.0, 2.0, 1.0, 1.0],
            vec![4.0, 3.0, 2.0, 1.0],
        ];

        // Squared Euclidean distances
        let squared: Vec<Vec<f64>> = unsquared
            .iter()
            .map(|row| row.iter().map(|d| d * d).collect())
            .collect();

        let perplexity = 2.0;

        let graph_unsq =
            gaussian_knn_affinities(&knn_indices, &unsquared, perplexity, 1e-5, 200, false)
                .unwrap();
        let graph_sq =
            gaussian_knn_affinities(&knn_indices, &squared, perplexity, 1e-5, 200, true).unwrap();

        let adj_unsq = graph_to_adj(&graph_unsq);
        let adj_sq = graph_to_adj(&graph_sq);

        // Results should be identical
        for i in 0..5 {
            assert_eq!(adj_unsq[i].len(), adj_sq[i].len());
            for (a, b) in adj_unsq[i].iter().zip(adj_sq[i].iter()) {
                assert_eq!(a.0, b.0); // same neighbour index
                assert_relative_eq!(a.1, b.1, epsilon = 1e-10);
            }
        }
        println!("Squared vs unsquared: results match!");
    }

    #[test]
    fn test_self_loops_excluded() {
        // kNN includes self (index i appears in knn_indices[i])
        let knn_indices = vec![
            vec![0, 1, 2, 3], // includes self
            vec![1, 0, 2, 3], // includes self
            vec![2, 0, 1, 3], // includes self
            vec![3, 0, 1, 2], // includes self
        ];
        let knn_dists = vec![
            vec![0.0, 1.0, 4.0, 9.0], // distance to self is 0
            vec![0.0, 1.0, 1.0, 4.0],
            vec![0.0, 4.0, 1.0, 1.0],
            vec![0.0, 9.0, 4.0, 1.0],
        ];

        let graph =
            gaussian_knn_affinities(&knn_indices, &knn_dists, 2.0, 1e-5, 200, true).unwrap();

        // Check no self-loops in output
        for (&i, &j) in graph.row_indices.iter().zip(&graph.col_indices) {
            assert_ne!(i, j, "Self-loop found: {} -> {}", i, j);
        }
        println!("No self-loops in output graph.");
    }

    #[test]
    fn test_closer_neighbours_have_higher_probability() {
        let knn_indices = vec![vec![1, 2, 3, 4]];
        // Strictly increasing squared distances
        let knn_dists = vec![vec![1.0, 4.0, 9.0, 16.0]];

        let graph =
            gaussian_knn_affinities(&knn_indices, &knn_dists, 2.0, 1e-5, 200, true).unwrap();
        let adj = graph_to_adj(&graph);

        let probs: Vec<(usize, f64)> = adj[0].clone();
        println!("Probabilities: {:?}", probs);

        // Closer neighbours should have higher probability
        // neighbour 1 (d²=1) > neighbour 2 (d²=4) > neighbour 3 (d²=9) > neighbour 4 (d²=16)
        let p1 = probs.iter().find(|(j, _)| *j == 1).unwrap().1;
        let p2 = probs.iter().find(|(j, _)| *j == 2).unwrap().1;
        let p3 = probs.iter().find(|(j, _)| *j == 3).unwrap().1;
        let p4 = probs.iter().find(|(j, _)| *j == 4).unwrap().1;

        assert!(p1 > p2, "p1={} should be > p2={}", p1, p2);
        assert!(p2 > p3, "p2={} should be > p3={}", p2, p3);
        assert!(p3 > p4, "p3={} should be > p4={}", p3, p4);
    }

    #[test]
    fn test_uniform_distances_give_uniform_probs() {
        // All neighbours at same distance → should get uniform distribution
        // Use perplexity = 4 so target entropy matches uniform entropy over 4 items
        let knn_indices = vec![vec![1, 2, 3, 4]];
        let knn_dists = vec![vec![4.0, 4.0, 4.0, 4.0]]; // all same squared distance

        let perplexity = 3.99999999; // to not throw an error
        let graph =
            gaussian_knn_affinities(&knn_indices, &knn_dists, perplexity, 1e-5, 200, true).unwrap();
        let adj = graph_to_adj(&graph);

        let probs: Vec<f64> = adj[0].iter().map(|(_, p)| *p).collect();
        let expected = 0.25; // uniform over 4 neighbours

        for (i, &p) in probs.iter().enumerate() {
            assert_relative_eq!(p, expected, epsilon = 1e-1); // generous threshold
            println!("Neighbour {}: p = {:.6}", i, p);
        }
    }

    #[test]
    fn test_perplexity_affects_distribution_spread() {
        let knn_indices = vec![vec![1, 2, 3, 4, 5, 6, 7]];
        let knn_dists = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]]; // unsquared

        // Low perplexity → more concentrated distribution
        let graph_low =
            gaussian_knn_affinities(&knn_indices, &knn_dists, 1.5, 1e-5, 200, false).unwrap();
        let adj_low = graph_to_adj(&graph_low);
        let probs_low: Vec<f64> = adj_low[0].iter().map(|(_, p)| *p).collect();
        let entropy_low = entropy(&probs_low);

        // High perplexity → more spread distribution
        let graph_high =
            gaussian_knn_affinities(&knn_indices, &knn_dists, 4.0, 1e-5, 200, false).unwrap();
        let adj_high = graph_to_adj(&graph_high);
        let probs_high: Vec<f64> = adj_high[0].iter().map(|(_, p)| *p).collect();
        let entropy_high = entropy(&probs_high);

        println!("Low perplexity (1.5): entropy = {:.4}", entropy_low);
        println!("High perplexity (4.0): entropy = {:.4}", entropy_high);

        assert!(
            entropy_high > entropy_low,
            "Higher perplexity should give higher entropy"
        );
    }

    ///////////
    // PHATE //
    ///////////

    #[test]
    fn test_phate_basic_affinity_computation() {
        // 4 points, each has 3 neighbours (including self at position 0)
        let knn_indices = vec![
            vec![0, 1, 2, 3],
            vec![1, 0, 2, 3],
            vec![2, 0, 1, 3],
            vec![3, 0, 1, 2],
        ];
        // Squared Euclidean distances (self at position 0 with distance 0)
        let knn_dists = vec![
            vec![0.0, 1.0, 4.0, 9.0],
            vec![0.0, 1.0, 4.0, 9.0],
            vec![0.0, 4.0, 1.0, 9.0],
            vec![0.0, 9.0, 4.0, 1.0],
        ];

        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,          // knn: use 2nd neighbor (indices[1]) for bandwidth
            Some(40.0), // decay
            1.0,        // bandwidth_scale
            1e-4,       // thresh
            "none",     // no symmetrisation
            true,
        );

        assert_eq!(graph.n_samples, 4);
        assert!(!graph.row_indices.is_empty());
        assert_eq!(graph.row_indices.len(), graph.col_indices.len());
        assert_eq!(graph.row_indices.len(), graph.values.len());

        // All affinities should be between 0 and 1
        for &v in &graph.values {
            assert!((0.0..=1.0).contains(&v), "Affinity {} out of range", v);
        }
        println!("Basic test passed: {} edges created", graph.values.len());
    }

    #[test]
    fn test_phate_self_loops_excluded() {
        // Self is at position 0 in the knn arrays
        let knn_indices = vec![vec![0, 1, 2], vec![1, 0, 2], vec![2, 0, 1]];
        let knn_dists = vec![
            vec![0.0, 1.0, 4.0],
            vec![0.0, 1.0, 4.0],
            vec![0.0, 4.0, 1.0],
        ];

        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(40.0),
            1.0,
            1e-4,
            "none",
            true,
        );

        // Check no self-loops
        for (&i, &j) in graph.row_indices.iter().zip(&graph.col_indices) {
            assert_ne!(i, j, "Self-loop found: {} -> {}", i, j);
        }
        println!("No self-loops in PHATE graph");
    }

    #[test]
    fn test_phate_closer_neighbours_higher_affinity() {
        // Single point with 4 neighbors at increasing distances
        let knn_indices = vec![vec![0, 1, 2, 3, 4]];
        // Squared distances: 0 (self), 1, 4, 9, 16
        let knn_dists = vec![vec![0.0, 1.0, 4.0, 9.0, 16.0]];

        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "none",
            true,
        );

        let adj = graph_to_adj(&graph);
        let affinities: Vec<(usize, f64)> = adj[0].clone();

        // Get affinities for each neighbor
        let a1 = affinities.iter().find(|(j, _)| *j == 1).unwrap().1;
        let a2 = affinities.iter().find(|(j, _)| *j == 2).unwrap().1;
        let a3 = affinities.iter().find(|(j, _)| *j == 3).unwrap().1;
        let a4 = affinities.iter().find(|(j, _)| *j == 4).unwrap().1;

        println!(
            "Affinities: a1={:.6}, a2={:.6}, a3={:.6}, a4={:.6}",
            a1, a2, a3, a4
        );

        // closer neighbors should have higher affinities
        assert!(a1 > a2, "a1={} should be > a2={}", a1, a2);
        assert!(a2 > a3, "a2={} should be > a3={}", a2, a3);
        assert!(a3 > a4, "a3={} should be > a4={}", a3, a4);
    }

    #[test]
    fn test_phate_bandwidth_from_kth_neighbor() {
        // Test that bandwidth is correctly computed from the kth neighbor
        let knn_indices = vec![vec![0, 1, 2, 3]];
        let knn_dists = vec![vec![0.0, 1.0, 4.0, 16.0]]; // squared distances

        // Use k=2: bandwidth should be sqrt(1.0) = 1.0
        let graph_k2 = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "none",
            true,
        );

        // Use k=3: bandwidth should be sqrt(4.0) = 2.0
        let graph_k3 = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            3,
            Some(2.0),
            1.0,
            1e-10,
            "none",
            true,
        );

        let adj_k2 = graph_to_adj(&graph_k2);
        let adj_k3 = graph_to_adj(&graph_k3);

        // Get affinities for near and far neighbors
        let a1_k2 = adj_k2[0].iter().find(|(j, _)| *j == 1).unwrap().1; // near
        let a2_k2 = adj_k2[0].iter().find(|(j, _)| *j == 2).unwrap().1; // far

        let a1_k3 = adj_k3[0].iter().find(|(j, _)| *j == 1).unwrap().1; // near
        let a2_k3 = adj_k3[0].iter().find(|(j, _)| *j == 2).unwrap().1; // far

        println!(
            "k=2: a1={:.6}, a2={:.6}, ratio={:.6}",
            a1_k2,
            a2_k2,
            a1_k2 / a2_k2
        );
        println!(
            "k=3: a1={:.6}, a2={:.6}, ratio={:.6}",
            a1_k3,
            a2_k3,
            a1_k3 / a2_k3
        );

        // Smaller bandwidth (k=2) should give LARGER ratio between near and far neighbors
        // (more peaked distribution)
        let ratio_k2 = a1_k2 / a2_k2;
        let ratio_k3 = a1_k3 / a2_k3;

        assert!(
        ratio_k2 > ratio_k3,
        "Smaller bandwidth should give higher ratio (more peaked): ratio_k2={:.6} vs ratio_k3={:.6}",
        ratio_k2, ratio_k3
    );

        // Also verify: larger bandwidth gives higher absolute affinity to nearest neighbor
        assert!(
            a1_k3 > a1_k2,
            "Larger bandwidth should give higher absolute affinity to nearest neighbor"
        );
    }

    #[test]
    fn test_phate_bandwidth_scale_effect() {
        let knn_indices = vec![vec![0, 1, 2, 3]];
        let knn_dists = vec![vec![0.0, 1.0, 4.0, 9.0]];

        // Test with bandwidth_scale = 1.0
        let graph_scale1 = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "none",
            true,
        );

        // Test with bandwidth_scale = 2.0 (wider kernel)
        let graph_scale2 = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            2.0,
            1e-10,
            "none",
            true,
        );

        let adj1 = graph_to_adj(&graph_scale1);
        let adj2 = graph_to_adj(&graph_scale2);

        // For a distant neighbor, wider bandwidth should give higher affinity
        let a3_scale1 = adj1[0].iter().find(|(j, _)| *j == 3).unwrap().1;
        let a3_scale2 = adj2[0].iter().find(|(j, _)| *j == 3).unwrap().1;

        println!(
            "scale=1.0: a3={:.6}, scale=2.0: a3={:.6}",
            a3_scale1, a3_scale2
        );
        assert!(
            a3_scale2 > a3_scale1,
            "Larger bandwidth scale should give higher affinity to distant neighbors"
        );
    }

    #[test]
    fn test_phate_decay_parameter_effect() {
        let knn_indices = vec![vec![0, 1, 2, 3, 4]];
        let knn_dists = vec![vec![0.0, 1.0, 4.0, 9.0, 16.0]];

        // Low decay (α=10): gentler falloff
        let graph_low = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(1.0),
            1.0,
            1e-10,
            "none",
            true,
        );

        // High decay (α=80): sharper falloff
        let graph_high = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "none",
            true,
        );

        let adj_low = graph_to_adj(&graph_low);
        let adj_high = graph_to_adj(&graph_high);

        // For distant neighbor, low decay should give higher affinity
        if let Some(&(_, a4_low)) = adj_low[0].iter().find(|(j, _)| *j == 4) {
            if let Some(&(_, a4_high)) = adj_high[0].iter().find(|(j, _)| *j == 4) {
                println!("decay=10: a4={:.6}, decay=80: a4={:.6}", a4_low, a4_high);
                assert!(
                    a4_low > a4_high,
                    "Lower decay should give higher affinity to distant neighbors"
                );
            }
        }
    }

    #[test]
    fn test_phate_thresholding() {
        let knn_indices = vec![vec![0, 1, 2, 3, 4, 5]];
        // Create distances that will result in some very small affinities
        let knn_dists = vec![vec![0.0, 1.0, 4.0, 9.0, 16.0, 25.0]];

        // Strict threshold - should exclude far neighbors
        let graph_strict = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-5, // higher threshold
            "none",
            true,
        );

        // Lenient threshold - should include more neighbors
        let graph_lenient = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10, // lower threshold
            "none",
            true,
        );

        println!(
            "Strict threshold edges: {}, Lenient threshold edges: {}",
            graph_strict.values.len(),
            graph_lenient.values.len()
        );

        assert!(
            graph_lenient.values.len() >= graph_strict.values.len(),
            "Lenient threshold should produce at least as many edges"
        );
    }

    #[test]
    fn test_phate_binary_connectivity() {
        let knn_indices = vec![
            vec![0, 1, 2, 3],
            vec![1, 0, 2, 3],
            vec![2, 0, 1, 3],
            vec![3, 0, 1, 2],
        ];
        let knn_dists = vec![
            vec![0.0, 1.0, 4.0, 9.0],
            vec![0.0, 1.0, 4.0, 9.0],
            vec![0.0, 4.0, 1.0, 9.0],
            vec![0.0, 9.0, 4.0, 1.0],
        ];

        // decay = None should give binary connectivity
        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            None, // binary mode
            1.0,
            1e-4,
            "none",
            true,
        );

        // All edges should have weight 1.0
        for &v in &graph.values {
            assert_relative_eq!(v, 1.0, epsilon = 1e-10);
        }
        println!(
            "Binary connectivity: all {} edges have weight 1.0",
            graph.values.len()
        );
    }

    #[test]
    fn test_phate_additive_symmetrisation() {
        // Create an asymmetric scenario
        let knn_indices = vec![
            vec![0, 1, 2],
            vec![1, 2, 0], // Note: 1 is closer to 2 than to 0
            vec![2, 1, 0],
        ];
        let knn_dists = vec![
            vec![0.0, 1.0, 4.0],
            vec![0.0, 1.0, 9.0],
            vec![0.0, 1.0, 4.0],
        ];

        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-4,
            "add", // additive symmetrisation
            true,
        );

        // Check symmetry: K[i,j] should equal K[j,i]
        let edges: std::collections::HashMap<(usize, usize), f64> = graph
            .row_indices
            .iter()
            .zip(&graph.col_indices)
            .zip(&graph.values)
            .map(|((&i, &j), &v)| ((i, j), v))
            .collect();

        for (i, j) in [(0, 1), (0, 2), (1, 2)] {
            if let (Some(&v_ij), Some(&v_ji)) = (edges.get(&(i, j)), edges.get(&(j, i))) {
                assert_relative_eq!(v_ij, v_ji, epsilon = 1e-6, max_relative = 1e-6);
                println!(
                    "Symmetric edge ({},{}): v_ij={:.6}, v_ji={:.6}",
                    i, j, v_ij, v_ji
                );
            }
        }
    }

    #[test]
    fn test_phate_multiplicative_symmetrisation() {
        let knn_indices = vec![vec![0, 1, 2], vec![1, 0, 2], vec![2, 0, 1]];
        let knn_dists = vec![
            vec![0.0, 1.0, 4.0],
            vec![0.0, 1.0, 4.0],
            vec![0.0, 4.0, 1.0],
        ];

        // Get asymmetric version first
        let graph_asym = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "none",
            true,
        );

        // Get multiplicative symmetrized version
        let graph_sym = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "multiply",
            true,
        );

        // Multiplicative should have fewer edges (only mutual neighbors)
        assert!(
            graph_sym.values.len() <= graph_asym.values.len(),
            "Multiplicative symmetrisation should have <= edges"
        );

        // Check symmetry
        let edges: std::collections::HashMap<(usize, usize), f64> = graph_sym
            .row_indices
            .iter()
            .zip(&graph_sym.col_indices)
            .zip(&graph_sym.values)
            .map(|((&i, &j), &v)| ((i, j), v))
            .collect();

        for (&(i, j), &v_ij) in &edges {
            if let Some(&v_ji) = edges.get(&(j, i)) {
                assert_relative_eq!(v_ij, v_ji, epsilon = 1e-6);
            }
        }
        println!(
            "Multiplicative symmetrisation: {} edges",
            graph_sym.values.len()
        );
    }

    #[test]
    fn test_phate_symmetrisation_comparison() {
        let knn_indices = vec![
            vec![0, 1, 2, 3],
            vec![1, 0, 2, 3],
            vec![2, 1, 0, 3],
            vec![3, 0, 1, 2],
        ];
        let knn_dists = vec![
            vec![0.0, 1.0, 4.0, 9.0],
            vec![0.0, 2.0, 3.0, 8.0],
            vec![0.0, 1.0, 5.0, 7.0],
            vec![0.0, 6.0, 7.0, 8.0],
        ];

        let graph_none = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "none",
            true,
        );
        let graph_add = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "add",
            true,
        );
        let graph_mult = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(2.0),
            1.0,
            1e-10,
            "multiply",
            true,
        );

        println!("Asymmetric: {} edges", graph_none.values.len());
        println!("Additive: {} edges", graph_add.values.len());
        println!("Multiplicative: {} edges", graph_mult.values.len());

        // Additive should preserve most edges
        assert!(graph_add.values.len() >= graph_mult.values.len());

        // All methods should produce reasonable number of edges
        assert!(!graph_add.values.is_empty());
        assert!(!graph_mult.values.is_empty());
    }

    #[test]
    fn test_phate_affinity_formula_verification() {
        // Manual verification of the affinity formula
        let knn_indices = vec![vec![0, 1]];
        let knn_dists = vec![vec![0.0, 4.0]]; // squared distance = 4, actual distance = 2

        let decay = 2.0;
        let bandwidth_scale = 1.0;

        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(decay),
            bandwidth_scale,
            1e-10,
            "none",
            true,
        );

        // Manual calculation:
        // bandwidth = sqrt(4.0) * 1.0 = 2.0
        // distance = sqrt(4.0) = 2.0
        // scaled = 2.0 / 2.0 = 1.0
        // powered = 1.0^40 = 1.0
        // affinity = exp(-1.0) ≈ 0.3678794411714423

        let expected = (-1.0_f64).exp();
        let actual = graph.values[0];

        println!("Expected affinity: {:.10}", expected);
        println!("Actual affinity: {:.10}", actual);

        assert_relative_eq!(actual, expected, epsilon = 1e-6);
    }

    #[test]
    fn test_phate_zero_distance_handling() {
        // Test that zero distances (duplicates) are handled correctly
        let knn_indices = vec![vec![0, 1, 2]];
        let knn_dists = vec![vec![0.0, 0.0, 1.0]]; // duplicate at distance 0

        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            2,
            Some(40.0),
            1.0,
            1e-4,
            "none",
            true,
        );

        // Find affinity for the zero-distance neighbor
        let adj = graph_to_adj(&graph);
        if let Some(&(j, v)) = adj[0].iter().find(|(j, _)| *j == 1) {
            println!("Zero-distance neighbor {} has affinity {:.6}", j, v);
            assert_relative_eq!(v, 1.0, epsilon = 1e-6,);
        }
    }

    #[test]
    fn test_phate_large_dataset_edge_count() {
        // Test with a larger dataset to ensure edge count is reasonable
        let n = 100;
        let k = 10;

        let knn_indices: Vec<Vec<usize>> = (0..n)
            .map(|i| {
                let mut indices: Vec<usize> = (0..k).map(|j| (i + j) % n).collect();
                indices[0] = i; // self at position 0
                indices
            })
            .collect();

        let knn_dists: Vec<Vec<f64>> = (0..n)
            .map(|_| {
                let mut dists: Vec<f64> = (0..k).map(|j| (j as f64).powi(2)).collect();
                dists[0] = 0.0; // self distance
                dists
            })
            .collect();

        let graph = phate_alpha_decay_affinities(
            &knn_indices,
            &knn_dists,
            5,
            Some(2.0),
            1.0,
            1e-10,
            "add",
            true,
        );

        println!(
            "Large dataset: {} vertices, {} edges",
            n,
            graph.values.len()
        );

        // Should have roughly n * (k-1) edges (minus self-loops and thresholding)
        assert!(!graph.values.is_empty());
        assert!(graph.values.len() <= n * (k - 1) * 2); // *2 for symmetrisation
    }
}