rustalign_aligner/

seed.rs

//! Seed-based alignment implementation

use crate::{Constraint, SeedPolicy, SeedType};
use rustalign_common::Nuc;
use rustalign_fmindex::FmIndex;

/// A seed (substring match) for alignment
#[derive(Debug, Clone)]
pub struct Seed {
    /// BWT row position (for position resolution)
    pub bwt_row: u64,

    /// Starting position in reference
    pub pos: usize,

    /// Starting position in read
    pub read_pos: usize,

    /// Length of the seed
    pub len: usize,

    /// Seed type
    pub seed_type: SeedType,

    /// Constraint used for this seed
    pub constraint: Constraint,

    /// Number of genome positions this seed sequence matches
    /// Lower values indicate more unique/better seeds
    pub hit_count: usize,
}

impl Seed {
    /// Create a new seed
    pub fn new(bwt_row: u64, pos: usize, read_pos: usize, len: usize, seed_type: SeedType) -> Self {
        Self {
            bwt_row,
            pos,
            read_pos,
            len,
            seed_type,
            constraint: Constraint::new(),
            hit_count: 1,
        }
    }

    /// Create a new seed with a specific hit count
    pub fn with_hit_count(
        bwt_row: u64,
        pos: usize,
        read_pos: usize,
        len: usize,
        seed_type: SeedType,
        hit_count: usize,
    ) -> Self {
        Self {
            bwt_row,
            pos,
            read_pos,
            len,
            seed_type,
            constraint: Constraint::new(),
            hit_count,
        }
    }
}

/// Results from seed search
#[derive(Debug, Clone, Default)]
pub struct SeedResults {
    /// Seed hits found
    pub hits: Vec<Seed>,

    /// Number of seeds searched
    pub seeds_searched: usize,
}

/// Seed aligner that searches for seed matches
pub struct SeedAligner {
    /// Seed policy
    policy: SeedPolicy,

    /// Number of seed sets to try (similar to -R option)
    seed_sets: usize,
}

impl SeedAligner {
    /// Create a new seed aligner
    pub fn new(policy: SeedPolicy) -> Self {
        Self {
            policy,
            seed_sets: 3, // Try 3 different seed offsets
        }
    }

    /// Create a seed aligner with custom number of seed sets
    pub fn with_seed_sets(policy: SeedPolicy, seed_sets: usize) -> Self {
        Self { policy, seed_sets }
    }

    /// Search for seeds in a read using multiple seed sets
    ///
    /// This implements a multi-set seed search.
    /// If the first seed set finds no matches, it tries additional sets with
    /// different offsets to handle reads with SNPs that break some seeds.
    pub fn search<I: FmIndex>(
        &mut self,
        index: &I,
        read: &[Nuc],
    ) -> rustalign_common::Result<SeedResults> {
        let mut best_results = SeedResults::default();
        let _read_len = read.len();

        // Try multiple seed sets with different offsets
        for set_idx in 0..self.seed_sets {
            // Calculate offset for this seed set
            // Set 0: offset 0 (default)
            // Set 1: offset seed_interval/2
            // Set 2: offset seed_interval/4
            let offset = if set_idx == 0 {
                0
            } else {
                (self.policy.seed_interval * set_idx) / (self.seed_sets)
            };

            let results = self.search_with_offset(index, read, offset)?;

            // If this set found matches, use them
            if !results.hits.is_empty() {
                // Merge with best results (prefer earlier sets for same uniqueness)
                for hit in results.hits {
                    // Avoid duplicates from different seed sets
                    let is_duplicate = best_results.hits.iter().any(|existing| {
                        existing.bwt_row == hit.bwt_row && existing.read_pos == hit.read_pos
                    });
                    if !is_duplicate {
                        best_results.hits.push(hit);
                    }
                }
                best_results.seeds_searched += results.seeds_searched;

                // If we have good results from the first set, stop early
                if set_idx == 0 {
                    break;
                }
            }
        }

        // Sort hits by uniqueness (fewer hits = more unique) as primary criterion,
        // then by read position to prefer seeds earlier in the read
        best_results.hits.sort_by(|a, b| {
            // Primary: prefer seeds with fewer genome hits (more unique)
            match a.hit_count.cmp(&b.hit_count) {
                std::cmp::Ordering::Equal => {
                    // Secondary: prefer seeds earlier in the read
                    a.read_pos.cmp(&b.read_pos)
                }
                other => other,
            }
        });

        Ok(best_results)
    }

    /// Search for seeds with a specific starting offset
    fn search_with_offset<I: FmIndex>(
        &mut self,
        index: &I,
        read: &[Nuc],
        offset: usize,
    ) -> rustalign_common::Result<SeedResults> {
        let mut results = SeedResults::default();
        let read_len = read.len();

        // Generate seed positions starting from offset
        let mut seed_pos = offset;
        while seed_pos + self.policy.seed_len <= read_len {
            let seed_seq = &read[seed_pos..seed_pos + self.policy.seed_len];

            // Search for exact matches
            match index.exact_match(seed_seq) {
                Ok(positions) => {
                    // Track how many genome positions this seed matches
                    let hit_count = positions.len();

                    // Limit to max_hits to avoid explosion with repetitive seeds
                    for bwt_row in positions.into_iter().take(self.policy.max_hits) {
                        let seed = Seed::with_hit_count(
                            bwt_row as u64,
                            bwt_row,
                            seed_pos,
                            self.policy.seed_len,
                            self.policy.seed_type,
                            hit_count,
                        );
                        results.hits.push(seed);
                    }
                }
                Err(_) => {
                    // Continue trying other seeds even if this one fails
                }
            }

            results.seeds_searched += 1;
            seed_pos += self.policy.seed_interval;
        }

        Ok(results)
    }
}

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

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

    #[test]
    fn test_seed_new() {
        let seed = Seed::new(100, 100, 0, 22, SeedType::Exact);
        assert_eq!(seed.bwt_row, 100);
        assert_eq!(seed.pos, 100);
        assert_eq!(seed.read_pos, 0);
        assert_eq!(seed.len, 22);
        assert_eq!(seed.hit_count, 1);
    }

    #[test]
    fn test_seed_with_hit_count() {
        let seed = Seed::with_hit_count(100, 100, 0, 22, SeedType::Exact, 5);
        assert_eq!(seed.bwt_row, 100);
        assert_eq!(seed.hit_count, 5);
    }

    #[test]
    fn test_seed_aligner_new() {
        let aligner = SeedAligner::new(SeedPolicy::new());
        assert_eq!(aligner.policy.seed_len, 22);
    }
}