genomicframe-core 0.2.0

//! FASTQ statistics computation
//!
//! Streaming statistics collection for sequencing read files.
//! Memory-efficient: processes millions of reads with O(1) memory for basic stats.

use crate::core::GenomicRecordIterator;
use crate::error::Result;
use crate::parallel::Mergeable;
use crate::stats::{CategoryCounter, RunningStats};

use super::reader::{FastqReader, FastqRecord, QualityEncoding};

/// Comprehensive statistics for FASTQ files
#[derive(Debug, Clone)]
pub struct FastqStats {
    /// Total number of reads
    pub total_reads: usize,

    /// Total number of bases sequenced
    pub total_bases: usize,

    /// Read length statistics
    pub length_stats: RunningStats,

    /// Quality score statistics (mean quality per read, then averaged)
    pub mean_quality_stats: RunningStats,

    /// GC content statistics (per read, then averaged)
    pub gc_content_stats: RunningStats,

    /// Distribution of read lengths
    pub length_distribution: CategoryCounter<usize>,

    /// Distribution of mean quality scores (binned)
    pub quality_distribution: CategoryCounter<u8>,

    /// Reads with mean quality >= 20 (common threshold)
    pub high_quality_reads_q20: usize,

    /// Reads with mean quality >= 30 (high quality threshold)
    pub high_quality_reads_q30: usize,

    /// Detected quality encoding
    pub encoding: Option<QualityEncoding>,

    /// N-base count (ambiguous bases)
    pub n_bases: usize,

    /// Base composition
    pub base_counts: CategoryCounter<char>,
}

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

impl FastqStats {
    /// Create a new statistics collector
    pub fn new() -> Self {
        Self {
            total_reads: 0,
            total_bases: 0,
            length_stats: RunningStats::new(),
            mean_quality_stats: RunningStats::new(),
            gc_content_stats: RunningStats::new(),
            length_distribution: CategoryCounter::new(),
            quality_distribution: CategoryCounter::new(),
            high_quality_reads_q20: 0,
            high_quality_reads_q30: 0,
            encoding: None,
            n_bases: 0,
            base_counts: CategoryCounter::new(),
        }
    }

    /// Update statistics with a single record (streaming)
    ///
    /// Optimized for single-pass processing: computes all statistics in one iteration
    /// over sequence and quality data, avoiding repeated allocations and iterations.
    pub fn update(&mut self, record: &FastqRecord) {
        self.total_reads += 1;

        let length = record.len();
        self.total_bases += length;
        self.length_stats.push(length as f64);
        self.length_distribution.increment(length);

        // Single pass over sequence and quality in parallel
        // This eliminates 3 separate iterations (quality, GC, base counts)
        let mut gc_count = 0;
        let mut qual_sum = 0u32;

        // Use bytes() for performance - avoids UTF-8 decoding overhead
        for (seq_byte, qual_byte) in record.sequence.bytes().zip(record.quality.bytes()) {
            // Calculate quality score (no Vec allocation!)
            qual_sum += (qual_byte.saturating_sub(33)) as u32;

            // Count GC bases
            if seq_byte == b'G' || seq_byte == b'C' || seq_byte == b'g' || seq_byte == b'c' {
                gc_count += 1;
            }

            // Track base composition (convert byte to char for counter)
            let base = seq_byte as char;
            self.base_counts.increment(base);

            // Track N bases
            if base == 'N' || base == 'n' {
                self.n_bases += 1;
            }
        }

        // Compute mean quality for this read
        let mean_qual = if length > 0 {
            qual_sum as f64 / length as f64
        } else {
            0.0
        };

        self.mean_quality_stats.push(mean_qual);
        self.quality_distribution.increment(mean_qual as u8);

        // Quality thresholds
        if mean_qual >= 20.0 {
            self.high_quality_reads_q20 += 1;
        }
        if mean_qual >= 30.0 {
            self.high_quality_reads_q30 += 1;
        }

        // Compute GC content
        let gc_content = if length > 0 {
            gc_count as f64 / length as f64
        } else {
            0.0
        };
        self.gc_content_stats.push(gc_content);
    }

    /// Compute statistics from a FASTQ reader (consumes the reader)
    pub fn compute(reader: &mut FastqReader) -> Result<Self> {
        let mut stats = Self::new();

        // Capture encoding from reader after first record is parsed
        while let Some(record) = reader.next_record()? {
            if stats.encoding.is_none() {
                stats.encoding = reader.header().encoding;
            }
            stats.update(&record);
        }

        Ok(stats)
    }

    /// Compute statistics in parallel using multiple threads
    ///
    /// This divides the FASTQ file into chunks and processes them in parallel
    /// using Rayon. Each thread computes partial statistics for its chunk,
    /// then the results are merged.
    ///
    /// # Performance
    /// - For small files (<10 MB): Use `compute()` instead (less overhead)
    /// - For large files (>100 MB): Can achieve near-linear speedup with cores
    ///
    /// # Arguments
    /// - `path`: Path to FASTQ file (can be gzipped)
    /// - `config`: Optional parallel configuration (defaults to all CPUs)
    ///
    /// # Example
    /// ```ignore
    /// use genomicframe_core::formats::fastq::FastqStats;
    /// use genomicframe_core::parallel::ParallelConfig;
    ///
    /// let config = ParallelConfig::new().with_threads(4);
    /// let stats = FastqStats::par_compute("large.fastq", Some(config))?;
    /// ```
    pub fn par_compute<P: AsRef<std::path::Path>>(
        path: P,
        config: Option<crate::parallel::ParallelConfig>,
    ) -> Result<Self> {
        use crate::parallel::{calculate_chunks, find_record_boundary};
        use rayon::prelude::*;
        use std::fs::File;
        use std::io::{BufReader, Seek};

        let config = config.unwrap_or_default();
        let path = path.as_ref();

        // Get file size
        let file_size = std::fs::metadata(path)?.len();

        // For small files, use sequential processing
        if file_size < config.min_chunk_size as u64 {
            let mut reader = FastqReader::from_path(path)?;
            return Self::compute(&mut reader);
        }

        // Calculate chunks
        let chunks = calculate_chunks(file_size, &config);

        // Get or create a cached thread pool with the configured number of threads
        let pool = crate::parallel::get_thread_pool(config.threads())?;

        // Process chunks in parallel using the cached thread pool
        let partial_stats: Vec<Self> = pool.install(|| {
            chunks
                .par_iter()
                .map(|chunk| -> Result<Self> {
                // Open file for this thread
                let file = File::open(path)?;
                let mut reader = BufReader::new(file);

                // Find the actual start of the first complete record in this chunk
                let start_pos = find_record_boundary(&mut reader, chunk.start)?;

                let start_pos = match start_pos {
                    Some(pos) => pos,
                    None => return Ok(Self::new()), // No records in this chunk
                };

                // Create FASTQ reader starting at this position
                reader.seek(std::io::SeekFrom::Start(start_pos))?;
                let mut fastq_reader = FastqReader::from_buffer(reader);

                // Process records until we hit the chunk end
                let mut stats = Self::new();
                let mut current_pos = start_pos;

                while let Some(record) = fastq_reader.next_record()? {
                    // Capture encoding from the first record
                    if stats.encoding.is_none() {
                        stats.encoding = fastq_reader.header().encoding;
                    }

                    stats.update(&record);

                    // Estimate current position (each FASTQ record is ~4 lines)
                    // This is approximate but avoids expensive tell() calls
                    current_pos += record.len() as u64 * 2 + 40; // Rough estimate

                    // Stop when we reach the end of our chunk
                    if current_pos >= chunk.end {
                        break;
                    }
                }

                Ok(stats)
            })
            .collect::<Result<Vec<_>>>()
        })?; // Close pool.install()

        // Merge all partial statistics
        Self::merge_all(partial_stats)
            .ok_or_else(|| crate::error::Error::InvalidInput("No statistics computed".to_string()))
    }

    /// Mean read length
    pub fn mean_length(&self) -> Option<f64> {
        self.length_stats.mean()
    }

    /// Standard deviation of read length
    pub fn std_length(&self) -> Option<f64> {
        self.length_stats.std_dev()
    }

    /// Min read length
    pub fn min_length(&self) -> Option<f64> {
        self.length_stats.min()
    }

    /// Max read length
    pub fn max_length(&self) -> Option<f64> {
        self.length_stats.max()
    }

    /// Mean of per-read mean quality scores
    pub fn mean_quality(&self) -> Option<f64> {
        self.mean_quality_stats.mean()
    }

    /// Standard deviation of quality scores
    pub fn std_quality(&self) -> Option<f64> {
        self.mean_quality_stats.std_dev()
    }

    /// Mean GC content across all reads
    pub fn mean_gc_content(&self) -> Option<f64> {
        self.gc_content_stats.mean()
    }

    /// Standard deviation of GC content
    pub fn std_gc_content(&self) -> Option<f64> {
        self.gc_content_stats.std_dev()
    }

    /// Percentage of high quality reads (Q >= 20)
    pub fn percent_q20(&self) -> f64 {
        if self.total_reads == 0 {
            0.0
        } else {
            (self.high_quality_reads_q20 as f64 / self.total_reads as f64) * 100.0
        }
    }

    /// Percentage of high quality reads (Q >= 30)
    pub fn percent_q30(&self) -> f64 {
        if self.total_reads == 0 {
            0.0
        } else {
            (self.high_quality_reads_q30 as f64 / self.total_reads as f64) * 100.0
        }
    }

    /// Percentage of N bases (ambiguous)
    pub fn percent_n(&self) -> f64 {
        if self.total_bases == 0 {
            0.0
        } else {
            (self.n_bases as f64 / self.total_bases as f64) * 100.0
        }
    }

    /// Get base composition as percentages
    pub fn base_composition(&self) -> BaseComposition {
        let total = self.total_bases as f64;
        if total == 0.0 {
            return BaseComposition::default();
        }

        BaseComposition {
            a_percent: (self.base_counts.get(&'A') as f64 / total) * 100.0,
            c_percent: (self.base_counts.get(&'C') as f64 / total) * 100.0,
            g_percent: (self.base_counts.get(&'G') as f64 / total) * 100.0,
            t_percent: (self.base_counts.get(&'T') as f64 / total) * 100.0,
            n_percent: self.percent_n(),
        }
    }

    /// Print a human-readable summary of statistics
    pub fn print_summary(&self) {
        println!("=== FASTQ Statistics ===\n");

        println!("Basic Counts:");
        println!("  Total reads:     {:>12}", self.total_reads);
        println!("  Total bases:     {:>12}", self.total_bases);
        println!();

        println!("Read Length:");
        println!("  Mean:            {:>12.2}", self.mean_length().unwrap_or(0.0));
        println!("  Std Dev:         {:>12.2}", self.std_length().unwrap_or(0.0));
        println!(
            "  Min:             {:>12.0}",
            self.min_length().unwrap_or(0.0)
        );
        println!(
            "  Max:             {:>12.0}",
            self.max_length().unwrap_or(0.0)
        );
        println!();

        println!("Quality Scores:");
        if let Some(encoding) = self.encoding {
            println!("  Encoding:        {:>12?}", encoding);
        }
        println!(
            "  Mean quality:    {:>12.2}",
            self.mean_quality().unwrap_or(0.0)
        );
        println!(
            "  Std Dev:         {:>12.2}",
            self.std_quality().unwrap_or(0.0)
        );
        println!("  Q >= 20:         {:>11.1}%", self.percent_q20());
        println!("  Q >= 30:         {:>11.1}%", self.percent_q30());
        println!();

        println!("Base Composition:");
        let composition = self.base_composition();
        println!("  A:               {:>11.1}%", composition.a_percent);
        println!("  C:               {:>11.1}%", composition.c_percent);
        println!("  G:               {:>11.1}%", composition.g_percent);
        println!("  T:               {:>11.1}%", composition.t_percent);
        println!("  N (ambiguous):   {:>11.1}%", composition.n_percent);
        println!();

        println!("GC Content:");
        println!(
            "  Mean:            {:>11.1}%",
            self.mean_gc_content().unwrap_or(0.0) * 100.0
        );
        println!(
            "  Std Dev:         {:>11.1}%",
            self.std_gc_content().unwrap_or(0.0) * 100.0
        );
        println!();

        // Top length bins (sorted by count descending)
        println!("Top Read Lengths:");
        let mut length_vec: Vec<_> = self.length_distribution.iter().collect();
        length_vec.sort_by(|a, b| b.1.cmp(a.1)); // Sort by count descending
        for (length, count) in length_vec.iter().take(10) {
            let percent = (**count as f64 / self.total_reads as f64) * 100.0;
            println!("  {:>4}bp: {:>10} ({:>5.1}%)", length, count, percent);
        }
        println!();

        // Quality distribution (binned, sorted by quality score)
        println!("Quality Distribution (binned):");
        let mut qual_vec: Vec<_> = self.quality_distribution.iter().collect();
        qual_vec.sort_by_key(|a| a.0); // Sort by quality score ascending
        for (qual_bin, count) in qual_vec.iter().take(10) {
            let percent = (**count as f64 / self.total_reads as f64) * 100.0;
            println!("  Q{:>2}: {:>10} ({:>5.1}%)", qual_bin, count, percent);
        }
    }
}

/// Base composition percentages
#[derive(Debug, Clone, Default)]
pub struct BaseComposition {
    pub a_percent: f64,
    pub c_percent: f64,
    pub g_percent: f64,
    pub t_percent: f64,
    pub n_percent: f64,
}

/// Implementation of Mergeable for parallel statistics computation
///
/// Merges FASTQ statistics from multiple threads/chunks.
/// All statistics are correctly combined using the underlying
/// Mergeable implementations for RunningStats and CategoryCounter.
impl Mergeable for FastqStats {
    fn merge(&mut self, other: Self) {
        // Merge simple counters
        self.total_reads += other.total_reads;
        self.total_bases += other.total_bases;
        self.high_quality_reads_q20 += other.high_quality_reads_q20;
        self.high_quality_reads_q30 += other.high_quality_reads_q30;
        self.n_bases += other.n_bases;

        // Merge running statistics
        self.length_stats.merge(other.length_stats);
        self.mean_quality_stats.merge(other.mean_quality_stats);
        self.gc_content_stats.merge(other.gc_content_stats);

        // Merge categorical data
        self.length_distribution.merge(other.length_distribution);
        self.quality_distribution.merge(other.quality_distribution);
        self.base_counts.merge(other.base_counts);

        // Merge encoding (prefer first non-None value)
        if self.encoding.is_none() {
            self.encoding = other.encoding;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    fn create_test_fastq() -> NamedTempFile {
        let mut temp_file = NamedTempFile::new().unwrap();

        // Write 3 reads with different qualities
        writeln!(temp_file, "@READ1").unwrap();
        writeln!(temp_file, "ACGTACGT").unwrap(); // 8bp, 50% GC
        writeln!(temp_file, "+").unwrap();
        writeln!(temp_file, "IIIIIIII").unwrap(); // High quality (Q40)

        writeln!(temp_file, "@READ2").unwrap();
        writeln!(temp_file, "AAAATTTT").unwrap(); // 8bp, 0% GC
        writeln!(temp_file, "+").unwrap();
        writeln!(temp_file, "########").unwrap(); // Medium quality (Q7)

        writeln!(temp_file, "@READ3").unwrap();
        writeln!(temp_file, "GGGGCCCCNNNN").unwrap(); // 12bp, 66% GC (excluding N)
        writeln!(temp_file, "+").unwrap();
        writeln!(temp_file, "IIIIIIIIIIII").unwrap(); // High quality

        temp_file.flush().unwrap();
        temp_file
    }

    #[test]
    fn test_fastq_stats_basic() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;
        let stats = FastqStats::compute(&mut reader)?;

        assert_eq!(stats.total_reads, 3);
        assert_eq!(stats.total_bases, 28); // 8 + 8 + 12

        Ok(())
    }

    #[test]
    fn test_fastq_stats_length() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;
        let stats = FastqStats::compute(&mut reader)?;

        // Mean length: (8 + 8 + 12) / 3 = 9.33
        let mean = stats.mean_length().unwrap();
        assert!((mean - 9.33).abs() < 0.1);

        assert_eq!(stats.min_length().unwrap(), 8.0);
        assert_eq!(stats.max_length().unwrap(), 12.0);

        Ok(())
    }

    #[test]
    fn test_fastq_stats_quality() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;
        let stats = FastqStats::compute(&mut reader)?;

        // At least 2 out of 3 reads should be Q30+ (READ1 and READ3)
        assert!(stats.high_quality_reads_q30 >= 2);
        assert!(stats.percent_q30() >= 60.0);

        Ok(())
    }

    #[test]
    fn test_fastq_stats_gc_content() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;
        let stats = FastqStats::compute(&mut reader)?;

        // READ1: 4/8 = 50%, READ2: 0/8 = 0%, READ3: 8/12 = 66.7%
        // Mean: (0.5 + 0.0 + 0.667) / 3 = 0.389 = 38.9%
        let mean_gc = stats.mean_gc_content().unwrap();
        assert!((mean_gc - 0.389).abs() < 0.01);

        Ok(())
    }

    #[test]
    fn test_fastq_stats_base_composition() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;
        let stats = FastqStats::compute(&mut reader)?;

        let composition = stats.base_composition();

        // READ1: 2A, 2C, 2G, 2T = 8 bases
        // READ2: 4A, 0C, 0G, 4T = 8 bases
        // READ3: 0A, 4C, 4G, 0T, 4N = 12 bases
        // Total: 6A, 6C, 6G, 6T, 4N = 28 bases
        assert!((composition.a_percent - 21.4).abs() < 1.0); // 6/28 ≈ 21.4%
        assert!((composition.c_percent - 21.4).abs() < 1.0); // 6/28 ≈ 21.4%
        assert!((composition.g_percent - 21.4).abs() < 1.0); // 6/28 ≈ 21.4%
        assert!((composition.t_percent - 21.4).abs() < 1.0); // 6/28 ≈ 21.4%
        assert!((composition.n_percent - 14.3).abs() < 1.0); // 4/28 ≈ 14.3%

        Ok(())
    }

    #[test]
    fn test_fastq_stats_n_bases() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;
        let stats = FastqStats::compute(&mut reader)?;

        assert_eq!(stats.n_bases, 4); // 4 N's in READ3
        assert!((stats.percent_n() - 14.3).abs() < 1.0); // 4/28 ≈ 14.3%

        Ok(())
    }

    #[test]
    fn test_fastq_stats_streaming() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;

        // Manually stream and update
        let mut stats = FastqStats::new();
        let mut count = 0;
        while let Some(record) = reader.next_record()? {
            stats.update(&record);
            count += 1;
        }

        assert_eq!(count, 3);
        assert_eq!(stats.total_reads, 3);

        Ok(())
    }

    #[test]
    fn test_fastq_stats_merge() -> Result<()> {
        let temp_file = create_test_fastq();
        let mut reader = FastqReader::from_path(temp_file.path())?;

        // Create two partial stats
        let mut stats1 = FastqStats::new();
        let mut stats2 = FastqStats::new();

        // Add first record to stats1
        if let Some(record) = reader.next_record()? {
            stats1.update(&record);
        }

        // Add remaining records to stats2
        while let Some(record) = reader.next_record()? {
            stats2.update(&record);
        }

        // Verify partial stats
        assert_eq!(stats1.total_reads, 1);
        assert_eq!(stats2.total_reads, 2);

        // Merge
        stats1.merge(stats2);

        // Verify merged stats
        assert_eq!(stats1.total_reads, 3);
        assert_eq!(stats1.total_bases, 28); // 8 + 8 + 12

        Ok(())
    }
}