md_analysis 0.1.0

/// 直方图分箱逻辑
///
/// 包含 Histogram、BaseHistogram 结构体以及从 BAM aux tag 收集值到 bins
/// 的所有函数。支持整体统计和 ACGT 分碱基统计。
use std::{collections::HashMap, vec};

// ── 整体直方图 ─────────────────────────────────────────────

/// 分子动力学直方图：dw 和 ar 两个 bin 数组
#[derive(Debug, Clone)]
pub struct Histogram<const N: usize> {
    pub dw: Vec<u64>,
    pub ar: Vec<u64>,
}

impl<const N: usize> Histogram<N> {
    /// 创建空的直方图
    pub fn new() -> Self {
        Self {
            dw: vec![0u64; N],
            ar: vec![0u64; N],
        }
    }

    pub fn collect_ar_dw<'a>(
        &mut self,
        ar: impl Iterator<Item = &'a u32>,
        dw: impl Iterator<Item = &'a u32>,
    ) {
        self.collect_ar(ar);
        self.collect_dw(dw);
    }

    pub fn collect_dw<'a>(&mut self, data: impl Iterator<Item = &'a u32>) {
        data.for_each(|&v| {
            let v = v.min((N - 1) as u32);
            self.dw[v as usize] += 1;
        });
    }

    pub fn collect_ar<'a>(&mut self, data: impl Iterator<Item = &'a u32>) {
        data.for_each(|&v| {
            let v = v.min((N - 1) as u32);
            self.ar[v as usize] += 1;
        });
    }

    pub fn merge(&mut self, hist: &Histogram<N>) {
        self.dw
            .iter_mut()
            .zip(hist.dw.iter())
            .for_each(|(dest, src)| {
                *dest += *src;
            });

        self.ar
            .iter_mut()
            .zip(hist.ar.iter())
            .for_each(|(dest, src)| {
                *dest += *src;
            });
    }

    /// 汇总统计摘要（每行一条字符串）
    pub fn summary(&self) -> Vec<String> {
        let mut lines = Vec::new();
        for (name, bins) in &[("dw", &self.dw), ("ar", &self.ar)] {
            let total: u64 = bins.iter().sum();
            if total == 0 {
                lines.push(format!("{}: 无数据", name));
                continue;
            }

            let nonzero_count: u64 = bins.iter().skip(1).sum();
            if nonzero_count == 0 {
                lines.push(format!("{}: {} 个零值", name, total));
                continue;
            }

            let mut min_v = N;
            let mut max_v = 0;
            for (i, &count) in bins.iter().enumerate() {
                if count > 0 && i > 0 {
                    if i < min_v {
                        min_v = i;
                    }
                    if i > max_v {
                        max_v = i;
                    }
                }
            }
            lines.push(format!(
                "{}: {} 个非零值 (范围: {} ~ {})",
                name, nonzero_count, min_v, max_v
            ));
        }
        lines
    }
}

// ── 分碱基直方图 ────────────────────────────────────────────

/// ACGT 单个碱基类型的直方图
#[derive(Debug, Clone)]
pub struct PerBaseHistogram<const N: usize> {
    pub hist_grams: HashMap<char, Histogram<N>>,
}

impl<const N: usize> PerBaseHistogram<N> {
    pub fn new() -> Self {
        let hist_grams: HashMap<char, Histogram<N>> = "ACGT"
            .chars()
            .into_iter()
            .map(|base| (base, Histogram::<N>::new()))
            .collect();
        Self { hist_grams }
    }

    pub fn collect_ar_dw(&mut self, bases: &str, ar: &[u32], dw: &[u32]) {
        self.collect_ar(bases, ar);
        self.collect_dw(bases, dw);
    }

    pub fn collect_dw(&mut self, bases: &str, data: &[u32]) {
        "ACGT".chars().into_iter().for_each(|cur_base| {
            let cur_base_hist = self.hist_grams.get_mut(&cur_base).unwrap();

            let iter = bases
                .chars()
                .into_iter()
                .zip(data.iter())
                .filter(|&(query_base, _)| query_base == cur_base)
                .map(|(_, v)| v);

            cur_base_hist.collect_dw(iter);
        });
    }

    pub fn collect_ar(&mut self, bases: &str, data: &[u32]) {
        "ACGT".chars().into_iter().for_each(|cur_base| {
            let cur_base_hist = self.hist_grams.get_mut(&cur_base).unwrap();

            let iter = bases
                .chars()
                .into_iter()
                .zip(data.iter())
                .filter(|&(query_base, _)| query_base == cur_base)
                .map(|(_, v)| v);

            cur_base_hist.collect_ar(iter);
        });
    }

    pub fn merge(&mut self, hist: &PerBaseHistogram<N>) {
        "ACGT".chars().into_iter().for_each(|cur_base| {
            let dest = self.hist_grams.get_mut(&cur_base).unwrap();
            let src = hist.hist_grams.get(&cur_base).unwrap();
            dest.merge(src);
        });
    }
}

pub struct FinalResult<const N: usize> {
    pub total_hist: Option<Histogram<N>>,
    pub total_per_base_hist: Option<PerBaseHistogram<N>>,
    pub first_n_total_hist: Option<Histogram<N>>,
    pub first_n_per_base_hist: Option<PerBaseHistogram<N>>,

    pub last_n_total_hist: Option<Histogram<N>>,
    pub last_n_per_base_hist: Option<PerBaseHistogram<N>>,
}

impl<const N: usize> FinalResult<N> {
    pub fn new_first_n_last_n() -> Self {
        Self {
            total_hist: None,
            total_per_base_hist: None,
            first_n_total_hist: Some(Histogram::new()),
            first_n_per_base_hist: Some(PerBaseHistogram::new()),
            last_n_total_hist: Some(Histogram::new()),
            last_n_per_base_hist: Some(PerBaseHistogram::new()),
        }
    }

    pub fn new_all() -> Self {
        Self {
            total_hist: Some(Histogram::new()),
            total_per_base_hist: Some(PerBaseHistogram::new()),
            first_n_total_hist: None,
            first_n_per_base_hist: None,
            last_n_total_hist: None,
            last_n_per_base_hist: None,
        }
    }

    pub fn get_total_hist_mut(&mut self) -> &mut Histogram<N> {
        self.total_hist.as_mut().unwrap()
    }

    pub fn get_total_per_base_hist_mut(&mut self) -> &mut PerBaseHistogram<N> {
        self.total_per_base_hist.as_mut().unwrap()
    }

    pub fn get_first_n_total_hist_mut(&mut self) -> &mut Histogram<N> {
        self.first_n_total_hist.as_mut().unwrap()
    }
    pub fn get_first_n_per_base_hist_mut(&mut self) -> &mut PerBaseHistogram<N> {
        self.first_n_per_base_hist.as_mut().unwrap()
    }

    pub fn get_last_n_total_hist_mut(&mut self) -> &mut Histogram<N> {
        self.last_n_total_hist.as_mut().unwrap()
    }
    pub fn get_last_n_per_base_hist_mut(&mut self) -> &mut PerBaseHistogram<N> {
        self.last_n_per_base_hist.as_mut().unwrap()
    }

    pub fn merge(&mut self, other: &FinalResult<N>) {
        /*
            pub total_hist: Option<Histogram<N>>,
        pub total_per_base_hist: Option<PerBaseHistogram<N>>,
        pub first_n_total_hist: Option<Histogram<N>>,
        pub first_n_per_base_hist: Option<PerBaseHistogram<N>>,

        pub last_n_total_hist: Option<Histogram<N>>,
        pub last_n_per_base_hist: Option<PerBaseHistogram<N>>,
             */

        if other.total_hist.is_some() {
            self.get_total_hist_mut()
                .merge(other.total_hist.as_ref().unwrap());
        }

        if other.total_per_base_hist.is_some() {
            self.get_total_per_base_hist_mut()
                .merge(other.total_per_base_hist.as_ref().unwrap());
        }

        if other.first_n_total_hist.is_some() {
            self.get_first_n_total_hist_mut()
                .merge(other.first_n_total_hist.as_ref().unwrap());
        }

        if other.first_n_per_base_hist.is_some() {
            self.get_first_n_per_base_hist_mut()
                .merge(other.first_n_per_base_hist.as_ref().unwrap());
        }

        if other.last_n_total_hist.is_some() {
            self.get_last_n_total_hist_mut()
                .merge(other.last_n_total_hist.as_ref().unwrap());
        }

        if other.last_n_per_base_hist.is_some() {
            self.get_last_n_per_base_hist_mut()
                .merge(other.last_n_per_base_hist.as_ref().unwrap());
        }
    }
}

// ── tests ────────────────────────────────────────────────────────
// Helper: turn &[u32] slice into a Vec<u32> so we can call .iter()
fn v(u: &[u32]) -> Vec<u32> {
    u.to_vec()
}

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

    // ── Histogram ────────────────────────────────────────────

    #[test]
    fn new_is_zero() {
        let h: Histogram<100> = Histogram::new();
        assert!(h.dw.iter().all(|&x| x == 0));
        assert!(h.ar.iter().all(|&x| x == 0));
    }

    #[test]
    fn collect_dw_increments() {
        let mut h: Histogram<100> = Histogram::new();
        h.collect_dw(v(&[5, 20, 99]).iter());
        assert_eq!(h.dw[5], 1);
        assert_eq!(h.dw[20], 1);
        assert_eq!(h.dw[99], 1);
    }

    #[test]
    fn collect_dw_clamps_over_max() {
        let mut h: Histogram<100> = Histogram::new();
        h.collect_dw(v(&[200, 500]).iter());
        assert_eq!(h.dw[99], 2); // clamped to last bin
    }

    #[test]
    fn collect_ar_increments() {
        let mut h: Histogram<100> = Histogram::new();
        h.collect_ar(v(&[3, 7]).iter());
        assert_eq!(h.ar[3], 1);
        assert_eq!(h.ar[7], 1);
    }

    #[test]
    fn collect_ar_dw_calls_both() {
        let mut h: Histogram<10> = Histogram::new();
        h.collect_ar_dw(v(&[1, 2]).iter(), v(&[3, 4]).iter());
        assert_eq!(h.ar[1], 1);
        assert_eq!(h.ar[2], 1);
        assert_eq!(h.dw[3], 1);
        assert_eq!(h.dw[4], 1);
    }

    #[test]
    fn merge_adds_pointwise() {
        let mut h1: Histogram<10> = Histogram::new();
        h1.collect_dw(v(&[2, 5]).iter());
        let mut h2: Histogram<10> = Histogram::new();
        h2.collect_dw(v(&[2, 3]).iter());
        h1.merge(&h2);
        assert_eq!(h1.dw[2], 2);
        assert_eq!(h1.dw[3], 1);
        assert_eq!(h1.dw[5], 1);
    }

    #[test]
    fn merge_empty_is_identity() {
        let mut h1: Histogram<10> = Histogram::new();
        h1.collect_dw(v(&[3]).iter());
        let h2: Histogram<10> = Histogram::new();
        h1.merge(&h2);
        assert_eq!(h1.dw[3], 1);
    }

    #[test]
    fn summary_no_data() {
        let h: Histogram<10> = Histogram::new();
        assert!(h.summary().iter().any(|l| l.contains("无数据")));
    }

    #[test]
    fn summary_with_values() {
        let mut h: Histogram<10> = Histogram::new();
        for _ in 0..5 {
            h.collect_dw(v(&[3, 7]).iter());
        }
        assert!(
            h.summary()
                .iter()
                .any(|l| l.contains("非零值") && l.contains("10"))
        );
    }

    // ── PerBaseHistogram ─────────────────────────────────────

    #[test]
    fn new_has_acgt() {
        let pb: PerBaseHistogram<50> = PerBaseHistogram::new();
        for c in "ACGT".chars() {
            assert!(pb.hist_grams.contains_key(&c));
        }
    }

    #[test]
    fn collect_dw_filters_by_base() {
        let mut pb: PerBaseHistogram<100> = PerBaseHistogram::new();
        pb.collect_dw("ATG", &[5, 20, 30]);
        assert_eq!(pb.hist_grams.get(&'A').unwrap().dw[5], 1);
        assert_eq!(pb.hist_grams.get(&'T').unwrap().dw[20], 1);
        assert_eq!(pb.hist_grams.get(&'G').unwrap().dw[30], 1);
    }

    #[test]
    fn collect_ar_filters_by_base() {
        let mut pb: PerBaseHistogram<100> = PerBaseHistogram::new();
        pb.collect_ar("CCGT", &[1, 2, 3]);
        assert_eq!(pb.hist_grams.get(&'C').unwrap().ar[1], 1);
        assert_eq!(pb.hist_grams.get(&'C').unwrap().ar[2], 1);
    }

    #[test]
    fn merge_across_bases() {
        let mut pb1: PerBaseHistogram<10> = PerBaseHistogram::new();
        pb1.collect_dw("A", &[3]);
        let mut pb2: PerBaseHistogram<10> = PerBaseHistogram::new();
        pb2.collect_dw("A", &[5]);
        pb1.merge(&pb2);
        assert_eq!(pb1.hist_grams.get(&'A').unwrap().dw[3], 1);
        assert_eq!(pb1.hist_grams.get(&'A').unwrap().dw[5], 1);
    }

    // ── FinalResult ──────────────────────────────────────────

    #[test]
    fn new_all_sets_total() {
        let r: FinalResult<50> = FinalResult::new_all();
        assert!(r.total_hist.is_some());
        assert!(r.total_per_base_hist.is_some());
        assert!(r.first_n_total_hist.is_none());
    }

    #[test]
    fn new_first_n_last_n_sets_regions() {
        let r: FinalResult<50> = FinalResult::new_first_n_last_n();
        assert!(r.first_n_total_hist.is_some());
        assert!(r.last_n_per_base_hist.is_some());
        assert!(r.total_hist.is_none());
    }

    #[test]
    fn merge_combines_data() {
        let mut a: FinalResult<10> = FinalResult::new_all();
        a.get_total_hist_mut().collect_dw(v(&[2]).iter());
        let mut b: FinalResult<10> = FinalResult::new_all();
        b.get_total_hist_mut().collect_dw(v(&[3]).iter());
        a.merge(&b);
        let total = a.total_hist.unwrap();
        assert_eq!(total.dw[2], 1);
        assert_eq!(total.dw[3], 1);
    }

    #[test]
    fn per_base_collect_ar_dw() {
        let mut pb: PerBaseHistogram<10> = PerBaseHistogram::new();
        pb.collect_ar_dw("AT", &[1, 2], &[3, 4]);
        assert_eq!(pb.hist_grams[&'A'].ar[1], 1);
        assert_eq!(pb.hist_grams[&'T'].ar[2], 1);
        assert_eq!(pb.hist_grams[&'A'].dw[3], 1);
        assert_eq!(pb.hist_grams[&'T'].dw[4], 1);
    }

    #[test]
    fn merge_across_all_regions() {
        let mut a: FinalResult<10> = FinalResult::new_first_n_last_n();
        a.get_first_n_total_hist_mut().collect_dw(v(&[1]).iter());
        a.get_last_n_per_base_hist_mut().collect_dw("C", &[2]);

        let mut b: FinalResult<10> = FinalResult::new_first_n_last_n();
        b.get_last_n_total_hist_mut().collect_dw(v(&[3]).iter());
        b.get_last_n_per_base_hist_mut().collect_dw("C", &[4]);

        a.merge(&b);
        assert_eq!(a.first_n_total_hist.unwrap().dw[1], 1);
        assert_eq!(a.last_n_total_hist.unwrap().dw[3], 1);
        assert_eq!(a.last_n_per_base_hist.unwrap().hist_grams[&'C'].dw[4], 1);
    }

    #[test]
    fn clamp_edge() {
        let mut h: Histogram<5> = Histogram::new();
        h.collect_dw(v(&[0, 4, 999]).iter());
        assert_eq!(h.dw[0], 1);
        assert_eq!(h.dw[4], 2); // 4 and 999 → last bin
    }

    #[test]
    fn empty_is_noop() {
        let mut h: Histogram<5> = Histogram::new();
        h.collect_dw(v(&[]).iter());
        h.collect_ar(v(&[]).iter());
        assert!(h.dw.iter().all(|&x| x == 0));
        assert!(h.ar.iter().all(|&x| x == 0));
    }
}