pandora_lib_promethion/src/annotation/mod.rs

use std::{
    collections::{HashMap, HashSet},
    fmt,
};

use crate::helpers::mean;
use dashmap::DashMap;
use rayon::prelude::*;

#[derive(Debug, Clone, PartialEq)]
pub enum Annotation {
    SoloDiag,
    SoloConstit,
    Callers(Caller),
    Germline,
    Somatic,
    ShannonEntropy(f64),
    ConstitDepth(u16),
    ConstitAlt(u16),
    LowConstitDepth,
    HighConstitAlt,
}

impl fmt::Display for Annotation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let str = match self {
            Annotation::SoloDiag => "SoloDiag",
            Annotation::SoloConstit => "SoloConstit",
            Annotation::Callers(caller) => &caller.to_string(),
            Annotation::Germline => "Germline",
            Annotation::Somatic => "Somatic",
            Annotation::ShannonEntropy(_) => "ShannonEntropy",
            Annotation::ConstitDepth(_) => "ConstitDepth",
            Annotation::ConstitAlt(_) => "ConstitAlt",
            Annotation::LowConstitDepth => "LowConstitDepth",
            Annotation::HighConstitAlt => "HighConstitAlt",
        };
        write!(f, "{}", str)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Caller {
    DeepVariant,
    ClairS,
}

impl fmt::Display for Caller {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Caller::DeepVariant => write!(f, "DeepVariant"),
            Caller::ClairS => write!(f, "ClairS"),
        }
    }
}

#[derive(Debug, Default, Clone)]
pub struct Annotations {
    pub store: DashMap<u128, Vec<Annotation>>,
}

impl Annotations {
    pub fn insert_update(&self, key: u128, add: &[Annotation]) {
        self.store
            .entry(key)
            .or_default()
            .extend(add.iter().cloned())
    }

    pub fn callers_stat(&self) {
        let map: DashMap<String, u64> = DashMap::new();
        let num_maps: DashMap<String, HashMap<String, Vec<f64>>> = DashMap::new();

        self.store.par_iter().for_each(|e| {
            let anns = e.value();

            let mut categorical = Vec::new();
            let mut numerical = Vec::new();
            for ann in anns {
                match ann {
                    Annotation::SoloDiag
                    | Annotation::SoloConstit
                    | Annotation::Germline
                    | Annotation::Somatic
                    | Annotation::LowConstitDepth
                    | Annotation::HighConstitAlt => categorical.push(ann.to_string()),
                    Annotation::Callers(caller) => categorical.push(caller.to_string()),
                    Annotation::ShannonEntropy(v) => numerical.push((ann.to_string(), *v)),
                    Annotation::ConstitDepth(v) | Annotation::ConstitAlt(v) => {
                        numerical.push((ann.to_string(), *v as f64));
                    }
                }
            }

            categorical.sort();
            categorical.dedup();
            let k = categorical.join(" + ");

            *map.entry(k.clone()).or_default() += 1;

            for (k_num, v_num) in numerical {
                num_maps
                    .entry(k.clone())
                    .or_default()
                    .entry(k_num)
                    .or_default()
                    .push(v_num);
            }
        });
        println!("\nCallers stats:");
        println!("\tcategories: {}", map.len());
        let mut n = 0;
        map.iter().for_each(|e| {
            let k = e.key();
            let v = e.value();
            n += v;
            let mut num_str = Vec::new();
            if let Some(nums) = num_maps.get(k) {
                num_str.extend(
                    nums.iter()
                        .map(|(k_n, v_n)| format!("{k_n} {:.2}", mean(v_n))),
                )
            }
            num_str.sort();
            println!("\t{k}\t{v}\t{}", num_str.join("\t"));
        });

        println!("Total\t{n}");
    }

    pub fn get_keys_filter(
        &self,
        filter: impl Fn(&Vec<Annotation>) -> bool + Send + Sync,
    ) -> Vec<u128> {
        self.store
            .par_iter()
            .filter(|entry| filter(entry.value()))
            .map(|entry| *entry.key())
            .collect()
    }

    pub fn retain_keys(&mut self, keys_to_keep: &HashSet<u128>) {
        self.store.retain(|key, _| keys_to_keep.contains(key));
    }

    pub fn solo_constit_boundaries(&self, max_alt_constit: u16, min_constit_depth: u16) {
        self.store
            .iter_mut()
            .filter(|anns| {
                let contains = anns.iter().any(|item| matches!(item, Annotation::SoloDiag));
                let contains_not = anns.iter().all(|item| !matches!(item, Annotation::Somatic));

                contains && contains_not
            })
            .for_each(|mut e| {
                let v = e.value_mut();
                let mut to_add = Vec::new();
                v.iter().for_each(|ann| match ann {
                    Annotation::ConstitDepth(v) => {
                        if *v < min_constit_depth {
                            to_add.push(Annotation::LowConstitDepth);
                        }
                    }
                    Annotation::ConstitAlt(v) => {
                        if *v > max_alt_constit {
                            to_add.push(Annotation::HighConstitAlt);
                        }
                    },
                    _ => (),
                });
                v.extend(to_add);
            });
    }
}