genomicframe_core/interval/

annotation.rs

//! Genomic interval annotation
//!
//! This module provides infrastructure for annotating genomic intervals
//! with contextual information from reference databases (genes, exons, etc.).
//!
//! # Design Philosophy
//!
//! - **Lazy by default**: Annotations are computed on-demand during iteration
//! - **Format-agnostic**: Works with any format that converts to GenomicInterval
//! - **Composable**: Multiple annotation sources can be chained
//! - **Fast queries**: O(log n + k) overlap queries using interval trees
//!
//! # Examples
//!
//! ```no_run
//! use genomicframe_core::interval::annotation::AnnotationIndex;
//! use genomicframe_core::formats::bed::BedReader;
//!
//! // Build annotation index from BED file
//! let exons = AnnotationIndex::from_bed("exons.bed", |record| {
//!     record.name.clone().unwrap_or_else(|| "unknown".to_string())
//! })?;
//!
//! // Query for overlapping annotations
//! let interval = GenomicInterval::new("chr1", 1000, 2000);
//! let genes = exons.query(&interval);
//! # use genomicframe_core::interval::GenomicInterval;
//! # Ok::<(), genomicframe_core::error::Error>(())
//! ```

use crate::core::GenomicRecordIterator;
use crate::error::Result;
use crate::formats::bed::{BedReader, BedRecord};
use crate::interval::GenomicInterval;
use std::collections::HashMap;
use std::path::Path;

/// Annotation metadata attached to an interval
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Annotation {
    /// The genomic interval
    pub interval: GenomicInterval,
    /// Free-form annotation data (gene name, feature type, etc.)
    pub data: String,
}

/// Interval tree node for O(log n + k) overlap queries
#[derive(Debug, Clone)]
pub struct IntervalTreeNode {
    center: u64,
    left_sorted: Vec<Annotation>,  // Sorted by start position
    right_sorted: Vec<Annotation>, // Sorted by end position
    left: Option<Box<IntervalTreeNode>>,
    right: Option<Box<IntervalTreeNode>>,
}

/// A standalone interval tree for efficient overlap queries
///
/// This provides direct access to the tree structure for advanced use cases.
#[derive(Debug, Clone)]
pub struct IntervalTree {
    root: Option<IntervalTreeNode>,
    count: usize,
}

impl IntervalTree {
    /// Create an empty interval tree
    pub fn new() -> Self {
        Self {
            root: None,
            count: 0,
        }
    }

    /// Build interval tree from a collection of annotations
    pub fn from_annotations(annotations: Vec<Annotation>) -> Self {
        let count = annotations.len();
        let root = IntervalTreeNode::from_annotations(annotations);
        Self { root, count }
    }

    /// Query for overlapping annotations
    pub fn query<'a>(&'a self, interval: &GenomicInterval) -> Vec<&'a str> {
        let mut results = Vec::new();
        if let Some(ref root) = self.root {
            root.query(interval, &mut results);
        }
        results
    }

    /// Query for overlapping annotations (full structs)
    pub fn query_annotations<'a>(&'a self, interval: &GenomicInterval) -> Vec<&'a Annotation> {
        let mut results = Vec::new();
        if let Some(ref root) = self.root {
            root.query_annotations(interval, &mut results);
        }
        results
    }

    /// Count of annotations in the tree
    pub fn len(&self) -> usize {
        self.count
    }

    /// Check if tree is empty
    pub fn is_empty(&self) -> bool {
        self.count == 0
    }
}

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

impl IntervalTreeNode {
    /// Build interval tree from a list of annotations
    fn from_annotations(annotations: Vec<Annotation>) -> Option<Self> {
        if annotations.is_empty() {
            return None;
        }

        // Base case: for small sets, stop recursing and store directly
        // This prevents stack overflow and improves performance for leaf nodes
        // With a threshold of 64, max tree depth is ~log2(N/64), so for 32K items
        // that's ~9 levels, well within stack limits
        if annotations.len() <= 64 {
            let mut left_sorted = annotations.clone();
            let mut right_sorted = annotations;
            left_sorted.sort_by_key(|a| a.interval.start);
            right_sorted.sort_by_key(|a| a.interval.end);

            let center = left_sorted[left_sorted.len() / 2].interval.start;

            return Some(IntervalTreeNode {
                center,
                left_sorted,
                right_sorted,
                left: None,
                right: None,
            });
        }

        // Find center point (median of all start/end coordinates)
        let mut coords: Vec<u64> = annotations
            .iter()
            .flat_map(|a| vec![a.interval.start, a.interval.end])
            .collect();
        coords.sort_unstable();
        let center = coords[coords.len() / 2];

        // Partition intervals
        let mut left_intervals = Vec::new();
        let mut right_intervals = Vec::new();
        let mut center_intervals = Vec::new();

        for ann in annotations {
            if ann.interval.end <= center {
                left_intervals.push(ann);
            } else if ann.interval.start > center {
                right_intervals.push(ann);
            } else {
                // Overlaps center
                center_intervals.push(ann);
            }
        }

        // Sort center intervals for efficient stabbing queries
        let mut left_sorted = center_intervals.clone();
        let mut right_sorted = center_intervals;

        left_sorted.sort_by_key(|a| a.interval.start);
        right_sorted.sort_by_key(|a| a.interval.end);

        Some(IntervalTreeNode {
            center,
            left_sorted,
            right_sorted,
            left: IntervalTreeNode::from_annotations(left_intervals).map(Box::new),
            right: IntervalTreeNode::from_annotations(right_intervals).map(Box::new),
        })
    }

    /// Query for overlapping intervals
    fn query<'a>(&'a self, interval: &GenomicInterval, results: &mut Vec<&'a str>) {
        // Check intervals that span the center
        if interval.end <= self.center {
            // Query left of center
            for ann in &self.left_sorted {
                if ann.interval.start >= interval.end {
                    break; // No more overlaps possible
                }
                if ann.interval.overlaps(interval) {
                    results.push(&ann.data);
                }
            }
        } else if interval.start > self.center {
            // Query right of center
            for ann in self.right_sorted.iter().rev() {
                if ann.interval.end <= interval.start {
                    break; // No more overlaps possible
                }
                if ann.interval.overlaps(interval) {
                    results.push(&ann.data);
                }
            }
        } else {
            // Query spans center - check all center intervals
            for ann in &self.left_sorted {
                if ann.interval.overlaps(interval) {
                    results.push(&ann.data);
                }
            }
        }

        // Recurse into subtrees
        if interval.start < self.center {
            if let Some(ref left) = self.left {
                left.query(interval, results);
            }
        }
        if interval.end > self.center {
            if let Some(ref right) = self.right {
                right.query(interval, results);
            }
        }
    }

    /// Query for overlapping annotations (full structs)
    fn query_annotations<'a>(
        &'a self,
        interval: &GenomicInterval,
        results: &mut Vec<&'a Annotation>,
    ) {
        // Check intervals that span the center
        if interval.end <= self.center {
            for ann in &self.left_sorted {
                if ann.interval.start >= interval.end {
                    break;
                }
                if ann.interval.overlaps(interval) {
                    results.push(ann);
                }
            }
        } else if interval.start > self.center {
            for ann in self.right_sorted.iter().rev() {
                if ann.interval.end <= interval.start {
                    break;
                }
                if ann.interval.overlaps(interval) {
                    results.push(ann);
                }
            }
        } else {
            for ann in &self.left_sorted {
                if ann.interval.overlaps(interval) {
                    results.push(ann);
                }
            }
        }

        if interval.start < self.center {
            if let Some(ref left) = self.left {
                left.query_annotations(interval, results);
            }
        }
        if interval.end > self.center {
            if let Some(ref right) = self.right {
                right.query_annotations(interval, results);
            }
        }
    }
}

/// Fast interval index using interval trees for O(log n + k) overlap queries
///
/// This implementation uses a centered interval tree per chromosome,
/// providing logarithmic query time instead of linear scanning.
#[derive(Debug, Clone)]
pub struct AnnotationIndex {
    /// Interval trees grouped by chromosome
    trees: HashMap<String, IntervalTreeNode>,
    /// Total annotation count (cached for efficiency)
    count: usize,
}

impl AnnotationIndex {
    /// Create an empty annotation index
    pub fn new() -> Self {
        Self {
            trees: HashMap::new(),
            count: 0,
        }
    }

    /// Build interval trees from collected annotations
    ///
    /// This should be called after all insertions are complete.
    /// For efficiency, batch all insertions then call this once.
    pub fn build_trees(&mut self, annotations: HashMap<String, Vec<Annotation>>) {
        self.count = annotations.values().map(|v| v.len()).sum();

        for (chrom, ann_list) in annotations {
            if let Some(tree) = IntervalTreeNode::from_annotations(ann_list) {
                self.trees.insert(chrom, tree);
            }
        }
    }

    /// Build annotation index from BED file
    ///
    /// The `extractor` function converts each BED record into annotation data.
    /// Typically this extracts the "name" field, but can be customized.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::interval::annotation::AnnotationIndex;
    ///
    /// // Extract gene names from BED name field
    /// let index = AnnotationIndex::from_bed("genes.bed", |record| {
    ///     record.name.clone().unwrap_or_default()
    /// })?;
    ///
    /// // Custom extraction: combine name and score
    /// let index = AnnotationIndex::from_bed("features.bed", |record| {
    ///     format!("{}:{}",
    ///         record.name.as_deref().unwrap_or("unknown"),
    ///         record.score.unwrap_or(0))
    /// })?;
    /// # Ok::<(), genomicframe_core::error::Error>(())
    /// ```
    pub fn from_bed<P, F>(path: P, extractor: F) -> Result<Self>
    where
        P: AsRef<Path>,
        F: Fn(&BedRecord) -> String,
    {
        let mut reader = BedReader::from_path(path)?;
        let mut annotations: HashMap<String, Vec<Annotation>> = HashMap::new();

        while let Some(record) = reader.next_record()? {
            let interval: GenomicInterval = (&record).into();
            let data = extractor(&record);

            annotations
                .entry(interval.chrom.clone())
                .or_insert_with(Vec::new)
                .push(Annotation { interval, data });
        }

        let mut index = Self::new();
        index.build_trees(annotations);
        Ok(index)
    }

    /// Build annotation index from an iterator of BED records
    pub fn from_records<'a, I, F>(records: I, extractor: F) -> Self
    where
        I: IntoIterator<Item = &'a BedRecord>,
        F: Fn(&BedRecord) -> String,
    {
        let mut annotations: HashMap<String, Vec<Annotation>> = HashMap::new();

        for record in records {
            let interval: GenomicInterval = record.into();
            let data = extractor(record);

            annotations
                .entry(interval.chrom.clone())
                .or_insert_with(Vec::new)
                .push(Annotation { interval, data });
        }

        let mut index = Self::new();
        index.build_trees(annotations);
        index
    }

    /// Build annotation index from a reader (possibly filtered)
    pub fn from_reader<I, F>(mut reader: I, extractor: F) -> Result<Self>
    where
        I: crate::core::GenomicRecordIterator<Record = BedRecord>,
        F: Fn(&BedRecord) -> String,
    {
        let mut annotations: HashMap<String, Vec<Annotation>> = HashMap::new();

        while let Some(record) = reader.next_record()? {
            let interval: GenomicInterval = (&record).into();
            let data = extractor(&record);

            annotations
                .entry(interval.chrom.clone())
                .or_insert_with(Vec::new)
                .push(Annotation { interval, data });
        }

        let mut index = Self::new();
        index.build_trees(annotations);
        Ok(index)
    }

    /// Query for all annotations overlapping the given interval
    ///
    /// Returns references to annotation data strings.
    ///
    /// # Performance
    /// O(log n + k) where n = number of annotations on the chromosome,
    /// k = number of overlapping annotations
    pub fn query(&self, interval: &GenomicInterval) -> Vec<&str> {
        let mut results = Vec::new();

        // Try exact chrom match first
        if let Some(tree) = self.trees.get(&interval.chrom) {
            tree.query(interval, &mut results);
            return results;
        }

        // Fallback: try alternate chrom name (add/remove leading "chr")
        let alt_chrom = if interval.chrom.starts_with("chr") {
            interval.chrom.trim_start_matches("chr").to_string()
        } else {
            format!("chr{}", interval.chrom)
        };

        if let Some(tree) = self.trees.get(&alt_chrom) {
            // Create a normalized query interval for overlap checking
            let normalized_query = GenomicInterval {
                chrom: alt_chrom,
                start: interval.start,
                end: interval.end,
            };

            tree.query(&normalized_query, &mut results);
        }

        results
    }

    /// Query for all annotations overlapping the given interval
    ///
    /// Returns full Annotation structs (includes intervals).
    pub fn query_annotations(&self, interval: &GenomicInterval) -> Vec<&Annotation> {
        let mut results = Vec::new();

        // Try exact chrom match first
        if let Some(tree) = self.trees.get(&interval.chrom) {
            tree.query_annotations(interval, &mut results);
            return results;
        }

        // Fallback: try alternate chrom name
        let alt_chrom = if interval.chrom.starts_with("chr") {
            interval.chrom.trim_start_matches("chr").to_string()
        } else {
            format!("chr{}", interval.chrom)
        };

        if let Some(tree) = self.trees.get(&alt_chrom) {
            let normalized_query = GenomicInterval {
                chrom: alt_chrom,
                start: interval.start,
                end: interval.end,
            };

            tree.query_annotations(&normalized_query, &mut results);
        }

        results
    }

    /// Count total annotations in the index
    pub fn len(&self) -> usize {
        self.count
    }

    /// Check if the index is empty
    pub fn is_empty(&self) -> bool {
        self.count == 0
    }

    /// Get chromosomes present in the index
    pub fn chromosomes(&self) -> Vec<&str> {
        self.trees.keys().map(|s| s.as_str()).collect()
    }
}

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

/// Multi-index query result for efficient batch queries
#[derive(Debug, Clone)]
pub struct MultiQueryResult<'a> {
    pub genes: Vec<&'a str>,
    pub exons: Vec<&'a str>,
}

/// Query multiple annotation indexes simultaneously for better cache locality
///
/// This is more efficient than calling query() on each index separately
/// when you need results from multiple sources.
pub fn multi_query<'a>(
    interval: &GenomicInterval,
    genes: &'a AnnotationIndex,
    exons: &'a AnnotationIndex,
) -> MultiQueryResult<'a> {
    // Query both indexes - compiler can potentially optimize this
    let gene_results = genes.query(interval);
    let exon_results = exons.query(interval);

    MultiQueryResult {
        genes: gene_results,
        exons: exon_results,
    }
}

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

    #[test]
    fn test_annotation_index() {
        let mut annotations: HashMap<String, Vec<Annotation>> = HashMap::new();

        annotations
            .entry("chr1".to_string())
            .or_insert_with(Vec::new)
            .push(Annotation {
                interval: GenomicInterval::new("chr1", 100, 200),
                data: "GeneA".to_string(),
            });

        annotations
            .entry("chr1".to_string())
            .or_insert_with(Vec::new)
            .push(Annotation {
                interval: GenomicInterval::new("chr1", 150, 250),
                data: "GeneB".to_string(),
            });

        annotations
            .entry("chr2".to_string())
            .or_insert_with(Vec::new)
            .push(Annotation {
                interval: GenomicInterval::new("chr2", 100, 200),
                data: "GeneC".to_string(),
            });

        let mut index = AnnotationIndex::new();
        index.build_trees(annotations);

        // Query overlapping interval
        let query = GenomicInterval::new("chr1", 175, 225);
        let results = index.query(&query);

        assert_eq!(results.len(), 2);
        assert!(results.contains(&"GeneA"));
        assert!(results.contains(&"GeneB"));

        // Query non-overlapping interval
        let query = GenomicInterval::new("chr1", 300, 400);
        let results = index.query(&query);
        assert_eq!(results.len(), 0);

        // Query different chromosome
        let query = GenomicInterval::new("chr2", 150, 250);
        let results = index.query(&query);
        assert_eq!(results.len(), 1);
        assert!(results.contains(&"GeneC"));
    }

    #[test]
    fn test_annotation_index_stats() {
        let mut annotations: HashMap<String, Vec<Annotation>> = HashMap::new();

        annotations
            .entry("chr1".to_string())
            .or_insert_with(Vec::new)
            .push(Annotation {
                interval: GenomicInterval::new("chr1", 100, 200),
                data: "GeneA".to_string(),
            });

        annotations
            .entry("chr2".to_string())
            .or_insert_with(Vec::new)
            .push(Annotation {
                interval: GenomicInterval::new("chr2", 100, 200),
                data: "GeneB".to_string(),
            });

        let mut index = AnnotationIndex::new();
        index.build_trees(annotations);

        assert!(!index.is_empty());
        assert_eq!(index.len(), 2);

        let chroms = index.chromosomes();
        assert_eq!(chroms.len(), 2);
        assert!(chroms.contains(&"chr1"));
        assert!(chroms.contains(&"chr2"));
    }

    #[test]
    fn test_multi_query() {
        let mut gene_annotations: HashMap<String, Vec<Annotation>> = HashMap::new();
        let mut exon_annotations: HashMap<String, Vec<Annotation>> = HashMap::new();

        gene_annotations
            .entry("chr1".to_string())
            .or_insert_with(Vec::new)
            .push(Annotation {
                interval: GenomicInterval::new("chr1", 100, 500),
                data: "GeneA".to_string(),
            });

        exon_annotations
            .entry("chr1".to_string())
            .or_insert_with(Vec::new)
            .push(Annotation {
                interval: GenomicInterval::new("chr1", 150, 200),
                data: "ExonA1".to_string(),
            });

        let mut genes = AnnotationIndex::new();
        genes.build_trees(gene_annotations);

        let mut exons = AnnotationIndex::new();
        exons.build_trees(exon_annotations);

        // Query both simultaneously
        let query = GenomicInterval::new("chr1", 175, 225);
        let results = multi_query(&query, &genes, &exons);

        assert_eq!(results.genes.len(), 1);
        assert_eq!(results.exons.len(), 1);
        assert_eq!(results.genes[0], "GeneA");
        assert_eq!(results.exons[0], "ExonA1");
    }
}