use rand::prelude::IndexedRandom;
use rand::{Rng, SeedableRng, rngs::StdRng};
use rustc_hash::{FxHashMap, FxHashSet};
use crate::prelude::*;
#[derive(Clone, Debug)]
pub struct BootstrappedMetaCellParams {
pub max_shared: usize,
pub target_no_metacells: usize,
pub max_iter: usize,
pub knn_params: KnnParams,
}
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, FxHashSet<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(|| nn_map[candidate].iter().copied().collect());
let mut max_overlap = 0;
for &existing in &chosen {
set_cache
.entry(existing)
.or_insert_with(|| nn_map[existing].iter().copied().collect());
let cs = &set_cache[&candidate];
let es = &set_cache[&existing];
let shared = k2 - cs.union(es).count();
max_overlap = max_overlap.max(shared);
}
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
}
pub fn assign_bootstrapped_meta_cells(centres: &[usize], nn_map: &[Vec<usize>]) -> Vec<Vec<usize>> {
centres.iter().map(|&c| nn_map[c].clone()).collect()
}