bixverse-sc 0.8.1

//! Meta-cell aggregation based on the bootstrapped approach from Morabito,
//! et al., Cell Rep. Methods, 2023

use rand::prelude::IndexedRandom;
use rand::{Rng, SeedableRng, rngs::StdRng};
use rustc_hash::FxHashMap;

use crate::prelude::*;

////////////////////////
// hdWGCNA meta cells //
////////////////////////

/// Structure for the MetaCell parameters
#[derive(Clone, Debug)]
pub struct BootstrappedMetaCellParams {
    /// Maximum number of shared cells for the meta cell aggregation
    pub max_shared: usize,
    /// Number of target meta cells.
    pub target_no_metacells: usize,
    /// Maximum iterations for the algorithm
    pub max_iter: usize,
    /// [KnnParams] for the various approximate nearest neighbour searches in
    /// ann-search-rs
    pub knn_params: KnnParams,
}

/// Sorted, deduplicated copy of a neighbour list.
///
/// ### Params
///
/// * `v` - Neighbour indices
///
/// ### Returns
///
/// The sorted unique indices.
fn sorted_unique(v: &[usize]) -> Vec<usize> {
    let mut out = v.to_vec();
    out.sort_unstable();
    out.dedup();
    out
}

/// Size of the intersection of two sorted unique slices.
///
/// ### Params
///
/// * `a` - First sorted unique slice
/// * `b` - Second sorted unique slice
///
/// ### Returns
///
/// The number of shared elements.
fn sorted_intersection_len(a: &[usize], b: &[usize]) -> usize {
    let (mut i, mut j, mut shared) = (0, 0, 0);
    while i < a.len() && j < b.len() {
        match a[i].cmp(&b[j]) {
            std::cmp::Ordering::Less => i += 1,
            std::cmp::Ordering::Greater => j += 1,
            std::cmp::Ordering::Equal => {
                shared += 1;
                i += 1;
                j += 1;
            }
        }
    }
    shared
}

/// Select meta cells
///
/// ### Params
///
/// * `nn_map` - Nearest neighbours with self
/// * `max_shared` - Maximum number of shared neighbours to allow
/// * `target_no` - Target number of meta cells
/// * `max_iter` - Maximum iterations for the algorithm
/// * `seed` - seed for reproducibility purposes
/// * `verbose` - If `0` -> silent or `1` for normal verbosity, `2` for detailed
///   verbosity.
///
/// ### Returns
///
/// Borrowed slice of the selected meta cell indices
pub fn identify_meta_cells(
    nn_map: &[Vec<usize>],
    max_shared: usize,
    target_no: usize,
    max_iter: usize,
    seed: usize,
    verbose: usize,
) -> Vec<usize> {
    let verbosity = parse_verbosity_level(verbose);

    let mut rng = StdRng::seed_from_u64(seed as u64);
    let k = nn_map[0].len();
    let k2 = k * 2;
    let mut good_choices: Vec<usize> = (0..nn_map.len()).collect();
    let mut chosen: Vec<usize> = Vec::new();

    if let Some(&first) = good_choices.choose(&mut rng) {
        chosen.push(first);
        good_choices.retain(|&x| x != first);
    }

    let mut it = 0;
    let mut set_cache: FxHashMap<usize, Vec<usize>> = FxHashMap::default();

    while !good_choices.is_empty() && chosen.len() < target_no && it < max_iter {
        it += 1;
        let choice_idx = rng.random_range(0..good_choices.len());
        let candidate = good_choices.swap_remove(choice_idx);

        set_cache
            .entry(candidate)
            .or_insert_with(|| sorted_unique(&nn_map[candidate]));

        let mut max_overlap = 0;
        for &existing in &chosen {
            set_cache
                .entry(existing)
                .or_insert_with(|| sorted_unique(&nn_map[existing]));
            let cs = &set_cache[&candidate];
            let es = &set_cache[&existing];
            let shared = k2 - (cs.len() + es.len() - sorted_intersection_len(cs, es));
            max_overlap = max_overlap.max(shared);
            // only `max_overlap <= max_shared` is read, so stop at the first miss
            if max_overlap > max_shared {
                break;
            }
        }

        if verbosity.normal_verbosity() && it % 10000 == 0 {
            println!(
                "Meta cell neighbour search - iter {} of {} max iters",
                it, max_iter
            );
        }

        if max_overlap <= max_shared {
            chosen.push(candidate);
        }
    }

    chosen
}

/// Assign cells to bootstrapped metacells with centre priority.
///
/// Centres always claim themselves first, so no metacell is empty. Remaining
/// cells go to the first centre whose kNN window contains them.
///
/// ### Params
///
/// * `centres` - The center cell.
/// * `nn_map` - The nearest neighbour map
///
/// ### Returns
///
/// The meta cell assignments
pub fn assign_bootstrapped_meta_cells(centres: &[usize], nn_map: &[Vec<usize>]) -> Vec<Vec<usize>> {
    centres.iter().map(|&c| nn_map[c].clone()).collect()
}