fastars 0.1.0

//! Quality-based trimming.
//!
//! This module provides trimming based on quality scores,
//! removing low-quality bases from read ends using a sliding window algorithm.

use super::TrimResult;

/// Configuration for quality-based trimming.
#[derive(Debug, Clone)]
pub struct QualityTrimConfig {
    /// Size of the sliding window for quality calculation.
    pub window_size: usize,
    /// Quality threshold (Phred score). Windows with mean quality below this are trimmed.
    pub threshold: u8,
    /// Whether to trim from the 5' (front) end.
    pub cut_front: bool,
    /// Whether to trim from the 3' (tail) end.
    pub cut_tail: bool,
    /// Whether to trim from the first low-quality position to the end.
    pub cut_right: bool,
    /// Window size for cut_right (if None, uses window_size).
    pub right_window_size: Option<usize>,
    /// Quality threshold for cut_right (if None, uses threshold).
    pub right_threshold: Option<u8>,
    /// Window size for cut_tail (if None, uses window_size).
    pub tail_window_size: Option<usize>,
    /// Quality threshold for cut_tail (if None, uses threshold).
    pub tail_threshold: Option<u8>,
    /// Window size for cut_front (if None, uses window_size).
    pub front_window_size: Option<usize>,
    /// Quality threshold for cut_front (if None, uses threshold).
    pub front_threshold: Option<u8>,
}

impl Default for QualityTrimConfig {
    fn default() -> Self {
        Self {
            window_size: 4,
            threshold: 15,
            cut_front: false,
            cut_tail: true,
            cut_right: false,
            right_window_size: None,
            right_threshold: None,
            tail_window_size: None,
            tail_threshold: None,
            front_window_size: None,
            front_threshold: None,
        }
    }
}

impl QualityTrimConfig {
    /// Create a new quality trim config with default settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create config for short read mode.
    pub fn short_read() -> Self {
        Self {
            window_size: 4,
            threshold: 15,
            cut_front: false,
            cut_tail: true,
            cut_right: false,
            right_window_size: None,
            right_threshold: None,
            tail_window_size: None,
            tail_threshold: None,
            front_window_size: None,
            front_threshold: None,
        }
    }

    /// Create config for long read mode.
    pub fn long_read() -> Self {
        Self {
            window_size: 20,
            threshold: 7,
            cut_front: false,
            cut_tail: true,
            cut_right: false,
            right_window_size: None,
            right_threshold: None,
            tail_window_size: None,
            tail_threshold: None,
            front_window_size: None,
            front_threshold: None,
        }
    }

    /// Set window size.
    pub fn with_window_size(mut self, size: usize) -> Self {
        self.window_size = size;
        self
    }

    /// Set quality threshold.
    pub fn with_threshold(mut self, threshold: u8) -> Self {
        self.threshold = threshold;
        self
    }

    /// Enable/disable front trimming.
    pub fn with_cut_front(mut self, enabled: bool) -> Self {
        self.cut_front = enabled;
        self
    }

    /// Enable/disable tail trimming.
    pub fn with_cut_tail(mut self, enabled: bool) -> Self {
        self.cut_tail = enabled;
        self
    }

    /// Enable/disable right trimming.
    pub fn with_cut_right(mut self, enabled: bool) -> Self {
        self.cut_right = enabled;
        self
    }

    /// Set window size for cut_right.
    pub fn with_right_window_size(mut self, size: usize) -> Self {
        self.right_window_size = Some(size);
        self
    }

    /// Set quality threshold for cut_right.
    pub fn with_right_threshold(mut self, threshold: u8) -> Self {
        self.right_threshold = Some(threshold);
        self
    }

    /// Set window size for cut_tail.
    pub fn with_tail_window_size(mut self, size: usize) -> Self {
        self.tail_window_size = Some(size);
        self
    }

    /// Set quality threshold for cut_tail.
    pub fn with_tail_threshold(mut self, threshold: u8) -> Self {
        self.tail_threshold = Some(threshold);
        self
    }

    /// Set window size for cut_front.
    pub fn with_front_window_size(mut self, size: usize) -> Self {
        self.front_window_size = Some(size);
        self
    }

    /// Set quality threshold for cut_front.
    pub fn with_front_threshold(mut self, threshold: u8) -> Self {
        self.front_threshold = Some(threshold);
        self
    }
}

/// Perform sliding window quality trimming.
///
/// Uses a sliding window approach where the window moves along the sequence,
/// calculating mean quality at each position. When the mean quality falls below
/// the threshold, the read is trimmed from that position.
///
/// # Arguments
/// * `qual` - Quality scores (Phred+33 encoded, ASCII 33-126)
/// * `config` - Trimming configuration
///
/// # Returns
/// A `TrimResult` containing the start and end indices of the good-quality region.
///
/// # Example
/// ```
/// use fastars::trim::quality::{sliding_window_trim, QualityTrimConfig};
///
/// let qual = b"IIIIIIIII55555!!!!!"; // High quality followed by low
/// let config = QualityTrimConfig::default();
/// let result = sliding_window_trim(qual, &config);
/// assert!(result.end < qual.len()); // Some trimming occurred
/// ```
pub fn sliding_window_trim(qual: &[u8], config: &QualityTrimConfig) -> TrimResult {
    let len = qual.len();

    // Handle edge cases
    if len == 0 {
        return TrimResult::empty();
    }

    if config.window_size == 0 || config.window_size > len {
        return TrimResult::full(len);
    }

    let mut start = 0;
    let mut end = len;

    // Trim from 3' end (tail)
    if config.cut_tail {
        let window_size = config.tail_window_size.unwrap_or(config.window_size);
        let threshold = config.tail_threshold.unwrap_or(config.threshold);
        end = trim_tail(qual, window_size, threshold);
    }

    // Trim from 5' end (front)
    if config.cut_front && end > 0 {
        let window_size = config.front_window_size.unwrap_or(config.window_size);
        let threshold = config.front_threshold.unwrap_or(config.threshold);
        start = trim_front(qual, end, window_size, threshold);
    }

    // Trim from the first low-quality window to the end (cut_right)
    if config.cut_right && end > 0 {
        let window_size = config.right_window_size.unwrap_or(config.window_size);
        let threshold = config.right_threshold.unwrap_or(config.threshold);
        let right_end = trim_right(qual, window_size, threshold);
        end = end.min(right_end);
    }

    if start >= end {
        return TrimResult::empty();
    }

    TrimResult::new(start, end)
}

/// Trim from the 3' end using sliding window.
///
/// This implements fastp's exact algorithm: when a good window is found,
/// keep only up to the start position of that window + 1.
/// (fastp does post-loop adjustment: t = t - w + 1, then rlen = t + 1)
fn trim_tail(qual: &[u8], window_size: usize, threshold: u8) -> usize {
    let len = qual.len();
    if len < window_size {
        return len;
    }

    // Calculate initial window sum at the end
    let mut window_sum: u32 = qual[len - window_size..]
        .iter()
        .map(|&q| phred_score(q) as u32)
        .sum();

    let threshold_sum = (threshold as u32) * (window_size as u32);

    // If last window is good, no trimming needed
    if window_sum >= threshold_sum {
        return len;
    }

    // Slide window from end towards start
    let mut end = len - window_size;

    for i in (0..len - window_size).rev() {
        // Slide window: add the new base, remove the one that exits
        window_sum = window_sum + phred_score(qual[i]) as u32
                   - phred_score(qual[i + window_size]) as u32;

        if window_sum >= threshold_sum {
            // Found a good window - fastp-compatible: keep up to window start + 1
            // (fastp does post-loop: t = t - w + 1, then rlen = t + 1)
            return i + 1;
        }
        end = i;
    }

    // All windows are bad
    end
}

/// Trim from the 5' end using sliding window.
fn trim_front(qual: &[u8], end: usize, window_size: usize, threshold: u8) -> usize {
    if end < window_size {
        return 0;
    }

    // Calculate initial window sum at the start
    let mut window_sum: u32 = qual[..window_size]
        .iter()
        .map(|&q| phred_score(q) as u32)
        .sum();

    let threshold_sum = (threshold as u32) * (window_size as u32);

    // If first window is good, no trimming needed
    if window_sum >= threshold_sum {
        return 0;
    }

    // Slide window from start towards end
    for i in 1..=end - window_size {
        // Slide window: add the new base, remove the one that exits
        window_sum = window_sum + phred_score(qual[i + window_size - 1]) as u32
                   - phred_score(qual[i - 1]) as u32;

        if window_sum >= threshold_sum {
            return i;
        }
    }

    // All windows are bad
    end
}

/// Trim from the first position where sliding window mean drops below threshold.
/// This scans left to right and trims from the first low-quality window to the end.
fn trim_right(qual: &[u8], window_size: usize, threshold: u8) -> usize {
    let len = qual.len();
    if len < window_size {
        return len;
    }

    // Calculate initial window sum at the start
    let mut window_sum: u32 = qual[..window_size]
        .iter()
        .map(|&q| phred_score(q) as u32)
        .sum();

    let threshold_sum = (threshold as u32) * (window_size as u32);

    // Check first window
    if window_sum < threshold_sum {
        return 0;
    }

    // Slide window from left to right
    for i in 1..=len - window_size {
        // Slide window: add the new base, remove the one that exits
        window_sum = window_sum + phred_score(qual[i + window_size - 1]) as u32
                   - phred_score(qual[i - 1]) as u32;

        if window_sum < threshold_sum {
            // Found low quality window, trim from here
            return i;
        }
    }

    // All windows are good
    len
}

/// Convert Phred+33 ASCII to Phred score.
#[inline]
fn phred_score(ascii: u8) -> u8 {
    ascii.saturating_sub(33)
}

/// Calculate mean quality score for a slice.
#[inline]
pub fn mean_quality(qual: &[u8]) -> f64 {
    if qual.is_empty() {
        return 0.0;
    }
    let sum: u32 = qual.iter().map(|&q| phred_score(q) as u32).sum();
    sum as f64 / qual.len() as f64
}

// Legacy type alias for compatibility
pub type QualityTrimmer = QualityTrimConfig;

/// Configuration for quality masking.
#[derive(Debug, Clone)]
pub struct MaskConfig {
    /// Window size for quality masking.
    pub window_size: usize,
    /// Mean quality threshold for masking.
    pub threshold: u8,
}

impl MaskConfig {
    /// Create a new mask config.
    pub fn new(window_size: usize, threshold: u8) -> Self {
        Self {
            window_size,
            threshold,
        }
    }
}

/// Configuration for read breaking.
#[derive(Debug, Clone)]
pub struct BreakConfig {
    /// Window size for read breaking.
    pub window_size: usize,
    /// Mean quality threshold for breaking.
    pub threshold: u8,
}

impl BreakConfig {
    /// Create a new break config.
    pub fn new(window_size: usize, threshold: u8) -> Self {
        Self {
            window_size,
            threshold,
        }
    }
}

/// Mask low quality regions in sequence and quality strings.
///
/// Detects low quality regions using a sliding window approach and replaces
/// bases in those regions with 'N' and quality with '!' (Phred 0).
///
/// # Arguments
/// * `seq` - Mutable sequence slice
/// * `qual` - Quality scores (Phred+33 encoded)
/// * `config` - Masking configuration
///
/// # Returns
/// Number of bases masked
pub fn mask_low_quality(seq: &mut [u8], qual: &[u8], config: &MaskConfig) -> usize {
    let len = seq.len();
    if len == 0 || len != qual.len() {
        return 0;
    }

    if config.window_size == 0 || config.window_size > len {
        return 0;
    }

    let threshold_sum = (config.threshold as u32) * (config.window_size as u32);
    let mut masked_count = 0;

    // Slide window across the sequence
    let mut window_sum: u32 = qual[..config.window_size]
        .iter()
        .map(|&q| phred_score(q) as u32)
        .sum();

    // Check first window
    if window_sum < threshold_sum {
        for i in 0..config.window_size {
            seq[i] = b'N';
            masked_count += 1;
        }
    }

    // Slide window from left to right
    for i in 1..=len - config.window_size {
        // Update window sum
        window_sum = window_sum + phred_score(qual[i + config.window_size - 1]) as u32
            - phred_score(qual[i - 1]) as u32;

        // If mean quality is below threshold, mask this window
        if window_sum < threshold_sum {
            for j in i..i + config.window_size {
                if seq[j] != b'N' {
                    seq[j] = b'N';
                    masked_count += 1;
                }
            }
        }
    }

    masked_count
}

/// Break a read at low quality regions.
///
/// Detects low quality regions using a sliding window approach and identifies
/// break points where the read should be split.
///
/// # Arguments
/// * `qual` - Quality scores (Phred+33 encoded)
/// * `config` - Breaking configuration
///
/// # Returns
/// Vector of (start, end) pairs representing high-quality segments
pub fn break_at_low_quality(qual: &[u8], config: &BreakConfig) -> Vec<(usize, usize)> {
    let len = qual.len();
    if len == 0 {
        return Vec::new();
    }

    if config.window_size == 0 || config.window_size > len {
        return vec![(0, len)];
    }

    let threshold_sum = (config.threshold as u32) * (config.window_size as u32);
    let mut low_quality_regions = Vec::new();

    // Slide window across the sequence to find low quality regions
    let mut window_sum: u32 = qual[..config.window_size]
        .iter()
        .map(|&q| phred_score(q) as u32)
        .sum();

    // Track low quality regions
    let mut in_low_quality = window_sum < threshold_sum;
    let mut low_quality_start = if in_low_quality { 0 } else { usize::MAX };

    for i in 1..=len - config.window_size {
        // Update window sum
        window_sum = window_sum + phred_score(qual[i + config.window_size - 1]) as u32
            - phred_score(qual[i - 1]) as u32;

        let is_low = window_sum < threshold_sum;

        if is_low && !in_low_quality {
            // Start of low quality region
            low_quality_start = i;
            in_low_quality = true;
        } else if !is_low && in_low_quality {
            // End of low quality region
            low_quality_regions.push((low_quality_start, i + config.window_size - 1));
            in_low_quality = false;
        }
    }

    // Handle case where read ends in low quality region
    if in_low_quality {
        low_quality_regions.push((low_quality_start, len));
    }

    // Convert low quality regions to high quality segments
    if low_quality_regions.is_empty() {
        return vec![(0, len)];
    }

    let mut segments = Vec::new();
    let mut current_pos = 0;

    for (lq_start, lq_end) in low_quality_regions {
        if current_pos < lq_start {
            segments.push((current_pos, lq_start));
        }
        current_pos = lq_end;
    }

    // Add final segment if there's high quality region at the end
    if current_pos < len {
        segments.push((current_pos, len));
    }

    segments
}

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

    fn make_qual(scores: &[u8]) -> Vec<u8> {
        scores.iter().map(|&s| s + 33).collect()
    }

    #[test]
    fn test_quality_trim_config_default() {
        let config = QualityTrimConfig::default();
        assert_eq!(config.window_size, 4);
        assert_eq!(config.threshold, 15);
        assert!(!config.cut_front);
        assert!(config.cut_tail);
    }

    #[test]
    fn test_quality_trim_config_short_read() {
        let config = QualityTrimConfig::short_read();
        assert_eq!(config.window_size, 4);
        assert_eq!(config.threshold, 15);
    }

    #[test]
    fn test_quality_trim_config_long_read() {
        let config = QualityTrimConfig::long_read();
        assert_eq!(config.window_size, 20);
        assert_eq!(config.threshold, 7);
    }

    #[test]
    fn test_sliding_window_empty() {
        let config = QualityTrimConfig::default();
        let result = sliding_window_trim(&[], &config);
        assert!(result.is_empty());
    }

    #[test]
    fn test_sliding_window_all_high_quality() {
        let qual = make_qual(&[30, 30, 30, 30, 30, 30, 30, 30]);
        let config = QualityTrimConfig::default();
        let result = sliding_window_trim(&qual, &config);
        assert_eq!(result.start, 0);
        assert_eq!(result.end, 8);
    }

    #[test]
    fn test_sliding_window_all_low_quality() {
        let qual = make_qual(&[2, 2, 2, 2, 2, 2, 2, 2]);
        let config = QualityTrimConfig::default();
        let result = sliding_window_trim(&qual, &config);
        // All windows are bad, should trim to beginning
        assert!(result.is_empty() || result.len() < qual.len());
    }

    #[test]
    fn test_sliding_window_tail_trim() {
        // High quality at start, low at end
        let qual = make_qual(&[30, 30, 30, 30, 30, 30, 5, 5, 5, 5]);
        let config = QualityTrimConfig::default().with_cut_tail(true).with_cut_front(false);
        let result = sliding_window_trim(&qual, &config);
        assert_eq!(result.start, 0);
        // Tail should be trimmed - the exact position depends on window algorithm
        assert!(result.end < qual.len());
    }

    #[test]
    fn test_sliding_window_front_trim() {
        // Low quality at start, high at end
        let qual = make_qual(&[5, 5, 5, 5, 30, 30, 30, 30, 30, 30]);
        let config = QualityTrimConfig::default().with_cut_front(true).with_cut_tail(false);
        let result = sliding_window_trim(&qual, &config);
        // Front should be trimmed - start should be > 0
        assert!(result.start > 0);
        assert_eq!(result.end, 10);
    }

    #[test]
    fn test_sliding_window_both_ends() {
        // Low quality at both ends
        let qual = make_qual(&[5, 5, 5, 30, 30, 30, 30, 5, 5, 5]);
        let config = QualityTrimConfig::default().with_cut_front(true).with_cut_tail(true);
        let result = sliding_window_trim(&qual, &config);
        // Both ends should be trimmed - result should be shorter
        assert!(result.len() < qual.len());
    }

    #[test]
    fn test_sliding_window_shorter_than_window() {
        let qual = make_qual(&[30, 30]);
        let config = QualityTrimConfig::default().with_window_size(4);
        let result = sliding_window_trim(&qual, &config);
        assert_eq!(result.start, 0);
        assert_eq!(result.end, 2);
    }

    #[test]
    fn test_mean_quality() {
        let qual = make_qual(&[30, 30, 30, 30]);
        let mean = mean_quality(&qual);
        assert!((mean - 30.0).abs() < 0.01);
    }

    #[test]
    fn test_mean_quality_empty() {
        let mean = mean_quality(&[]);
        assert_eq!(mean, 0.0);
    }

    #[test]
    fn test_phred_score() {
        assert_eq!(phred_score(b'!'), 0);  // ASCII 33 = Phred 0
        assert_eq!(phred_score(b'I'), 40); // ASCII 73 = Phred 40
        assert_eq!(phred_score(b'~'), 93); // ASCII 126 = Phred 93
    }

    #[test]
    fn test_config_builder() {
        let config = QualityTrimConfig::new()
            .with_window_size(10)
            .with_threshold(20)
            .with_cut_front(true)
            .with_cut_tail(false);
        assert_eq!(config.window_size, 10);
        assert_eq!(config.threshold, 20);
        assert!(config.cut_front);
        assert!(!config.cut_tail);
    }
}