Skip to main content

gtars_genomicdist/
models.rs

1use crate::errors::GtarsGenomicDistError;
2use bio::io::fasta;
3use gtars_core::models::{CoordinateMode, Region, RegionSet};
4use memmap2::Mmap;
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7use std::fmt::Debug;
8use std::fs::File;
9use std::io::{BufWriter, Write};
10use std::path::Path;
11
12// `SortedRegionSet` is defined in gtars-core; re-exported here so the
13// genomicdist public path (`gtars_genomicdist::SortedRegionSet`) is unchanged.
14pub use gtars_core::models::SortedRegionSet;
15
16/// Genomic strand orientation.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
18pub enum Strand {
19    Plus,
20    Minus,
21    Unstranded,
22}
23
24impl Strand {
25    pub fn from_char(c: char) -> Self {
26        match c {
27            '+' => Strand::Plus,
28            '-' => Strand::Minus,
29            _ => Strand::Unstranded,
30        }
31    }
32}
33
34/// A `RegionSet` paired with a parallel `Vec<Strand>`.
35///
36/// Follows the `SortedRegionSet` wrapper pattern: wraps a `RegionSet` and
37/// adds strand information without modifying `Region` itself (which lives
38/// in gtars-core). Strand-aware operations (promoters, reduce, setdiff)
39/// are implemented as methods on this type.
40#[derive(Clone, Serialize, Deserialize)]
41pub struct StrandedRegionSet {
42    pub inner: RegionSet,
43    pub strands: Vec<Strand>,
44}
45
46impl StrandedRegionSet {
47    /// Create a new StrandedRegionSet from a RegionSet and parallel strand vector.
48    ///
49    /// # Panics
50    /// Panics if `strands.len() != rs.regions.len()`.
51    pub fn new(rs: RegionSet, strands: Vec<Strand>) -> Self {
52        assert_eq!(
53            rs.regions.len(),
54            strands.len(),
55            "StrandedRegionSet: regions and strands must have the same length"
56        );
57        StrandedRegionSet {
58            inner: rs,
59            strands,
60        }
61    }
62
63    /// Wrap a RegionSet with all-Unstranded. Preserves existing behavior.
64    pub fn unstranded(rs: RegionSet) -> Self {
65        let n = rs.regions.len();
66        StrandedRegionSet {
67            strands: vec![Strand::Unstranded; n],
68            inner: rs,
69        }
70    }
71
72    /// Consume into the inner RegionSet, dropping strand information.
73    pub fn into_regionset(self) -> RegionSet {
74        self.inner
75    }
76
77    pub fn len(&self) -> usize {
78        self.inner.regions.len()
79    }
80
81    pub fn is_empty(&self) -> bool {
82        self.inner.regions.is_empty()
83    }
84}
85
86/// Statistics summary for regions on a single chromosome.
87///
88/// Contains counts, bounds, and descriptive statistics for region lengths.
89#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct ChromosomeStatistics {
91    /// Chromosome name
92    pub chromosome: String,
93    /// Total number of regions on this chromosome
94    pub number_of_regions: u32,
95    /// Leftmost start position across all regions
96    pub start_nucleotide_position: u32,
97    /// Rightmost end position across all regions
98    pub end_nucleotide_position: u32,
99    /// Length of the shortest region
100    pub minimum_region_length: u32,
101    /// Length of the longest region
102    pub maximum_region_length: u32,
103    /// Average region length
104    pub mean_region_length: f64,
105    /// Median region length
106    pub median_region_length: f64,
107}
108
109/// A genomic bin with a count of overlapping regions.
110///
111/// Used to represent distribution of regions across fixed-size windows.
112#[derive(Debug, Clone, Serialize, Deserialize)]
113pub struct RegionBin {
114    /// Chromosome name
115    pub chr: String,
116    /// Start position of the bin
117    pub start: u32,
118    /// End position of the bin
119    pub end: u32,
120    /// Number of regions overlapping this bin
121    pub n: u32,
122    /// Rid: needed for plot to have correct order
123    pub rid: u32,
124}
125
126/// Trait for types that provide sequence access to a reference genome.
127///
128/// Implemented by [`GenomeAssembly`] (in-memory HashMap) and
129/// [`BinaryGenomeAssembly`] (mmap .fab binary). Functions like
130/// `calc_gc_content` accept `&impl SequenceAccess` to work with either.
131pub trait SequenceAccess {
132    /// Get the sequence for a genomic region. Returns owned bytes.
133    fn get_sequence(&self, coords: &Region) -> Result<Vec<u8>, GtarsGenomicDistError>;
134
135    /// Check whether a chromosome exists in the assembly.
136    fn contains_chr(&self, chr: &str) -> bool;
137}
138
139/// In-memory genome assembly backed by a HashMap of chromosome sequences.
140///
141/// Loads the entire FASTA into memory on construction. Slower to construct
142/// (~2s for hg38) but provides zero-copy `&[u8]` access to sequences,
143/// making per-region operations like GC content highly vectorizable.
144/// No `.fai` index required.
145pub struct GenomeAssembly {
146    seq_map: HashMap<String, Vec<u8>>,
147}
148
149impl TryFrom<&str> for GenomeAssembly {
150    type Error = GtarsGenomicDistError;
151    fn try_from(value: &str) -> Result<Self, GtarsGenomicDistError> {
152        GenomeAssembly::try_from(Path::new(value))
153    }
154}
155
156impl TryFrom<String> for GenomeAssembly {
157    type Error = GtarsGenomicDistError;
158    fn try_from(value: String) -> Result<Self, GtarsGenomicDistError> {
159        GenomeAssembly::try_from(Path::new(&value))
160    }
161}
162
163impl TryFrom<&Path> for GenomeAssembly {
164    type Error = GtarsGenomicDistError;
165
166    fn try_from(value: &Path) -> Result<GenomeAssembly, GtarsGenomicDistError> {
167        let file = File::open(value)?;
168        let genome = fasta::Reader::new(file);
169        let records = genome.records();
170
171        let mut seq_map: HashMap<String, Vec<u8>> = HashMap::new();
172        for record in records {
173            match record {
174                Ok(record) => {
175                    seq_map.insert(record.id().to_string(), record.seq().to_owned());
176                }
177                Err(e) => {
178                    return Err(GtarsGenomicDistError::CustomError(format!(
179                        "Error reading genome file: {}",
180                        e
181                    )));
182                }
183            }
184        }
185        Ok(GenomeAssembly { seq_map })
186    }
187}
188
189impl GenomeAssembly {
190    pub fn seq_from_region(&self, coords: &Region) -> Result<&[u8], GtarsGenomicDistError> {
191        let chr = &coords.chr;
192        let start = coords.start as usize;
193        let end = coords.end as usize;
194
195        if let Some(seq) = self.seq_map.get(chr) {
196            if end <= seq.len() && start <= end {
197                Ok(&seq[start..end])
198            } else {
199                Err(GtarsGenomicDistError::CustomError(format!(
200                    "Invalid range: start={}, end={} for chromosome {} with length {}",
201                    start, end, chr, seq.len()
202                )))
203            }
204        } else {
205            Err(GtarsGenomicDistError::CustomError(format!(
206                "Unknown chromosome found in region set: {}",
207                chr
208            )))
209        }
210    }
211
212    pub fn contains_chr(&self, chr: &str) -> bool {
213        self.seq_map.contains_key(chr)
214    }
215}
216
217impl SequenceAccess for GenomeAssembly {
218    fn get_sequence(&self, coords: &Region) -> Result<Vec<u8>, GtarsGenomicDistError> {
219        self.seq_from_region(coords).map(|s| s.to_vec())
220    }
221
222    fn contains_chr(&self, chr: &str) -> bool {
223        self.seq_map.contains_key(chr)
224    }
225}
226
227// --- Binary FASTA (.fab) format ---
228
229const FAB_MAGIC: &[u8; 4] = b"GFAB";
230const FAB_VERSION: u8 = 1;
231
232/// Memory-mapped genome assembly backed by a binary FASTA (.fab) file.
233///
234/// The .fab format stores sequences contiguously without line wrapping,
235/// enabling mmap + zero-copy `&[u8]` access with both instant construction
236/// and zero-copy per-region performance.
237///
238/// Create .fab files with [`BinaryGenomeAssembly::write_from_fasta`] or
239/// via `gtars prep --fasta <file>`.
240#[derive(Debug)]
241pub struct BinaryGenomeAssembly {
242    mmap: Mmap,
243    /// name → (offset_in_file, sequence_length)
244    index: HashMap<String, (usize, usize)>,
245}
246
247impl BinaryGenomeAssembly {
248    /// Open a .fab binary FASTA file via mmap.
249    pub fn from_file(path: &Path) -> Result<Self, GtarsGenomicDistError> {
250        let file = File::open(path).map_err(|e| {
251            GtarsGenomicDistError::CustomError(format!(
252                "Failed to open .fab file '{}': {}",
253                path.display(), e
254            ))
255        })?;
256        let mmap = unsafe { Mmap::map(&file) }.map_err(|e| {
257            GtarsGenomicDistError::CustomError(format!(
258                "Failed to mmap .fab file '{}': {}",
259                path.display(), e
260            ))
261        })?;
262
263        // Parse header
264        if mmap.len() < 9 {
265            return Err(GtarsGenomicDistError::CustomError(
266                "Invalid .fab file: too short".into(),
267            ));
268        }
269        if &mmap[0..4] != FAB_MAGIC {
270            return Err(GtarsGenomicDistError::CustomError(
271                "Invalid .fab file: bad magic bytes".into(),
272            ));
273        }
274        let version = mmap[4];
275        if version != FAB_VERSION {
276            return Err(GtarsGenomicDistError::CustomError(format!(
277                "Unsupported .fab version: {} (expected {})",
278                version, FAB_VERSION
279            )));
280        }
281        let n_chroms = u32::from_le_bytes(mmap[5..9].try_into().unwrap()) as usize;
282
283        // Parse index
284        let mut pos = 9;
285        let mut index = HashMap::with_capacity(n_chroms);
286        for _ in 0..n_chroms {
287            if pos + 2 > mmap.len() {
288                return Err(GtarsGenomicDistError::CustomError(
289                    "Invalid .fab file: truncated index".into(),
290                ));
291            }
292            let name_len = u16::from_le_bytes(mmap[pos..pos + 2].try_into().unwrap()) as usize;
293            pos += 2;
294            if pos + name_len + 16 > mmap.len() {
295                return Err(GtarsGenomicDistError::CustomError(
296                    "Invalid .fab file: truncated index entry".into(),
297                ));
298            }
299            let name = std::str::from_utf8(&mmap[pos..pos + name_len])
300                .map_err(|e| {
301                    GtarsGenomicDistError::CustomError(format!(
302                        "Invalid .fab file: non-UTF8 chromosome name: {}",
303                        e
304                    ))
305                })?
306                .to_string();
307            pos += name_len;
308            let offset =
309                u64::from_le_bytes(mmap[pos..pos + 8].try_into().unwrap()) as usize;
310            pos += 8;
311            let length =
312                u64::from_le_bytes(mmap[pos..pos + 8].try_into().unwrap()) as usize;
313            pos += 8;
314            index.insert(name, (offset, length));
315        }
316
317        Ok(BinaryGenomeAssembly { mmap, index })
318    }
319
320    /// Get the sequence for a region as a zero-copy `&[u8]` slice.
321    pub fn seq_from_region(&self, coords: &Region) -> Result<&[u8], GtarsGenomicDistError> {
322        let chr = &coords.chr;
323        let start = coords.start as usize;
324        let end = coords.end as usize;
325
326        let &(offset, length) = self.index.get(chr).ok_or_else(|| {
327            GtarsGenomicDistError::CustomError(format!(
328                "Unknown chromosome found in region set: {}",
329                chr
330            ))
331        })?;
332
333        if end > length || start > end {
334            return Err(GtarsGenomicDistError::CustomError(format!(
335                "Invalid range: start={}, end={} for chromosome {} with length {}",
336                start, end, chr, length
337            )));
338        }
339
340        let file_start = offset + start;
341        let file_end = offset + end;
342        if file_end > self.mmap.len() {
343            return Err(GtarsGenomicDistError::CustomError(format!(
344                "Corrupted .fab file: sequence data for {} extends beyond file boundary",
345                chr
346            )));
347        }
348
349        Ok(&self.mmap[file_start..file_end])
350    }
351
352    pub fn contains_chr(&self, chr: &str) -> bool {
353        self.index.contains_key(chr)
354    }
355
356    /// Convert a FASTA file to .fab binary format.
357    pub fn write_from_fasta(
358        fasta_path: &Path,
359        output_path: &Path,
360    ) -> Result<(), GtarsGenomicDistError> {
361        // Read all sequences into memory (same as GenomeAssembly)
362        let file = File::open(fasta_path)?;
363        let reader = fasta::Reader::new(file);
364
365        let mut chroms: Vec<(String, Vec<u8>)> = Vec::new();
366        for record in reader.records() {
367            let record = record.map_err(|e| {
368                GtarsGenomicDistError::CustomError(format!(
369                    "Error reading FASTA: {}", e
370                ))
371            })?;
372            chroms.push((record.id().to_string(), record.seq().to_owned()));
373        }
374
375        // Compute index: header size first
376        let mut header_size: usize = 4 + 1 + 4; // magic + version + n_chroms
377        for (name, _) in &chroms {
378            header_size += 2 + name.len() + 8 + 8; // name_len + name + offset + length
379        }
380
381        // Write
382        let out = File::create(output_path).map_err(|e| {
383            GtarsGenomicDistError::CustomError(format!(
384                "Failed to create .fab file '{}': {}",
385                output_path.display(), e
386            ))
387        })?;
388        let mut w = BufWriter::new(out);
389
390        // Header
391        w.write_all(FAB_MAGIC)?;
392        w.write_all(&[FAB_VERSION])?;
393        w.write_all(&(chroms.len() as u32).to_le_bytes())?;
394
395        // Index
396        let mut offset = header_size;
397        for (name, seq) in &chroms {
398            w.write_all(&(name.len() as u16).to_le_bytes())?;
399            w.write_all(name.as_bytes())?;
400            w.write_all(&(offset as u64).to_le_bytes())?;
401            w.write_all(&(seq.len() as u64).to_le_bytes())?;
402            offset += seq.len();
403        }
404
405        // Sequences
406        for (_, seq) in &chroms {
407            w.write_all(seq)?;
408        }
409
410        w.flush()?;
411        Ok(())
412    }
413}
414
415impl TryFrom<&str> for BinaryGenomeAssembly {
416    type Error = GtarsGenomicDistError;
417    fn try_from(value: &str) -> Result<Self, GtarsGenomicDistError> {
418        BinaryGenomeAssembly::from_file(Path::new(value))
419    }
420}
421
422impl TryFrom<String> for BinaryGenomeAssembly {
423    type Error = GtarsGenomicDistError;
424    fn try_from(value: String) -> Result<Self, GtarsGenomicDistError> {
425        BinaryGenomeAssembly::from_file(Path::new(&value))
426    }
427}
428
429impl TryFrom<&Path> for BinaryGenomeAssembly {
430    type Error = GtarsGenomicDistError;
431    fn try_from(value: &Path) -> Result<Self, GtarsGenomicDistError> {
432        BinaryGenomeAssembly::from_file(value)
433    }
434}
435
436impl SequenceAccess for BinaryGenomeAssembly {
437    fn get_sequence(&self, coords: &Region) -> Result<Vec<u8>, GtarsGenomicDistError> {
438        self.seq_from_region(coords).map(|s| s.to_vec())
439    }
440
441    fn contains_chr(&self, chr: &str) -> bool {
442        self.index.contains_key(chr)
443    }
444}
445
446#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
447pub enum Dinucleotide {
448    Aa,
449    Ac,
450    Ag,
451    At,
452    Ca,
453    Cc,
454    Cg,
455    Ct,
456    Ga,
457    Gc,
458    Gg,
459    Gt,
460    Ta,
461    Tc,
462    Tg,
463    Tt,
464}
465
466impl Dinucleotide {
467    pub fn from_bytes(bytes: &[u8]) -> Option<Dinucleotide> {
468        if bytes.len() != 2 {
469            return None;
470        }
471        // Normalize to uppercase for case-insensitive matching
472        let normalized = [bytes[0].to_ascii_uppercase(), bytes[1].to_ascii_uppercase()];
473        match &normalized {
474            b"AA" => Some(Dinucleotide::Aa),
475            b"AC" => Some(Dinucleotide::Ac),
476            b"AG" => Some(Dinucleotide::Ag),
477            b"AT" => Some(Dinucleotide::At),
478            b"CA" => Some(Dinucleotide::Ca),
479            b"CC" => Some(Dinucleotide::Cc),
480            b"CG" => Some(Dinucleotide::Cg),
481            b"CT" => Some(Dinucleotide::Ct),
482            b"GA" => Some(Dinucleotide::Ga),
483            b"GC" => Some(Dinucleotide::Gc),
484            b"GG" => Some(Dinucleotide::Gg),
485            b"GT" => Some(Dinucleotide::Gt),
486            b"TA" => Some(Dinucleotide::Ta),
487            b"TC" => Some(Dinucleotide::Tc),
488            b"TG" => Some(Dinucleotide::Tg),
489            b"TT" => Some(Dinucleotide::Tt),
490            _ => None,
491        }
492    }
493
494    pub fn to_string(&self) -> Result<String, GtarsGenomicDistError> {
495        match self {
496            Dinucleotide::Aa => Ok("Aa".to_string()),
497            Dinucleotide::Ac => Ok("Ac".to_string()),
498            Dinucleotide::Ag => Ok("Ag".to_string()),
499            Dinucleotide::At => Ok("At".to_string()),
500            Dinucleotide::Ca => Ok("Ca".to_string()),
501            Dinucleotide::Cc => Ok("Cc".to_string()),
502            Dinucleotide::Cg => Ok("Cg".to_string()),
503            Dinucleotide::Ct => Ok("Ct".to_string()),
504            Dinucleotide::Ga => Ok("Ga".to_string()),
505            Dinucleotide::Gc => Ok("Gc".to_string()),
506            Dinucleotide::Gg => Ok("Gg".to_string()),
507            Dinucleotide::Gt => Ok("Gt".to_string()),
508            Dinucleotide::Ta => Ok("Ta".to_string()),
509            Dinucleotide::Tc => Ok("Tc".to_string()),
510            Dinucleotide::Tg => Ok("Tg".to_string()),
511            Dinucleotide::Tt => Ok("Tt".to_string()),
512        }
513    }
514}
515
516///
517/// Struct to hold Tss information (RegionSet with additionally indexing) that is initialized from
518/// RegionSet or BED file that holds tss regions
519///
520pub struct TssIndex {
521    pub region_set: RegionSet,
522    pub mid_points: HashMap<String, Vec<u32>>,
523}
524
525impl TryFrom<RegionSet> for TssIndex {
526    type Error = GtarsGenomicDistError;
527    fn try_from(value: RegionSet) -> Result<Self, GtarsGenomicDistError> {
528        TssIndex::from_region_set(value, CoordinateMode::Bed)
529    }
530}
531
532impl TssIndex {
533    /// Create a TssIndex from a RegionSet using the specified coordinate mode for midpoints.
534    pub fn from_region_set(
535        value: RegionSet,
536        mode: CoordinateMode,
537    ) -> Result<Self, GtarsGenomicDistError> {
538        let mut mid_points = value.calc_mid_points_with_mode(mode);
539
540        for points in mid_points.values_mut() {
541            points.sort_unstable();
542        }
543
544        Ok(TssIndex {
545            region_set: value,
546            mid_points,
547        })
548    }
549}
550
551impl TryFrom<&Path> for TssIndex {
552    type Error = GtarsGenomicDistError;
553    fn try_from(value: &Path) -> Result<Self, GtarsGenomicDistError> {
554        let region_set = match RegionSet::try_from(value) {
555            Ok(region_set) => region_set,
556            Err(_e) => {
557                return Err(GtarsGenomicDistError::TSSContentError(String::from(
558                    "Unable to open Tss file",
559                )));
560            }
561        };
562        TssIndex::try_from(region_set)
563    }
564}
565
566impl TryFrom<&str> for TssIndex {
567    type Error = GtarsGenomicDistError;
568    fn try_from(value: &str) -> Result<Self, GtarsGenomicDistError> {
569        let region_set = match RegionSet::try_from(value) {
570            Ok(region_set) => region_set,
571            Err(_e) => {
572                return Err(GtarsGenomicDistError::TSSContentError(String::from(
573                    "Unable to open Tss file",
574                )));
575            }
576        };
577        TssIndex::try_from(region_set)
578    }
579}
580
581impl TryFrom<String> for TssIndex {
582    type Error = GtarsGenomicDistError;
583    fn try_from(value: String) -> Result<Self, GtarsGenomicDistError> {
584        let region_set = match RegionSet::try_from(value) {
585            Ok(region_set) => region_set,
586            Err(_e) => {
587                return Err(GtarsGenomicDistError::TSSContentError(String::from(
588                    "Unable to open Tss file",
589                )));
590            }
591        };
592        TssIndex::try_from(region_set)
593    }
594}
595
596impl TssIndex {
597    ///
598    /// Calculate the distance from each region to the nearest TSS mid-point.
599    ///
600    /// Uses binary search for O(R * log M) complexity instead of O(R * M),
601    /// where R is the number of regions and M is the number of TSS midpoints.
602    ///
603    pub fn calc_tss_distances(
604        &self,
605        rs: &RegionSet,
606        mode: CoordinateMode,
607    ) -> Result<Vec<u32>, GtarsGenomicDistError> {
608        let mut distances: Vec<u32> = Vec::with_capacity(rs.len());
609
610        for chromosome in rs.iter_chroms() {
611            if let Some(chr_midpoints) = self.mid_points.get(chromosome.as_str()) {
612                for region in rs.iter_chr_regions(chromosome.as_str()) {
613                    let target = region.mid_point_with_mode(mode);
614
615                    let min_distance = match chr_midpoints.binary_search(&target) {
616                        Ok(_) => 0,
617                        Err(idx) => {
618                            let left = idx
619                                .checked_sub(1)
620                                .map(|i| target.abs_diff(chr_midpoints[i]));
621                            let right = chr_midpoints.get(idx).map(|&v| target.abs_diff(v));
622
623                            match (left, right) {
624                                (Some(l), Some(r)) => l.min(r),
625                                (Some(l), None) => l,
626                                (None, Some(r)) => r,
627                                (None, None) => continue,
628                            }
629                        }
630                    };
631                    distances.push(min_distance);
632                }
633            } else {
634                // No features on this chromosome — push u32::MAX for each region
635                for _ in rs.iter_chr_regions(chromosome.as_str()) {
636                    distances.push(u32::MAX);
637                }
638            }
639        }
640        Ok(distances)
641    }
642
643    /// Calculate signed distances from each region to its nearest feature.
644    ///
645    /// Like `calc_tss_distances` but returns signed distances where:
646    /// - Positive: nearest feature is downstream (right) of the query
647    /// - Negative: nearest feature is upstream (left) of the query
648    ///
649    /// Sign convention: `nearest_feature_midpoint - query_midpoint`
650    /// (matches R GenomicDistributions `calcFeatureDist()`).
651    pub fn calc_feature_distances(
652        &self,
653        rs: &RegionSet,
654        mode: CoordinateMode,
655    ) -> Result<Vec<i64>, GtarsGenomicDistError> {
656        let mut distances: Vec<i64> = Vec::with_capacity(rs.len());
657
658        for chromosome in rs.iter_chroms() {
659            if let Some(chr_midpoints) = self.mid_points.get(chromosome.as_str()) {
660                for region in rs.iter_chr_regions(chromosome.as_str()) {
661                    let target = region.mid_point_with_mode(mode) as i64;
662
663                    let distance = match chr_midpoints.binary_search(&(target as u32)) {
664                        Ok(_) => 0i64,
665                        Err(idx) => {
666                            // distance = feature_mid - query_mid
667                            let left = idx
668                                .checked_sub(1)
669                                .map(|i| chr_midpoints[i] as i64 - target);
670                            let right =
671                                chr_midpoints.get(idx).map(|&v| v as i64 - target);
672
673                            match (left, right) {
674                                (Some(l), Some(r)) => {
675                                    if l.unsigned_abs() <= r.unsigned_abs() {
676                                        l
677                                    } else {
678                                        r
679                                    }
680                                }
681                                (Some(l), None) => l,
682                                (None, Some(r)) => r,
683                                (None, None) => continue,
684                            }
685                        }
686                    };
687                    distances.push(distance);
688                }
689            } else {
690                // No features on this chromosome — push i64::MAX for each region
691                for _ in rs.iter_chr_regions(chromosome.as_str()) {
692                    distances.push(i64::MAX);
693                }
694            }
695        }
696        Ok(distances)
697    }
698}
699
700#[cfg(test)]
701mod tests {
702    use super::*;
703    use std::io::Error;
704    use std::path::PathBuf;
705
706    use pretty_assertions::assert_eq;
707    use rstest::*;
708
709    fn get_test_path(file_name: &str) -> Result<PathBuf, Error> {
710        let file_path: PathBuf = std::env::current_dir()
711            .unwrap()
712            .join("../tests/data/regionset")
713            .join(file_name);
714        Ok(file_path)
715    }
716
717    fn get_fasta_path(file_name: &str) -> PathBuf {
718        std::env::current_dir()
719            .unwrap()
720            .join("../tests/data/fasta")
721            .join(file_name)
722    }
723
724    // --- Strand ---
725
726    #[test]
727    fn test_strand_from_char() {
728        assert_eq!(Strand::from_char('+'), Strand::Plus);
729        assert_eq!(Strand::from_char('-'), Strand::Minus);
730        assert_eq!(Strand::from_char('.'), Strand::Unstranded);
731        assert_eq!(Strand::from_char('?'), Strand::Unstranded);
732    }
733
734    // --- Dinucleotide ---
735
736    #[test]
737    fn test_dinucleotide_from_bytes_all_variants() {
738        let pairs = [
739            (b"AA", Dinucleotide::Aa), (b"AC", Dinucleotide::Ac),
740            (b"AG", Dinucleotide::Ag), (b"AT", Dinucleotide::At),
741            (b"CA", Dinucleotide::Ca), (b"CC", Dinucleotide::Cc),
742            (b"CG", Dinucleotide::Cg), (b"CT", Dinucleotide::Ct),
743            (b"GA", Dinucleotide::Ga), (b"GC", Dinucleotide::Gc),
744            (b"GG", Dinucleotide::Gg), (b"GT", Dinucleotide::Gt),
745            (b"TA", Dinucleotide::Ta), (b"TC", Dinucleotide::Tc),
746            (b"TG", Dinucleotide::Tg), (b"TT", Dinucleotide::Tt),
747        ];
748        for (bytes, expected) in &pairs {
749            assert_eq!(Dinucleotide::from_bytes(&bytes[..]), Some(*expected));
750        }
751    }
752
753    #[test]
754    fn test_dinucleotide_case_insensitive() {
755        assert_eq!(Dinucleotide::from_bytes(b"aa"), Some(Dinucleotide::Aa));
756        assert_eq!(Dinucleotide::from_bytes(b"cG"), Some(Dinucleotide::Cg));
757        assert_eq!(Dinucleotide::from_bytes(b"Tc"), Some(Dinucleotide::Tc));
758    }
759
760    #[test]
761    fn test_dinucleotide_invalid() {
762        assert_eq!(Dinucleotide::from_bytes(b"AN"), None);
763        assert_eq!(Dinucleotide::from_bytes(b"A"), None);  // too short
764        assert_eq!(Dinucleotide::from_bytes(b"ACG"), None); // too long
765    }
766
767    #[test]
768    fn test_dinucleotide_to_string_round_trip() {
769        let all = [
770            Dinucleotide::Aa, Dinucleotide::Ac, Dinucleotide::Ag, Dinucleotide::At,
771            Dinucleotide::Ca, Dinucleotide::Cc, Dinucleotide::Cg, Dinucleotide::Ct,
772            Dinucleotide::Ga, Dinucleotide::Gc, Dinucleotide::Gg, Dinucleotide::Gt,
773            Dinucleotide::Ta, Dinucleotide::Tc, Dinucleotide::Tg, Dinucleotide::Tt,
774        ];
775        for d in &all {
776            let s = d.to_string().unwrap();
777            assert_eq!(s.len(), 2);
778            let round_tripped = Dinucleotide::from_bytes(s.as_bytes()).unwrap();
779            assert_eq!(*d, round_tripped);
780        }
781    }
782
783    // --- SortedRegionSet ---
784
785    #[test]
786    fn test_sorted_regionset_sorts_in_place() {
787        let regions = vec![
788            Region { chr: "chr1".into(), start: 100, end: 200, rest: None },
789            Region { chr: "chr1".into(), start: 10, end: 20, rest: None },
790            Region { chr: "chr2".into(), start: 5, end: 15, rest: None },
791        ];
792        let sorted = SortedRegionSet::new(RegionSet::from(regions));
793        let starts: Vec<u32> = sorted.0.regions.iter().map(|r| r.start).collect();
794        // chr1 regions should come first (sorted by chr then start)
795        assert_eq!(starts, vec![10, 100, 5]);
796    }
797
798    // --- StrandedRegionSet ---
799
800    #[test]
801    fn test_stranded_regionset_new() {
802        let regions = vec![
803            Region { chr: "chr1".into(), start: 10, end: 20, rest: None },
804            Region { chr: "chr1".into(), start: 30, end: 40, rest: None },
805        ];
806        let strands = vec![Strand::Plus, Strand::Minus];
807        let srs = StrandedRegionSet::new(RegionSet::from(regions), strands);
808        assert_eq!(srs.len(), 2);
809        assert!(!srs.is_empty());
810        assert_eq!(srs.strands[0], Strand::Plus);
811        assert_eq!(srs.strands[1], Strand::Minus);
812    }
813
814    #[test]
815    fn test_stranded_regionset_unstranded() {
816        let regions = vec![
817            Region { chr: "chr1".into(), start: 10, end: 20, rest: None },
818        ];
819        let srs = StrandedRegionSet::unstranded(RegionSet::from(regions));
820        assert_eq!(srs.strands, vec![Strand::Unstranded]);
821    }
822
823    #[test]
824    #[should_panic(expected = "regions and strands must have the same length")]
825    fn test_stranded_regionset_mismatched_lengths() {
826        let regions = vec![
827            Region { chr: "chr1".into(), start: 10, end: 20, rest: None },
828        ];
829        StrandedRegionSet::new(RegionSet::from(regions), vec![]);
830    }
831
832    #[test]
833    fn test_stranded_regionset_into_regionset() {
834        let regions = vec![
835            Region { chr: "chr1".into(), start: 10, end: 20, rest: None },
836        ];
837        let srs = StrandedRegionSet::unstranded(RegionSet::from(regions));
838        let rs = srs.into_regionset();
839        assert_eq!(rs.regions.len(), 1);
840    }
841
842    // --- GenomeAssembly ---
843
844    #[test]
845    fn test_genome_assembly_from_fasta() {
846        // base.fa: chrX=TTGGGGAA, chr1=GGAA, chr2=GCGC
847        let path = get_fasta_path("base.fa");
848        let ga = GenomeAssembly::try_from(path.as_path()).unwrap();
849        assert!(ga.contains_chr("chr1"));
850        assert!(ga.contains_chr("chr2"));
851        assert!(ga.contains_chr("chrX"));
852        assert!(!ga.contains_chr("chr3"));
853    }
854
855    #[test]
856    fn test_genome_assembly_seq_from_region() {
857        let path = get_fasta_path("base.fa");
858        let ga = GenomeAssembly::try_from(path.as_path()).unwrap();
859
860        let region = Region { chr: "chr1".into(), start: 0, end: 4, rest: None };
861        let seq = ga.seq_from_region(&region).unwrap();
862        assert_eq!(seq, b"GGAA");
863
864        let region2 = Region { chr: "chrX".into(), start: 2, end: 6, rest: None };
865        let seq2 = ga.seq_from_region(&region2).unwrap();
866        assert_eq!(seq2, b"GGGG");
867    }
868
869    #[test]
870    fn test_genome_assembly_seq_unknown_chrom() {
871        let path = get_fasta_path("base.fa");
872        let ga = GenomeAssembly::try_from(path.as_path()).unwrap();
873
874        let region = Region { chr: "chr99".into(), start: 0, end: 1, rest: None };
875        assert!(ga.seq_from_region(&region).is_err());
876    }
877
878    #[test]
879    fn test_genome_assembly_seq_out_of_bounds() {
880        let path = get_fasta_path("base.fa");
881        let ga = GenomeAssembly::try_from(path.as_path()).unwrap();
882
883        // chr1 is only 4bp, request 0-100
884        let region = Region { chr: "chr1".into(), start: 0, end: 100, rest: None };
885        assert!(ga.seq_from_region(&region).is_err());
886    }
887
888    #[test]
889    fn test_genome_assembly_try_from_str() {
890        let path = get_fasta_path("base.fa");
891        let ga = GenomeAssembly::try_from(path.to_str().unwrap());
892        assert!(ga.is_ok());
893    }
894
895    #[test]
896    fn test_genome_assembly_try_from_string() {
897        let path = get_fasta_path("base.fa");
898        let ga = GenomeAssembly::try_from(path.to_str().unwrap().to_string());
899        assert!(ga.is_ok());
900        assert!(ga.unwrap().contains_chr("chr1"));
901    }
902
903    #[test]
904    fn test_binary_genome_assembly_round_trip() {
905        // Write .fab from base.fa, then read it back and verify sequences match
906        let fasta_path = get_fasta_path("base.fa");
907        let fab_path = fasta_path.with_extension("fa.test.fab");
908
909        BinaryGenomeAssembly::write_from_fasta(&fasta_path, &fab_path).unwrap();
910        let bga = BinaryGenomeAssembly::from_file(&fab_path).unwrap();
911
912        // Verify chromosomes exist
913        assert!(bga.contains_chr("chr1"));
914        assert!(bga.contains_chr("chr2"));
915        assert!(bga.contains_chr("chrX"));
916        assert!(!bga.contains_chr("chr3"));
917
918        // Verify sequences match HashMap GenomeAssembly
919        let ga = GenomeAssembly::try_from(fasta_path.as_path()).unwrap();
920
921        let region1 = Region { chr: "chr1".into(), start: 0, end: 4, rest: None };
922        assert_eq!(bga.seq_from_region(&region1).unwrap(), ga.seq_from_region(&region1).unwrap());
923        assert_eq!(bga.seq_from_region(&region1).unwrap(), b"GGAA");
924
925        let region2 = Region { chr: "chrX".into(), start: 2, end: 6, rest: None };
926        assert_eq!(bga.seq_from_region(&region2).unwrap(), ga.seq_from_region(&region2).unwrap());
927        assert_eq!(bga.seq_from_region(&region2).unwrap(), b"GGGG");
928
929        // Out-of-bounds error
930        let bad_region = Region { chr: "chr1".into(), start: 0, end: 100, rest: None };
931        assert!(bga.seq_from_region(&bad_region).is_err());
932
933        // Unknown chromosome error
934        let unk_region = Region { chr: "chr99".into(), start: 0, end: 1, rest: None };
935        assert!(bga.seq_from_region(&unk_region).is_err());
936
937        // Clean up
938        std::fs::remove_file(&fab_path).ok();
939    }
940
941    #[test]
942    fn test_binary_genome_assembly_bad_magic() {
943        let dir = tempfile::tempdir().unwrap();
944        let fab_path = dir.path().join("bad.fab");
945        // Need at least 9 bytes to pass the "too short" check
946        std::fs::write(&fab_path, b"XXXX\x01\x00\x00\x00\x00").unwrap();
947        let result = BinaryGenomeAssembly::from_file(&fab_path);
948        assert!(result.is_err());
949        assert!(result.unwrap_err().to_string().contains("bad magic"));
950    }
951
952    #[test]
953    fn test_binary_genome_assembly_gc_parity() {
954        // Verify calc_gc_content produces identical results from .fab and HashMap
955        use crate::statistics::calc_gc_content;
956
957        let fasta_path = get_fasta_path("base.fa");
958        let fab_path = fasta_path.with_extension("fa.test2.fab");
959        BinaryGenomeAssembly::write_from_fasta(&fasta_path, &fab_path).unwrap();
960
961        let ga = GenomeAssembly::try_from(fasta_path.as_path()).unwrap();
962        let bga = BinaryGenomeAssembly::from_file(&fab_path).unwrap();
963
964        let regions = vec![
965            Region { chr: "chr1".into(), start: 0, end: 4, rest: None },
966            Region { chr: "chr2".into(), start: 0, end: 4, rest: None },
967        ];
968        let rs = RegionSet::from(regions);
969
970        let gc_hashmap = calc_gc_content(&rs, &ga, false).unwrap();
971        let gc_fab = calc_gc_content(&rs, &bga, false).unwrap();
972        assert_eq!(gc_hashmap, gc_fab);
973
974        std::fs::remove_file(&fab_path).ok();
975    }
976
977    #[test]
978    fn test_tss_index_try_from_path() {
979        let path = get_test_path("dummy_tss.bed").unwrap();
980        let tss = TssIndex::try_from(path.as_path());
981        assert!(tss.is_ok());
982    }
983
984    #[test]
985    fn test_tss_index_try_from_path_invalid() {
986        let path = PathBuf::from("/nonexistent/file.bed");
987        let tss = TssIndex::try_from(path.as_path());
988        assert!(tss.is_err());
989    }
990
991    #[test]
992    fn test_tss_index_try_from_string() {
993        let path = get_test_path("dummy_tss.bed").unwrap();
994        let tss = TssIndex::try_from(path.to_str().unwrap().to_string());
995        assert!(tss.is_ok());
996    }
997
998    // --- TssIndex sentinel behavior ---
999
1000    #[test]
1001    fn test_tss_distances_sentinel_for_missing_chrom() {
1002        // TSS features only on chr1
1003        let tss_regions = vec![
1004            Region { chr: "chr1".into(), start: 50, end: 51, rest: None },
1005        ];
1006        let tss_index = TssIndex::try_from(RegionSet::from(tss_regions)).unwrap();
1007
1008        // Query has regions on chr1 and chr2 (no TSS on chr2)
1009        let query = RegionSet::from(vec![
1010            Region { chr: "chr1".into(), start: 40, end: 45, rest: None },
1011            Region { chr: "chr2".into(), start: 10, end: 20, rest: None },
1012        ]);
1013
1014        let distances = tss_index.calc_tss_distances(&query, CoordinateMode::Bed).unwrap();
1015        assert_eq!(distances.len(), 2); // one per input region
1016        // One should be a real distance, the other should be u32::MAX sentinel
1017        // (order depends on HashSet iteration of iter_chroms)
1018        assert_eq!(distances.iter().filter(|&&d| d == u32::MAX).count(), 1);
1019        assert_eq!(distances.iter().filter(|&&d| d < u32::MAX).count(), 1);
1020    }
1021
1022    #[test]
1023    fn test_feature_distances_sentinel_for_missing_chrom() {
1024        let tss_regions = vec![
1025            Region { chr: "chr1".into(), start: 50, end: 51, rest: None },
1026        ];
1027        let tss_index = TssIndex::try_from(RegionSet::from(tss_regions)).unwrap();
1028
1029        let query = RegionSet::from(vec![
1030            Region { chr: "chr1".into(), start: 40, end: 45, rest: None },
1031            Region { chr: "chr2".into(), start: 10, end: 20, rest: None },
1032        ]);
1033
1034        let distances = tss_index.calc_feature_distances(&query, CoordinateMode::Bed).unwrap();
1035        assert_eq!(distances.len(), 2);
1036        // One real distance, one i64::MAX sentinel
1037        assert_eq!(distances.iter().filter(|&&d| d == i64::MAX).count(), 1);
1038        assert_eq!(distances.iter().filter(|&&d| d != i64::MAX).count(), 1);
1039    }
1040
1041    // --- Existing tests ---
1042
1043    #[rstest]
1044    fn test_calc_tss_distances() {
1045        let file_path = get_test_path("dummy.narrowPeak").unwrap();
1046        let tss_path = get_test_path("dummy_tss.bed").unwrap();
1047        let region_set = RegionSet::try_from(file_path.to_str().unwrap()).unwrap();
1048        let tss_index = TssIndex::try_from(tss_path.to_str().unwrap()).unwrap();
1049
1050        let distances = tss_index.calc_tss_distances(&region_set, CoordinateMode::Bed).unwrap();
1051
1052        assert_eq!(distances.len(), 9);
1053        assert_eq!(distances.iter().min(), Some(&2));
1054    }
1055
1056    #[rstest]
1057    fn test_calc_feature_distances() {
1058        let file_path = get_test_path("dummy.narrowPeak").unwrap();
1059        let tss_path = get_test_path("dummy_tss.bed").unwrap();
1060        let region_set = RegionSet::try_from(file_path.to_str().unwrap()).unwrap();
1061        let tss_index = TssIndex::try_from(tss_path.to_str().unwrap()).unwrap();
1062
1063        let signed_distances = tss_index.calc_feature_distances(&region_set, CoordinateMode::Bed).unwrap();
1064        let abs_distances = tss_index.calc_tss_distances(&region_set, CoordinateMode::Bed).unwrap();
1065
1066        // same number of results
1067        assert_eq!(signed_distances.len(), abs_distances.len());
1068        // absolute values should match calc_tss_distances
1069        for (signed, abs) in signed_distances.iter().zip(abs_distances.iter()) {
1070            assert_eq!(signed.unsigned_abs() as u32, *abs);
1071        }
1072        // should have both positive and negative distances
1073        assert!(signed_distances.iter().any(|d| *d > 0));
1074        assert!(signed_distances.iter().any(|d| *d < 0));
1075    }
1076}