rustalign_aligner/

paired_end.rs

//! Paired-end alignment policy and utilities
//!
//! This module implements paired-end alignment support, including:
//! - Orientation policies (FR, RF, FF, RR)
//! - Insert size constraints
//! - Mate finding logic
//! - Concordant/discordant pair classification

use rustalign_common::Score;

/// Paired-end orientation policy
///
/// Defines the expected orientation of mate pairs.
/// This matches standard paired-end sequencing orientation types.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PePolicy {
    /// Forward-Reverse: R1 forward, R2 reverse (most common for Illumina)
    #[default]
    FR,
    /// Reverse-Forward: R1 reverse, R2 forward
    RF,
    /// Forward-Forward: both mates in forward orientation
    FF,
    /// Reverse-Reverse: both mates in reverse orientation
    RR,
}

/// Paired-end constraints
///
/// Defines the constraints for valid paired-end alignments.
#[derive(Debug, Clone, Copy)]
pub struct PeConstraints {
    /// Minimum fragment/insert size (outer distance between mates)
    pub min_insert: u32,
    /// Maximum fragment/insert size
    pub max_insert: u32,
    /// Orientation policy
    pub policy: PePolicy,
    /// Whether to allow overlapping mates
    pub allow_overlap: bool,
    /// Whether to allow dovetailing (one mate extends past the other)
    pub allow_dovetail: bool,
    /// Whether to allow one mate containing the other
    pub allow_contain: bool,
}

impl Default for PeConstraints {
    fn default() -> Self {
        Self {
            min_insert: 0,
            max_insert: 500,
            policy: PePolicy::FR,
            allow_overlap: true,
            allow_dovetail: false,
            allow_contain: false,
        }
    }
}

impl PeConstraints {
    /// Create new paired-end constraints
    pub fn new(min_insert: u32, max_insert: u32, policy: PePolicy) -> Self {
        Self {
            min_insert,
            max_insert,
            policy,
            ..Default::default()
        }
    }

    /// Check if a pair of alignments is concordant (properly paired)
    ///
    /// # Arguments
    /// * `pos1` - Position of mate 1 (0-based)
    /// * `len1` - Length of mate 1
    /// * `strand1` - Strand of mate 1 (true = forward, false = reverse)
    /// * `chr1` - Chromosome of mate 1
    /// * `pos2` - Position of mate 2 (0-based)
    /// * `len2` - Length of mate 2
    /// * `strand2` - Strand of mate 2
    /// * `chr2` - Chromosome of mate 2
    ///
    /// # Returns
    /// `Some(insert_size)` if concordant, `None` if discordant
    #[allow(clippy::too_many_arguments)]
    pub fn is_concordant(
        &self,
        pos1: u64,
        len1: u32,
        strand1: bool,
        chr1: u32,
        pos2: u64,
        len2: u32,
        strand2: bool,
        chr2: u32,
    ) -> Option<i32> {
        // Must be on same chromosome
        if chr1 != chr2 {
            return None;
        }

        // Check orientation
        let orientation_ok = match self.policy {
            PePolicy::FR => strand1 && !strand2,  // R1 forward, R2 reverse
            PePolicy::RF => !strand1 && strand2,  // R1 reverse, R2 forward
            PePolicy::FF => strand1 && strand2,   // Both forward
            PePolicy::RR => !strand1 && !strand2, // Both reverse
        };

        if !orientation_ok {
            return None;
        }

        // Calculate insert size (outer distance between mates)
        // Insert size = distance from leftmost base to rightmost base
        let (left_pos, _left_len, right_pos) = if pos1 <= pos2 {
            (pos1, len1, pos2 + len2 as u64)
        } else {
            (pos2, len2, pos1 + len1 as u64)
        };

        let insert_size = (right_pos - left_pos) as i32;

        // Check insert size constraints
        if insert_size < self.min_insert as i32 || insert_size > self.max_insert as i32 {
            return None;
        }

        // Check for containment (one read entirely within the other)
        if !self.allow_contain {
            let r1_end = pos1 + len1 as u64;
            let r2_end = pos2 + len2 as u64;

            // R1 contains R2
            if pos1 <= pos2 && r1_end >= r2_end {
                return None;
            }
            // R2 contains R1
            if pos2 <= pos1 && r2_end >= r1_end {
                return None;
            }
        }

        // Check for dovetailing (one read extends past the end of the other)
        if !self.allow_dovetail {
            // For FR policy with normal orientation, R1 should be left of R2
            // Dovetailing means R2 starts before R1 ends significantly
            if self.policy == PePolicy::FR && strand1 && !strand2 {
                let r1_end = pos1 + len1 as u64;
                if pos2 < r1_end {
                    // Mates overlap - check if it's just overlap or dovetail
                    let overlap = r1_end - pos2;
                    if overlap > len2 as u64 / 2 {
                        // Significant dovetail
                        return None;
                    }
                }
            }
        }

        Some(insert_size)
    }

    /// Calculate the expected position range for the opposite mate
    ///
    /// Given an anchor alignment, calculate where the opposite mate
    /// should align based on insert size constraints.
    ///
    /// # Arguments
    /// * `anchor_pos` - Position of the anchor mate (0-based)
    /// * `anchor_len` - Length of the anchor mate
    /// * `anchor_is_r1` - True if anchor is read 1, false if read 2
    /// * `anchor_strand` - True if anchor is forward strand
    ///
    /// # Returns
    /// (start_pos, end_pos, expected_strand) for the opposite mate search
    pub fn mate_search_range(
        &self,
        anchor_pos: u64,
        anchor_len: u32,
        anchor_is_r1: bool,
        anchor_strand: bool,
    ) -> (u64, u64, bool) {
        // Determine expected strand for opposite mate
        let opp_strand = match self.policy {
            PePolicy::FR => !anchor_strand, // Opposite strand
            PePolicy::RF => !anchor_strand, // Opposite strand
            PePolicy::FF => anchor_strand,  // Same strand
            PePolicy::RR => anchor_strand,  // Same strand
        };

        // Calculate search range based on policy and anchor position
        let anchor_end = anchor_pos + anchor_len as u64;

        let (start, end) = match self.policy {
            PePolicy::FR => {
                // FR: R1 forward left of R2 reverse
                if anchor_is_r1 {
                    // Anchor is R1 (forward), mate should be to the right
                    let start = anchor_pos.saturating_sub(self.max_insert as u64);
                    let end = anchor_end + self.max_insert as u64;
                    (start, end)
                } else {
                    // Anchor is R2 (reverse), mate should be to the left
                    let start = anchor_end.saturating_sub(self.max_insert as u64);
                    let end = anchor_pos + self.max_insert as u64;
                    (start, end)
                }
            }
            PePolicy::RF => {
                // RF: R1 reverse left of R2 forward
                if anchor_is_r1 {
                    // Anchor is R1 (reverse), mate should be to the right
                    let start = anchor_pos.saturating_sub(self.max_insert as u64);
                    let end = anchor_end + self.max_insert as u64;
                    (start, end)
                } else {
                    // Anchor is R2 (forward), mate should be to the left
                    let start = anchor_end.saturating_sub(self.max_insert as u64);
                    let end = anchor_pos + self.max_insert as u64;
                    (start, end)
                }
            }
            PePolicy::FF | PePolicy::RR => {
                // FF/RR: Mates can be on either side
                let start = anchor_end.saturating_sub(self.max_insert as u64);
                let end = anchor_pos + self.max_insert as u64;
                (start, end)
            }
        };

        (start, end, opp_strand)
    }
}

/// Classification of paired-end alignment types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PeAlignmentType {
    /// Normal concordant alignment
    #[default]
    Concordant,
    /// Mates overlap but neither contains the other
    Overlap,
    /// One mate contains the other
    Contain,
    /// Mates dovetail (one extends past the other)
    Dovetail,
    /// Discordant alignment (doesn't meet constraints)
    Discordant,
}

/// Result of a paired-end alignment
#[derive(Debug, Clone, Default)]
pub struct PairedAlignment {
    /// Chromosome index for mate 1
    pub chr1: u32,
    /// Position for mate 1 (0-based)
    pub pos1: u64,
    /// Strand for mate 1 (true = forward)
    pub strand1: bool,
    /// CIGAR for mate 1
    pub cigar1: String,
    /// Score for mate 1
    pub score1: Score,
    /// Number of edits for mate 1
    pub edits1: u32,

    /// Chromosome index for mate 2
    pub chr2: u32,
    /// Position for mate 2 (0-based)
    pub pos2: u64,
    /// Strand for mate 2 (true = forward)
    pub strand2: bool,
    /// CIGAR for mate 2
    pub cigar2: String,
    /// Score for mate 2
    pub score2: Score,
    /// Number of edits for mate 2
    pub edits2: u32,

    /// Insert size (TLEN in SAM)
    pub insert_size: i32,
    /// Alignment type
    pub alignment_type: PeAlignmentType,
    /// Mapping quality
    pub mapq: u8,
}

impl PairedAlignment {
    /// Create a new paired alignment
    pub fn new() -> Self {
        Self::default()
    }

    /// Check if this is a concordant alignment
    pub fn is_concordant(&self) -> bool {
        self.alignment_type == PeAlignmentType::Concordant
    }

    /// Calculate SAM flags for mate 1
    pub fn sam_flag1(&self) -> u16 {
        let mut flag: u16 = 0;
        flag |= 1; // PAIRED
        flag |= 2; // PROPER_PAIR (if concordant)
        if !self.strand1 {
            flag |= 16; // REVERSE
        }
        if !self.strand2 {
            flag |= 32; // MATE_REVERSE
        }
        flag |= 64; // FIRST_IN_PAIR
        flag
    }

    /// Calculate SAM flags for mate 2
    pub fn sam_flag2(&self) -> u16 {
        let mut flag: u16 = 0;
        flag |= 1; // PAIRED
        flag |= 2; // PROPER_PAIR (if concordant)
        if !self.strand2 {
            flag |= 16; // REVERSE
        }
        if !self.strand1 {
            flag |= 32; // MATE_REVERSE
        }
        flag |= 128; // SECOND_IN_PAIR
        flag
    }
}

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

    #[test]
    fn test_pe_policy_default() {
        let policy = PePolicy::default();
        assert_eq!(policy, PePolicy::FR);
    }

    #[test]
    fn test_pe_constraints_default() {
        let constraints = PeConstraints::default();
        assert_eq!(constraints.min_insert, 0);
        assert_eq!(constraints.max_insert, 500);
        assert_eq!(constraints.policy, PePolicy::FR);
    }

    #[test]
    fn test_concordant_fr_pair() {
        let constraints = PeConstraints::new(100, 500, PePolicy::FR);

        // FR pair: R1 forward at pos 1000, R2 reverse at pos 1200
        // Insert size = 1200 + 100 - 1000 = 300
        let result = constraints.is_concordant(
            1000, 100, true, 0, // R1: pos=1000, len=100, forward, chr0
            1200, 100, false, 0, // R2: pos=1200, len=100, reverse, chr0
        );

        assert_eq!(result, Some(300));
    }

    #[test]
    fn test_discordant_wrong_orientation() {
        let constraints = PeConstraints::new(100, 500, PePolicy::FR);

        // Wrong orientation: both forward
        let result = constraints.is_concordant(
            1000, 100, true, 0, 1200, 100, true, 0, // Wrong: R2 should be reverse
        );

        assert_eq!(result, None);
    }

    #[test]
    fn test_discordant_wrong_insert_size() {
        let constraints = PeConstraints::new(100, 500, PePolicy::FR);

        // Insert size too large
        let result = constraints.is_concordant(
            1000, 100, true, 0, 2000, 100, false, 0, // Insert = 2000 + 100 - 1000 = 1100 > 500
        );

        assert_eq!(result, None);
    }

    #[test]
    fn test_mate_search_range_fr_r1() {
        let constraints = PeConstraints::new(100, 500, PePolicy::FR);

        // R1 forward at position 1000, length 100
        let (start, end, strand) = constraints.mate_search_range(1000, 100, true, true);

        assert_eq!(start, 500); // 1000 - 500
        assert_eq!(end, 1600); // 1100 + 500
        assert!(!strand); // R2 should be reverse
    }

    #[test]
    fn test_paired_alignment_flags() {
        let mut aln = PairedAlignment::new();
        aln.strand1 = true;
        aln.strand2 = false;
        aln.alignment_type = PeAlignmentType::Concordant;

        let flag1 = aln.sam_flag1();
        let flag2 = aln.sam_flag2();

        assert_eq!(flag1 & 1, 1); // PAIRED
        assert_eq!(flag1 & 2, 2); // PROPER_PAIR
        assert_eq!(flag1 & 64, 64); // FIRST_IN_PAIR
        assert_eq!(flag1 & 16, 0); // R1 forward
        assert_eq!(flag1 & 32, 32); // Mate reverse

        assert_eq!(flag2 & 1, 1); // PAIRED
        assert_eq!(flag2 & 2, 2); // PROPER_PAIR
        assert_eq!(flag2 & 128, 128); // SECOND_IN_PAIR
        assert_eq!(flag2 & 16, 16); // R2 reverse
        assert_eq!(flag2 & 32, 0); // Mate forward
    }
}