fastars 0.1.0

//! Read length analysis.
//!
//! This module provides read length distribution analysis.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Read length statistics.
///
/// Tracks length distribution using a HashMap for sparse storage,
/// which is efficient for long reads with highly variable lengths.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LengthStats {
    /// Length distribution: length -> count
    distribution: HashMap<usize, u64>,
    /// Minimum length seen
    min_length: usize,
    /// Maximum length seen
    max_length: usize,
    /// Total bases across all reads
    total_bases: u64,
    /// Total reads processed
    total_reads: u64,
}

impl Default for LengthStats {
    fn default() -> Self {
        Self::new()
    }
}

impl LengthStats {
    /// Create a new length statistics container.
    pub fn new() -> Self {
        Self {
            distribution: HashMap::new(),
            min_length: usize::MAX,
            max_length: 0,
            total_bases: 0,
            total_reads: 0,
        }
    }

    /// Update statistics with a read length.
    ///
    /// This is a hot path - optimized for minimal overhead.
    #[inline]
    pub fn update(&mut self, len: usize) {
        if len == 0 {
            return;
        }

        *self.distribution.entry(len).or_insert(0) += 1;
        self.min_length = self.min_length.min(len);
        self.max_length = self.max_length.max(len);
        self.total_bases += len as u64;
        self.total_reads += 1;
    }

    /// Get mean read length.
    pub fn mean_length(&self) -> f64 {
        if self.total_reads == 0 {
            0.0
        } else {
            self.total_bases as f64 / self.total_reads as f64
        }
    }

    /// Get minimum read length.
    ///
    /// Returns 0 if no reads have been processed.
    pub fn min_length(&self) -> usize {
        if self.total_reads == 0 {
            0
        } else {
            self.min_length
        }
    }

    /// Get maximum read length.
    pub fn max_length(&self) -> usize {
        self.max_length
    }

    /// Get total bases across all reads.
    pub fn total_bases(&self) -> u64 {
        self.total_bases
    }

    /// Get total reads processed.
    pub fn total_reads(&self) -> u64 {
        self.total_reads
    }

    /// Get the length distribution.
    pub fn distribution(&self) -> &HashMap<usize, u64> {
        &self.distribution
    }

    /// Calculate N50 (critical for long reads).
    ///
    /// N50 is the length where 50% of total bases are in reads >= this length.
    pub fn n50(&self) -> usize {
        if self.total_reads == 0 {
            return 0;
        }

        // Get sorted lengths in descending order
        let mut lengths: Vec<(usize, u64)> = self.distribution.iter()
            .map(|(&len, &count)| (len, count))
            .collect();
        lengths.sort_by(|a, b| b.0.cmp(&a.0)); // Sort by length descending

        let half_bases = self.total_bases / 2;
        let mut cumulative_bases: u64 = 0;

        for (len, count) in lengths {
            cumulative_bases += (len as u64) * count;
            if cumulative_bases >= half_bases {
                return len;
            }
        }

        0
    }

    /// Calculate N90.
    ///
    /// N90 is the length where 90% of total bases are in reads >= this length.
    pub fn n90(&self) -> usize {
        self.nx(90)
    }

    /// Calculate Nx for any percentage.
    ///
    /// Nx is the length where x% of total bases are in reads >= this length.
    pub fn nx(&self, x: u8) -> usize {
        if self.total_reads == 0 || x == 0 || x > 100 {
            return 0;
        }

        let mut lengths: Vec<(usize, u64)> = self.distribution.iter()
            .map(|(&len, &count)| (len, count))
            .collect();
        lengths.sort_by(|a, b| b.0.cmp(&a.0));

        let target_bases = (self.total_bases as f64 * (x as f64 / 100.0)) as u64;
        let mut cumulative_bases: u64 = 0;

        for (len, count) in lengths {
            cumulative_bases += (len as u64) * count;
            if cumulative_bases >= target_bases {
                return len;
            }
        }

        0
    }

    /// Calculate median read length.
    pub fn median_length(&self) -> usize {
        if self.total_reads == 0 {
            return 0;
        }

        let mut lengths: Vec<(usize, u64)> = self.distribution.iter()
            .map(|(&len, &count)| (len, count))
            .collect();
        lengths.sort_by_key(|&(len, _)| len);

        // For median, we need the element at position (n+1)/2 (1-indexed)
        let half = self.total_reads / 2;
        let mut cumulative: u64 = 0;

        for (len, count) in lengths {
            cumulative += count;
            if cumulative > half {
                return len;
            }
        }

        0
    }

    /// Merge statistics from another LengthStats instance.
    ///
    /// Used for combining results from multiple workers.
    pub fn merge(&mut self, other: &LengthStats) {
        for (&len, &count) in &other.distribution {
            *self.distribution.entry(len).or_insert(0) += count;
        }

        if other.total_reads > 0 {
            self.min_length = self.min_length.min(other.min_length);
            self.max_length = self.max_length.max(other.max_length);
        }
        self.total_bases += other.total_bases;
        self.total_reads += other.total_reads;
    }

    /// Get length at a specific percentile.
    pub fn percentile(&self, p: f64) -> usize {
        if self.total_reads == 0 || !(0.0..=100.0).contains(&p) {
            return 0;
        }

        let mut lengths: Vec<(usize, u64)> = self.distribution.iter()
            .map(|(&len, &count)| (len, count))
            .collect();
        lengths.sort_by_key(|&(len, _)| len);

        let target = ((self.total_reads as f64 * p) / 100.0) as u64;
        let mut cumulative: u64 = 0;

        for (len, count) in lengths {
            cumulative += count;
            if cumulative >= target {
                return len;
            }
        }

        self.max_length
    }

    /// Create a LengthStats from raw data.
    ///
    /// This is used for converting from FastQcStats to QcStats.
    pub fn from_raw(
        distribution: HashMap<usize, u64>,
        min_length: usize,
        max_length: usize,
        total_bases: u64,
        total_reads: u64,
    ) -> Self {
        Self {
            distribution,
            min_length,
            max_length,
            total_bases,
            total_reads,
        }
    }
}

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

    #[test]
    fn test_length_stats_new() {
        let ls = LengthStats::new();
        assert_eq!(ls.total_reads(), 0);
        assert_eq!(ls.total_bases(), 0);
        assert_eq!(ls.min_length(), 0);
        assert_eq!(ls.max_length(), 0);
    }

    #[test]
    fn test_length_stats_update_single() {
        let mut ls = LengthStats::new();
        ls.update(100);

        assert_eq!(ls.total_reads(), 1);
        assert_eq!(ls.total_bases(), 100);
        assert_eq!(ls.min_length(), 100);
        assert_eq!(ls.max_length(), 100);
        assert!((ls.mean_length() - 100.0).abs() < 0.001);
    }

    #[test]
    fn test_length_stats_update_multiple() {
        let mut ls = LengthStats::new();
        ls.update(100);
        ls.update(200);
        ls.update(300);

        assert_eq!(ls.total_reads(), 3);
        assert_eq!(ls.total_bases(), 600);
        assert_eq!(ls.min_length(), 100);
        assert_eq!(ls.max_length(), 300);
        assert!((ls.mean_length() - 200.0).abs() < 0.001);
    }

    #[test]
    fn test_length_stats_update_zero() {
        let mut ls = LengthStats::new();
        ls.update(0);

        assert_eq!(ls.total_reads(), 0);
        assert_eq!(ls.total_bases(), 0);
    }

    #[test]
    fn test_length_stats_n50() {
        let mut ls = LengthStats::new();
        // 5 reads: 100, 200, 300, 400, 500
        // Total bases: 1500
        // Half = 750
        // 500 = 500 bases (cumulative: 500)
        // 400 = 400 bases (cumulative: 900 >= 750) -> N50 = 400
        ls.update(100);
        ls.update(200);
        ls.update(300);
        ls.update(400);
        ls.update(500);

        assert_eq!(ls.n50(), 400);
    }

    #[test]
    fn test_length_stats_n50_long_reads() {
        let mut ls = LengthStats::new();
        // Simulate long read distribution
        // 10 reads of 1000bp, 5 reads of 10000bp, 1 read of 50000bp
        for _ in 0..10 {
            ls.update(1000);
        }
        for _ in 0..5 {
            ls.update(10000);
        }
        ls.update(50000);

        // Total bases: 10*1000 + 5*10000 + 50000 = 10000 + 50000 + 50000 = 110000
        // Half = 55000
        // 50000 contributes 50000 (cumulative: 50000)
        // 10000 contributes 10000 each
        // After 50000 and one 10000: 60000 >= 55000 -> N50 = 10000
        assert_eq!(ls.n50(), 10000);
    }

    #[test]
    fn test_length_stats_median() {
        let mut ls = LengthStats::new();
        ls.update(100);
        ls.update(200);
        ls.update(300);

        assert_eq!(ls.median_length(), 200);
    }

    #[test]
    fn test_length_stats_percentile() {
        let mut ls = LengthStats::new();
        for i in 1..=100 {
            ls.update(i);
        }

        // 25th percentile should be around 25
        assert!(ls.percentile(25.0) >= 24 && ls.percentile(25.0) <= 26);
        // 75th percentile should be around 75
        assert!(ls.percentile(75.0) >= 74 && ls.percentile(75.0) <= 76);
    }

    #[test]
    fn test_length_stats_merge() {
        let mut ls1 = LengthStats::new();
        ls1.update(100);
        ls1.update(200);

        let mut ls2 = LengthStats::new();
        ls2.update(300);
        ls2.update(400);

        ls1.merge(&ls2);

        assert_eq!(ls1.total_reads(), 4);
        assert_eq!(ls1.total_bases(), 1000);
        assert_eq!(ls1.min_length(), 100);
        assert_eq!(ls1.max_length(), 400);
    }

    #[test]
    fn test_length_stats_merge_empty() {
        let mut ls1 = LengthStats::new();
        ls1.update(100);

        let ls2 = LengthStats::new();

        ls1.merge(&ls2);

        assert_eq!(ls1.total_reads(), 1);
        assert_eq!(ls1.min_length(), 100);
    }

    #[test]
    fn test_length_stats_distribution() {
        let mut ls = LengthStats::new();
        ls.update(100);
        ls.update(100);
        ls.update(200);

        let dist = ls.distribution();
        assert_eq!(dist.get(&100), Some(&2));
        assert_eq!(dist.get(&200), Some(&1));
    }

    #[test]
    fn test_length_stats_nx() {
        let mut ls = LengthStats::new();
        ls.update(100);
        ls.update(200);
        ls.update(300);
        ls.update(400);
        ls.update(500);

        // N90 should be smaller than N50
        assert!(ls.n90() <= ls.n50());
    }

    #[test]
    fn test_length_stats_serialize() {
        let mut ls = LengthStats::new();
        ls.update(100);
        ls.update(200);

        let json = serde_json::to_string(&ls).unwrap();
        let ls2: LengthStats = serde_json::from_str(&json).unwrap();

        assert_eq!(ls.total_reads(), ls2.total_reads());
        assert_eq!(ls.total_bases(), ls2.total_bases());
    }

    #[test]
    fn test_length_stats_empty_n50() {
        let ls = LengthStats::new();
        assert_eq!(ls.n50(), 0);
        assert_eq!(ls.median_length(), 0);
    }
}