Skip to main content

orphos_core/sequence/
mod.rs

1//! Sequence encoding and manipulation utilities.
2//!
3//! This module provides functions for encoding DNA sequences into compact bitmap
4//! representations and performing sequence analysis operations.
5//!
6//! ## Overview
7//!
8//! DNA sequences are encoded using a 2-bit representation where:
9//! - A (adenine): 00
10//! - C (cytosine): 01
11//! - G (guanine): 10
12//! - T/U (thymine/uracil): 11
13//!
14//! This encoding reduces memory usage by 75% compared to ASCII representation
15//! and enables fast bitwise operations for sequence analysis.
16//!
17//! ## Modules
18//!
19//! - [`encoded`]: Encoded sequence structures with forward and reverse-complement
20//! - [`io`]: FASTA file reading and parsing
21//! - [`processing`]: Sequence analysis functions (GC content, codon detection)
22//!
23//! ## Examples
24//!
25//! ### Encode a sequence
26//!
27//! ```rust
28//! use orphos_core::sequence::encoded::EncodedSequence;
29//!
30//! let sequence = b"ATGAAACGCATTAGCACCACCATT";
31//! let encoded = EncodedSequence::without_masking(sequence);
32//!
33//! println!("Length: {} bp", encoded.sequence_length);
34//! println!("GC content: {:.2}%", encoded.gc_content * 100.0);
35//! ```
36//!
37//! ### Test for specific nucleotides
38//!
39//! ```rust
40//! use orphos_core::sequence::{is_a, is_gc};
41//! use orphos_core::sequence::encoded::EncodedSequence;
42//!
43//! let sequence = b"ATGC";
44//! let encoded = EncodedSequence::without_masking(sequence);
45//!
46//! assert!(is_a(&encoded.forward_sequence, 0)); // Position 0 is 'A'
47//! assert!(is_gc(&encoded.forward_sequence, 3)); // Position 3 is 'C'
48//! ```
49
50use crate::bitmap;
51use crate::constants::MASK_SIZE;
52use crate::types::*;
53
54pub mod encoded;
55pub mod io;
56pub mod processing;
57
58use crate::bitmap::test_bit;
59use rayon::prelude::*;
60
61pub use io::*;
62pub use processing::*;
63use wide::u8x32;
64
65/// Converts nucleotide character to 2-bit encoding for bitmap storage.
66///
67/// Maps nucleotides to compact 2-bit representation for efficient storage
68/// and fast sequence operations.
69///
70/// # Encoding
71///
72/// - A: 00 (0)
73/// - C: 01 (1)
74/// - G: 10 (2)
75/// - T/U: 11 (3)
76/// - Other: 4 (invalid marker)
77///
78/// # Arguments
79///
80/// * `c` - ASCII nucleotide character (case-insensitive)
81///
82/// # Returns
83///
84/// A value 0-3 for valid nucleotides, 4 for invalid characters.
85///
86/// # Examples
87///
88/// ```rust
89/// use orphos_core::sequence::char_to_nuc;
90///
91/// assert_eq!(char_to_nuc(b'A'), 0);
92/// assert_eq!(char_to_nuc(b'a'), 0);
93/// assert_eq!(char_to_nuc(b'C'), 1);
94/// assert_eq!(char_to_nuc(b'G'), 2);
95/// assert_eq!(char_to_nuc(b'T'), 3);
96/// assert_eq!(char_to_nuc(b'N'), 4); // Invalid
97/// ```
98#[must_use]
99pub const fn char_to_nuc(c: u8) -> u8 {
100    match c.to_ascii_uppercase() {
101        b'A' => 0,
102        b'C' => 1,
103        b'G' => 2,
104        b'T' | b'U' => 3,
105        _ => 4,
106    }
107}
108
109const ENCODED_A: u8 = 0b00;
110const ENCODED_G: u8 = 0b01;
111const ENCODED_C: u8 = 0b10;
112const ENCODED_T: u8 = 0b11;
113
114/// Test if nucleotide at given position is adenine (A)
115#[must_use]
116#[inline]
117pub fn is_a(encoded_sequence: &[u8], n: usize) -> bool {
118    encoded_base_code(encoded_sequence, n) == ENCODED_A
119}
120
121/// Test if nucleotide at given position is cytosine (C)
122#[must_use]
123#[inline]
124pub fn is_c(encoded_sequence: &[u8], n: usize) -> bool {
125    encoded_base_code(encoded_sequence, n) == ENCODED_C
126}
127
128/// Test if nucleotide at given position is guanine (G)
129#[must_use]
130#[inline]
131pub fn is_g(encoded_sequence: &[u8], n: usize) -> bool {
132    encoded_base_code(encoded_sequence, n) == ENCODED_G
133}
134
135/// Test if nucleotide at given position is thymine (T)
136#[must_use]
137#[inline]
138pub fn is_t(encoded_sequence: &[u8], n: usize) -> bool {
139    encoded_base_code(encoded_sequence, n) == ENCODED_T
140}
141
142/// Test if position contains an unknown nucleotide (N)
143pub fn is_n(unknown_sequence: &[u8], n: usize) -> bool {
144    if n >= unknown_sequence.len() * 8 {
145        return false;
146    }
147    test_bit(unknown_sequence, n)
148}
149
150/// Test if nucleotide at given position is G or C (high GC content indicator)
151#[inline]
152pub fn is_gc(encoded_sequence: &[u8], n: usize) -> bool {
153    matches!(
154        encoded_base_code(encoded_sequence, n),
155        ENCODED_G | ENCODED_C
156    )
157}
158
159/// Check if genetic code table uses only ATG as start codon
160const fn uses_only_atg(trans_table: i32) -> bool {
161    matches!(trans_table, 6 | 10 | 14 | 15 | 16 | 22)
162}
163
164/// Check if GTG is not used as start codon in given translation table
165const fn gtg_not_start(trans_table: i32) -> bool {
166    matches!(trans_table, 1 | 3 | 12 | 22)
167}
168
169/// Check if TTG is not used as start codon in given translation table
170fn ttg_not_start(trans_table: i32) -> bool {
171    trans_table < 4 || trans_table == 9 || (21..25).contains(&trans_table)
172}
173
174/// Test if codon at given position is a valid start codon
175///
176/// Checks ATG, GTG, and TTG based on the genetic code table rules.
177/// ATG is universally accepted as a start codon across all tables.
178pub fn is_start(encoded_sequence: &[u8], pos: usize, training: &Training) -> bool {
179    // ATG is always a start
180    if is_atg(encoded_sequence, pos) {
181        return true;
182    }
183
184    // Tables that only use ATG
185    if uses_only_atg(training.translation_table) {
186        return false;
187    }
188
189    // GTG
190    if is_gtg(encoded_sequence, pos) && !gtg_not_start(training.translation_table) {
191        return true;
192    }
193
194    // TTG
195    if is_ttg(encoded_sequence, pos) && !ttg_not_start(training.translation_table) {
196        return true;
197    }
198
199    false
200}
201
202/// Test if codon at position is ATG (methionine start codon)
203#[inline]
204pub fn is_atg(encoded_sequence: &[u8], pos: usize) -> bool {
205    encoded_base_code(encoded_sequence, pos) == ENCODED_A
206        && encoded_base_code(encoded_sequence, pos + 1) == ENCODED_T
207        && encoded_base_code(encoded_sequence, pos + 2) == ENCODED_G
208}
209
210/// Test if codon at position is GTG (valine start codon)
211#[inline]
212pub fn is_gtg(encoded_sequence: &[u8], pos: usize) -> bool {
213    encoded_base_code(encoded_sequence, pos) == ENCODED_G
214        && encoded_base_code(encoded_sequence, pos + 1) == ENCODED_T
215        && encoded_base_code(encoded_sequence, pos + 2) == ENCODED_G
216}
217
218/// Test if codon at position is TTG (leucine start codon)
219#[inline]
220pub fn is_ttg(encoded_sequence: &[u8], pos: usize) -> bool {
221    encoded_base_code(encoded_sequence, pos) == ENCODED_T
222        && encoded_base_code(encoded_sequence, pos + 1) == ENCODED_T
223        && encoded_base_code(encoded_sequence, pos + 2) == ENCODED_G
224}
225
226/// Check if TAG is recognized as stop codon in the given translation table
227const fn is_tag_stop(trans_table: i32) -> bool {
228    !matches!(trans_table, 6 | 15 | 16 | 22)
229}
230
231/// Check if TGA is recognized as stop codon in the given translation table
232const fn is_tga_stop(trans_table: i32) -> bool {
233    !matches!(trans_table, 2..=5 | 9 | 10 | 13 | 14 | 21 | 25)
234}
235
236/// Check if TAA is recognized as stop codon in the given translation table
237const fn is_taa_stop(trans_table: i32) -> bool {
238    !matches!(trans_table, 6 | 14)
239}
240
241/// Test if codon at given position is a stop codon
242///
243/// Checks for standard stop codons (TAA, TAG, TGA) and special cases
244/// based on the genetic code translation table being used.
245#[inline]
246pub fn is_stop(encoded_sequence: &[u8], pos: usize, training: &Training) -> bool {
247    let first = encoded_base_code(encoded_sequence, pos);
248
249    if first == ENCODED_T {
250        let second = encoded_base_code(encoded_sequence, pos + 1);
251
252        if second == ENCODED_A {
253            let third = encoded_base_code(encoded_sequence, pos + 2);
254            if third == ENCODED_G {
255                return is_tag_stop(training.translation_table);
256            }
257            if third == ENCODED_A {
258                return is_taa_stop(training.translation_table);
259            }
260        } else if second == ENCODED_G {
261            return encoded_base_code(encoded_sequence, pos + 2) == ENCODED_A
262                && is_tga_stop(training.translation_table);
263        } else if (training.translation_table == 22 && second == ENCODED_C)
264            || (training.translation_table == 23 && second == ENCODED_T)
265        {
266            return encoded_base_code(encoded_sequence, pos + 2) == ENCODED_A;
267        }
268
269        return false;
270    }
271
272    if training.translation_table == 2 && first == ENCODED_A {
273        let second = encoded_base_code(encoded_sequence, pos + 1);
274        if second != ENCODED_G {
275            return false;
276        }
277
278        return matches!(
279            encoded_base_code(encoded_sequence, pos + 2),
280            ENCODED_A | ENCODED_G
281        );
282    }
283
284    false
285}
286
287/// Calculate the GC content of a sequence region
288///
289/// Returns the fraction of nucleotides that are G or C within
290/// the specified range (inclusive).
291pub fn gc_content(encoded_sequence: &[u8], start: usize, end: usize) -> f64 {
292    if start > end {
293        return 0.0;
294    }
295
296    let total = end - start + 1;
297    let gc_count = (start..=end)
298        .filter(|&position| is_gc(encoded_sequence, position))
299        .count();
300
301    gc_count as f64 / total as f64
302}
303
304/// Convert a forward strand reading frame to its corresponding reverse strand frame
305///
306/// Maps reading frames between forward and reverse strands accounting for
307/// sequence length and frame relationships.
308pub const fn reverse_strand_reading_frame(forward_frame: usize, sequence_length: usize) -> usize {
309    let frame_modulus = if sequence_length.is_multiple_of(3) {
310        3
311    } else {
312        sequence_length % 3
313    };
314    (frame_modulus - 1 - forward_frame) % 3
315}
316
317/// Determine which of three reading frames has the highest score
318///
319/// Returns the index (0, 1, or 2) of the frame with maximum value.
320pub const fn find_max_reading_frame(
321    frame_0_value: i32,
322    frame_1_value: i32,
323    frame_2_value: i32,
324) -> usize {
325    if frame_0_value > frame_1_value {
326        if frame_0_value > frame_2_value { 0 } else { 2 }
327    } else if frame_1_value > frame_2_value {
328        1
329    } else {
330        2
331    }
332}
333
334#[inline]
335const fn encoded_base_code(encoded_sequence: &[u8], position: usize) -> u8 {
336    let bit_index = position * 2;
337    (encoded_sequence[bit_index >> 3] >> (bit_index & 0x07)) & 0b11
338}
339
340#[inline]
341fn write_encoded_base_code(encoded_sequence: &mut [u8], position: usize, code: u8) {
342    let bit_index = position * 2;
343    let byte_index = bit_index >> 3;
344    let shift = bit_index & 0x07;
345    encoded_sequence[byte_index] =
346        (encoded_sequence[byte_index] & !(0b11 << shift)) | ((code & 0b11) << shift);
347}
348
349#[inline]
350const fn reverse_complement_byte(byte: u8) -> u8 {
351    let base_0 = (!byte >> 6) & 0b11;
352    let base_1 = (!byte >> 4) & 0b11;
353    let base_2 = (!byte >> 2) & 0b11;
354    let base_3 = !byte & 0b11;
355
356    base_0 | (base_1 << 2) | (base_2 << 4) | (base_3 << 6)
357}
358
359fn copy_unknown_encoded_bases(
360    reverse_complement_encoded_sequence: &mut [u8],
361    forward_sequence: &[u8],
362    unknown_sequence: &[u8],
363    nucleotide_length: usize,
364) {
365    for (byte_index, &unknown_byte) in unknown_sequence.iter().enumerate() {
366        let mut remaining_unknowns = unknown_byte;
367
368        while remaining_unknowns != 0 {
369            let bit_offset = remaining_unknowns.trailing_zeros() as usize;
370            let source_position = byte_index * 8 + bit_offset;
371
372            if source_position >= nucleotide_length {
373                break;
374            }
375
376            let target_position = nucleotide_length - source_position - 1;
377            let source_code = encoded_base_code(forward_sequence, source_position);
378            write_encoded_base_code(
379                reverse_complement_encoded_sequence,
380                target_position,
381                source_code,
382            );
383
384            remaining_unknowns &= remaining_unknowns - 1;
385        }
386    }
387}
388
389/// Generate the reverse complement of an encoded DNA sequence
390///
391/// Creates a reverse complement sequence using 2-bit encoding,
392/// handling both known and unknown nucleotides properly.
393pub fn create_reverse_complement_sequence(
394    forward_sequence: &[u8],
395    unknown_sequence: &[u8],
396    nucleotide_length: usize,
397) -> Vec<u8> {
398    let mut reverse_complement_encoded_sequence = vec![0; forward_sequence.len()];
399
400    if nucleotide_length.is_multiple_of(4) {
401        let encoded_bytes = nucleotide_length / 4;
402
403        for (source_byte_index, &source_byte) in
404            forward_sequence.iter().take(encoded_bytes).enumerate()
405        {
406            let target_byte_index = encoded_bytes - source_byte_index - 1;
407            reverse_complement_encoded_sequence[target_byte_index] =
408                reverse_complement_byte(source_byte);
409        }
410
411        copy_unknown_encoded_bases(
412            &mut reverse_complement_encoded_sequence,
413            forward_sequence,
414            unknown_sequence,
415            nucleotide_length,
416        );
417    } else {
418        for source_position in 0..nucleotide_length {
419            let target_position = nucleotide_length - source_position - 1;
420            let source_code = encoded_base_code(forward_sequence, source_position);
421            let target_code = if test_bit(unknown_sequence, source_position) {
422                source_code
423            } else {
424                source_code ^ 0b11
425            };
426
427            write_encoded_base_code(
428                &mut reverse_complement_encoded_sequence,
429                target_position,
430                target_code,
431            );
432        }
433    }
434
435    reverse_complement_encoded_sequence
436}
437
438/// Return the minimum of two integers (utility function)
439#[inline]
440pub fn min_of_two_integers(first_value: i32, second_value: i32) -> i32 {
441    first_value.min(second_value)
442}
443
444/// Calculate k-mer index from sequence position for frequency analysis
445///
446/// Converts a sequence position to a numeric index representing the
447/// k-mer pattern, used for codon usage and frequency calculations.
448#[must_use]
449#[inline]
450pub fn calculate_kmer_index(kmer_length: usize, encoded_sequence: &[u8], position: usize) -> usize {
451    let mut kmer_index = 0;
452    for offset in 0..kmer_length {
453        kmer_index |=
454            usize::from(encoded_base_code(encoded_sequence, position + offset)) << (offset * 2);
455    }
456    kmer_index
457}
458
459/// Calculate background k-mer frequencies for both strands
460///
461/// Computes frequency distributions of k-mers across the entire sequence,
462/// used for statistical modeling of codon usage patterns.
463pub fn calculate_background_mer_frequencies(
464    length: usize,
465    encoded_sequence: &[u8],
466    reverse_complement_encoded_sequence: &[u8],
467    sequence_length: usize,
468    bg: &mut [f64],
469) {
470    let mut size = 1usize;
471
472    for _i in 1..=length {
473        size *= 4;
474    }
475
476    // Use parallel processing to count k-mers
477    let chunk_size = std::cmp::max(
478        1000,
479        (sequence_length - length + 1) / rayon::current_num_threads(),
480    );
481    let total_counts: Vec<i32> = (0..(sequence_length - length + 1))
482        .into_par_iter()
483        .chunks(chunk_size)
484        .map(|chunk| {
485            let mut local_counts = vec![0i32; size];
486
487            for i in chunk {
488                let seq_idx = calculate_kmer_index(length, encoded_sequence, i);
489                if seq_idx < size {
490                    local_counts[seq_idx] += 1;
491                }
492
493                let rseq_idx = calculate_kmer_index(length, reverse_complement_encoded_sequence, i);
494                if rseq_idx < size {
495                    local_counts[rseq_idx] += 1;
496                }
497            }
498
499            local_counts
500        })
501        .reduce(
502            || vec![0i32; size],
503            |mut acc, local_counts| {
504                for (i, &count) in local_counts.iter().enumerate() {
505                    acc[i] += count;
506                }
507                acc
508            },
509        );
510
511    let glob = (sequence_length - length + 1) * 2;
512
513    bg.par_iter_mut()
514        .enumerate()
515        .take(size)
516        .for_each(|(i, bg_val)| {
517            *bg_val = f64::from(total_counts[i]) / (glob as f64);
518        });
519}
520
521/// Convert k-mer index back to nucleotide sequence representation
522///
523/// Decodes a numeric k-mer index back to its original DNA sequence
524/// for display and debugging purposes.
525pub fn mer_text(len: usize, bit_index: usize) -> String {
526    use crate::constants::NUCLEOTIDE_LETTERS;
527
528    if len == 0 {
529        return "None".to_string();
530    }
531
532    let mut result = String::with_capacity(len);
533    let index = bit_index;
534
535    for i in 0..len {
536        // Extract 2 bits for position i
537        let val = (index & (1 << (2 * i))) | (index & (1 << (2 * i + 1)));
538        let val = val >> (i * 2);
539        let base_idx = val & 0b11; // Ensure we only get 2 bits
540
541        if base_idx < 4 {
542            result.push(NUCLEOTIDE_LETTERS[base_idx]);
543        } else {
544            result.push('N'); // Fallback for invalid values
545        }
546    }
547
548    result
549}
550
551/// Encode a DNA sequence into compact 2-bit representation
552///
553/// Converts raw DNA sequence to bitmap format for efficient storage and processing.
554/// Also handles masking of low-complexity regions and tracks unknown nucleotides.
555/// Returns the GC content of the encoded sequence.
556pub fn encode_sequence(
557    sequence: &[u8],
558    encoded_sequence: &mut [u8],
559    unknown_sequence: &mut [u8],
560    masks: &mut Vec<Mask>,
561    do_mask: bool,
562) -> Result<f64, OrphosError> {
563    let mut gc_count = 0;
564    let mut total_count = 0;
565    let mut mask_start: Option<usize> = None;
566
567    for (i, &byte) in sequence.iter().enumerate() {
568        if i * 2 + 1 >= encoded_sequence.len() * 8 {
569            break;
570        }
571
572        // Handle masking for runs of N's
573        if do_mask {
574            if let Some(start) = mask_start {
575                if byte != b'N' && byte != b'n' {
576                    if i - start >= MASK_SIZE {
577                        masks.push(Mask {
578                            begin: start,
579                            end: i - 1,
580                        });
581                    }
582                    mask_start = None;
583                }
584            } else if byte == b'N' || byte == b'n' {
585                mask_start = Some(i);
586            }
587        }
588
589        // Encode nucleotide in bitmap format
590        let bctr = i * 2;
591        match byte.to_ascii_uppercase() {
592            b'A' => {
593                total_count += 1;
594            }
595            b'C' => {
596                bitmap::set_bit(encoded_sequence, bctr + 1);
597                gc_count += 1;
598                total_count += 1;
599            }
600            b'G' => {
601                bitmap::set_bit(encoded_sequence, bctr);
602                gc_count += 1;
603                total_count += 1;
604            }
605            b'T' | b'U' => {
606                bitmap::set_bit(encoded_sequence, bctr);
607                bitmap::set_bit(encoded_sequence, bctr + 1);
608                total_count += 1;
609            }
610            _ => {
611                bitmap::set_bit(encoded_sequence, bctr + 1);
612                bitmap::set_bit(unknown_sequence, i);
613                total_count += 1;
614            }
615        }
616    }
617
618    // Handle final mask if sequence ends with N's
619    if do_mask
620        && let Some(start) = mask_start
621        && sequence.len() - start >= MASK_SIZE
622    {
623        masks.push(Mask {
624            begin: start,
625            end: sequence.len() - 1,
626        });
627    }
628
629    let gc_content = if total_count > 0 {
630        f64::from(gc_count) / f64::from(total_count)
631    } else {
632        0.0
633    };
634
635    Ok(gc_content)
636}
637
638#[inline]
639const fn expand_4_lanes_to_even_bits(bits: u32) -> u8 {
640    ((bits & 0x1) | ((bits & 0x2) << 1) | ((bits & 0x4) << 2) | ((bits & 0x8) << 3)) as u8
641}
642
643/// Optimized packed encoding version with u8x32 and batch bit operations
644pub fn encode_sequence_simd_wide_packed(
645    sequence: &[u8],
646    encoded_sequence: &mut [u8],
647    unknown_sequence: &mut [u8],
648) -> Result<f64, OrphosError> {
649    let mut gc_count = 0u32;
650    let mut total_count = 0u32;
651
652    // Process 32 bytes at a time
653    use crate::constants::CHUNK_SIZE;
654    let chunks = sequence.len() / CHUNK_SIZE;
655
656    // SIMD splat constants hoisted outside the loop
657    let a_upper = u8x32::splat(b'A');
658    let c_upper = u8x32::splat(b'C');
659    let g_upper = u8x32::splat(b'G');
660    let t_upper = u8x32::splat(b'T');
661    let u_upper = u8x32::splat(b'U');
662
663    let a_lower = u8x32::splat(b'a');
664    let c_lower = u8x32::splat(b'c');
665    let g_lower = u8x32::splat(b'g');
666    let t_lower = u8x32::splat(b't');
667    let u_lower = u8x32::splat(b'u');
668
669    for chunk_idx in 0..chunks {
670        let chunk_start = chunk_idx * CHUNK_SIZE;
671
672        // Load 32 bytes
673        let input_slice = &sequence[chunk_start..chunk_start + CHUNK_SIZE];
674        let mut input_array = [0u8; 32];
675        input_array.copy_from_slice(input_slice);
676        let input = u8x32::from(input_array);
677
678        // SIMD nucleotide detection
679        let is_a = input.simd_eq(a_upper) | input.simd_eq(a_lower);
680        let is_c = input.simd_eq(c_upper) | input.simd_eq(c_lower);
681        let is_g = input.simd_eq(g_upper) | input.simd_eq(g_lower);
682        let is_t = input.simd_eq(t_upper)
683            | input.simd_eq(t_lower)
684            | input.simd_eq(u_upper)
685            | input.simd_eq(u_lower);
686
687        let gc_mask = is_g | is_c;
688        let valid_mask = is_a | is_c | is_g | is_t;
689
690        gc_count += gc_mask.to_bitmask().count_ones();
691        total_count += CHUNK_SIZE as u32;
692
693        // Extract SIMD results for bit setting - to_bitmask() returns u32
694        // let is_a_mask: u32 = is_a.to_bitmask();
695        let is_c_mask: u32 = is_c.to_bitmask();
696        let is_g_mask: u32 = is_g.to_bitmask();
697        let is_t_mask: u32 = is_t.to_bitmask();
698        let unknown_mask: u32 = !valid_mask.to_bitmask();
699
700        let encoded_byte_start = chunk_start / 4;
701        let unknown_byte_start = chunk_start / 8;
702
703        if encoded_byte_start + 8 <= encoded_sequence.len()
704            && unknown_byte_start + 4 <= unknown_sequence.len()
705        {
706            for byte_offset in 0..8 {
707                let shift = byte_offset * 4;
708                let even_bits =
709                    expand_4_lanes_to_even_bits((is_g_mask | is_t_mask) >> shift & 0x0f);
710                let odd_bits = expand_4_lanes_to_even_bits(
711                    (is_c_mask | is_t_mask | unknown_mask) >> shift & 0x0f,
712                ) << 1;
713                encoded_sequence[encoded_byte_start + byte_offset] |= even_bits | odd_bits;
714            }
715
716            for byte_offset in 0..4 {
717                unknown_sequence[unknown_byte_start + byte_offset] |=
718                    (unknown_mask >> (byte_offset * 8)) as u8;
719            }
720        } else {
721            for i in 0..CHUNK_SIZE {
722                let pos = chunk_start + i;
723                if pos >= sequence.len() || pos * 2 + 1 >= encoded_sequence.len() * 8 {
724                    break;
725                }
726
727                let bit_pos = pos * 2;
728                let bit_flag = 1u32 << i;
729
730                if (is_c_mask & bit_flag) != 0 {
731                    crate::bitmap::set_bit(encoded_sequence, bit_pos + 1);
732                } else if (is_g_mask & bit_flag) != 0 {
733                    crate::bitmap::set_bit(encoded_sequence, bit_pos);
734                } else if (is_t_mask & bit_flag) != 0 {
735                    crate::bitmap::set_bit(encoded_sequence, bit_pos);
736                    crate::bitmap::set_bit(encoded_sequence, bit_pos + 1);
737                } else if (unknown_mask & bit_flag) != 0 {
738                    crate::bitmap::set_bit(encoded_sequence, bit_pos + 1);
739                    crate::bitmap::set_bit(unknown_sequence, pos);
740                }
741            }
742        }
743    }
744
745    // Handle remaining bytes with scalar fallback
746    // for pos in (chunks * CHUNK_SIZE)..sequence.len() {
747    for (pos, byte) in sequence.iter().enumerate().skip(chunks * CHUNK_SIZE) {
748        if pos * 2 + 1 >= encoded_sequence.len() * 8 {
749            break;
750        }
751
752        // let byte = sequence[pos];
753        let bit_pos = pos * 2;
754
755        match byte.to_ascii_uppercase() {
756            b'A' => {
757                total_count += 1;
758            }
759            b'C' => {
760                crate::bitmap::set_bit(encoded_sequence, bit_pos + 1);
761                gc_count += 1;
762                total_count += 1;
763            }
764            b'G' => {
765                crate::bitmap::set_bit(encoded_sequence, bit_pos);
766                gc_count += 1;
767                total_count += 1;
768            }
769            b'T' | b'U' => {
770                crate::bitmap::set_bit(encoded_sequence, bit_pos);
771                crate::bitmap::set_bit(encoded_sequence, bit_pos + 1);
772                total_count += 1;
773            }
774            _ => {
775                crate::bitmap::set_bit(encoded_sequence, bit_pos + 1);
776                crate::bitmap::set_bit(unknown_sequence, pos);
777                total_count += 1;
778            }
779        }
780    }
781
782    let gc_content = if total_count > 0 {
783        gc_count as f64 / total_count as f64
784    } else {
785        0.0
786    };
787
788    Ok(gc_content)
789}
790
791#[cfg(test)]
792mod tests {
793    use super::*;
794    use crate::types::Training;
795
796    #[test]
797    fn test_char_to_nuc_valid_bases() {
798        assert_eq!(char_to_nuc(b'A'), 0);
799        assert_eq!(char_to_nuc(b'a'), 0);
800        assert_eq!(char_to_nuc(b'C'), 1);
801        assert_eq!(char_to_nuc(b'c'), 1);
802        assert_eq!(char_to_nuc(b'G'), 2);
803        assert_eq!(char_to_nuc(b'g'), 2);
804        assert_eq!(char_to_nuc(b'T'), 3);
805        assert_eq!(char_to_nuc(b't'), 3);
806        assert_eq!(char_to_nuc(b'U'), 3);
807        assert_eq!(char_to_nuc(b'u'), 3);
808    }
809
810    #[test]
811    fn test_char_to_nuc_invalid_bases() {
812        assert_eq!(char_to_nuc(b'N'), 4);
813        assert_eq!(char_to_nuc(b'n'), 4);
814        assert_eq!(char_to_nuc(b'X'), 4);
815        assert_eq!(char_to_nuc(b'-'), 4);
816        assert_eq!(char_to_nuc(b' '), 4);
817    }
818
819    #[test]
820    fn test_nucleotide_check_functions() {
821        let mut seq = vec![0u8; 10];
822
823        // Encode ATCG at positions 0,1,2,3
824        // A = 00 (default)
825        crate::bitmap::set_bit(&mut seq, 2); // T = 11
826        crate::bitmap::set_bit(&mut seq, 3);
827        crate::bitmap::set_bit(&mut seq, 5); // C = 01
828        crate::bitmap::set_bit(&mut seq, 6); // G = 10
829
830        assert!(is_a(&seq, 0));
831        assert!(!is_a(&seq, 1));
832        assert!(!is_a(&seq, 2));
833        assert!(!is_a(&seq, 3));
834
835        assert!(!is_t(&seq, 0));
836        assert!(is_t(&seq, 1));
837        assert!(!is_t(&seq, 2));
838        assert!(!is_t(&seq, 3));
839
840        assert!(!is_c(&seq, 0));
841        assert!(!is_c(&seq, 1));
842        assert!(is_c(&seq, 2));
843        assert!(!is_c(&seq, 3));
844
845        assert!(!is_g(&seq, 0));
846        assert!(!is_g(&seq, 1));
847        assert!(!is_g(&seq, 2));
848        assert!(is_g(&seq, 3));
849    }
850
851    #[test]
852    fn test_start_codon_functions() {
853        let mut seq = vec![0u8; 20];
854
855        // Encode ATG at position 0: A(00) T(11) G(10)
856        crate::bitmap::set_bit(&mut seq, 2); // T
857        crate::bitmap::set_bit(&mut seq, 3);
858        crate::bitmap::set_bit(&mut seq, 4); // G
859
860        assert!(is_atg(&seq, 0));
861        assert!(!is_gtg(&seq, 0));
862        assert!(!is_ttg(&seq, 0));
863
864        // Encode GTG at position 3: G(10) T(11) G(10)
865        crate::bitmap::set_bit(&mut seq, 6); // G
866        crate::bitmap::set_bit(&mut seq, 8); // T
867        crate::bitmap::set_bit(&mut seq, 9);
868        crate::bitmap::set_bit(&mut seq, 10); // G
869
870        assert!(!is_atg(&seq, 3));
871        assert!(is_gtg(&seq, 3));
872        assert!(!is_ttg(&seq, 3));
873    }
874
875    #[test]
876    fn test_stop_codon_functions() {
877        let mut seq = vec![0u8; 20];
878        let training = Training::default();
879
880        // Encode TAA at position 0: T(11) A(00) A(00)
881        crate::bitmap::set_bit(&mut seq, 0); // T
882        crate::bitmap::set_bit(&mut seq, 1);
883        // A and A are default (00)
884
885        assert!(is_stop(&seq, 0, &training));
886
887        // Encode TAG at position 3: T(11) A(00) G(10)
888        crate::bitmap::set_bit(&mut seq, 6); // T
889        crate::bitmap::set_bit(&mut seq, 7);
890        crate::bitmap::set_bit(&mut seq, 10); // G
891
892        assert!(is_stop(&seq, 3, &training));
893    }
894
895    #[test]
896    fn test_gc_content_calculation() {
897        let mut seq = vec![0u8; 20];
898
899        // Encode ATCG: A(00) T(11) C(01) G(10)
900        crate::bitmap::set_bit(&mut seq, 2); // T
901        crate::bitmap::set_bit(&mut seq, 3);
902        crate::bitmap::set_bit(&mut seq, 5); // C
903        crate::bitmap::set_bit(&mut seq, 6); // G
904
905        let gc = gc_content(&seq, 0, 3);
906        assert!((gc - 0.5).abs() < 0.001); // 2 GC out of 4 = 50%
907    }
908
909    #[test]
910    fn test_reverse_strand_reading_frame() {
911        // Test with sequence length 9 (frame_modulus = 0, since 9 % 3 == 0, so frame_modulus = 3)
912        assert_eq!(reverse_strand_reading_frame(0, 9), 2); // (3 - 1 - 0) % 3 = 2
913        assert_eq!(reverse_strand_reading_frame(1, 9), 1); // (3 - 1 - 1) % 3 = 1
914        assert_eq!(reverse_strand_reading_frame(2, 9), 0); // (3 - 1 - 2) % 3 = 0
915
916        // Test with sequence length 10 (frame_modulus = 1, since 10 % 3 == 1)
917        assert_eq!(reverse_strand_reading_frame(0, 10), 0); // (1 - 1 - 0) % 3 = 0
918
919        // Test with sequence length 8 (frame_modulus = 2, since 8 % 3 == 2)
920        assert_eq!(reverse_strand_reading_frame(0, 8), 1); // (2 - 1 - 0) % 3 = 1
921        assert_eq!(reverse_strand_reading_frame(1, 8), 0); // (2 - 1 - 1) % 3 = 0
922    }
923
924    #[test]
925    fn test_find_max_reading_frame() {
926        assert_eq!(find_max_reading_frame(10, 5, 3), 0);
927        assert_eq!(find_max_reading_frame(5, 10, 3), 1);
928        assert_eq!(find_max_reading_frame(5, 3, 10), 2);
929        assert_eq!(find_max_reading_frame(5, 5, 3), 1); // tie goes to second
930    }
931
932    #[test]
933    fn test_min_of_two_integers() {
934        assert_eq!(min_of_two_integers(5, 3), 3);
935        assert_eq!(min_of_two_integers(3, 5), 3);
936        assert_eq!(min_of_two_integers(-1, 5), -1);
937        assert_eq!(min_of_two_integers(5, 5), 5);
938    }
939
940    #[test]
941    fn test_calculate_kmer_index() {
942        let mut seq = vec![0u8; 20];
943
944        // Encode C at position 0 using the physical bitmap representation.
945        crate::bitmap::set_bit(&mut seq, 1); // C
946
947        let idx = calculate_kmer_index(2, &seq, 0);
948        assert_eq!(idx, 2);
949    }
950
951    fn calculate_kmer_index_reference(
952        kmer_length: usize,
953        encoded_sequence: &[u8],
954        position: usize,
955    ) -> usize {
956        let mut kmer_index = 0;
957        for i in 0..(2 * kmer_length) {
958            let bit_pos = position * 2 + i;
959            kmer_index |= usize::from(test_bit(encoded_sequence, bit_pos)) << i;
960        }
961        kmer_index
962    }
963
964    #[test]
965    fn test_calculate_kmer_index_matches_bit_reference() {
966        let sequence = b"ATCGNNGCATCGATGCGTACGATCGATCG";
967        let nucleotide_length = sequence.len();
968        let encoded_len = (nucleotide_length * 2).div_ceil(8);
969        let unknown_len = nucleotide_length.div_ceil(8);
970        let mut encoded = vec![0u8; encoded_len];
971        let mut unknown_sequence = vec![0u8; unknown_len];
972        let mut masks = Vec::new();
973
974        encode_sequence(
975            sequence,
976            &mut encoded,
977            &mut unknown_sequence,
978            &mut masks,
979            false,
980        )
981        .unwrap();
982
983        for kmer_length in [1, 2, 3, 6] {
984            for position in 0..=nucleotide_length - kmer_length {
985                assert_eq!(
986                    calculate_kmer_index(kmer_length, &encoded, position),
987                    calculate_kmer_index_reference(kmer_length, &encoded, position),
988                    "kmer_length={kmer_length}, position={position}",
989                );
990            }
991        }
992    }
993
994    #[test]
995    fn test_mer_text() {
996        assert_eq!(mer_text(0, 0), "None");
997        assert_eq!(mer_text(2, 0), "AA");
998        assert_eq!(mer_text(2, 1), "GA");
999        assert_eq!(mer_text(2, 2), "CA");
1000        assert_eq!(mer_text(2, 3), "TA");
1001    }
1002
1003    #[test]
1004    fn test_encode_sequence_basic() {
1005        let sequence = b"ATCG";
1006        let mut encoded = vec![0u8; 10];
1007        let mut unknown_sequence = vec![0u8; 10];
1008        let mut masks = Vec::new();
1009
1010        let gc = encode_sequence(
1011            sequence,
1012            &mut encoded,
1013            &mut unknown_sequence,
1014            &mut masks,
1015            false,
1016        )
1017        .unwrap();
1018        assert!((gc - 0.5).abs() < 0.001); // 2 GC out of 4 = 50%
1019    }
1020
1021    #[test]
1022    fn test_encode_sequence_with_n() {
1023        let sequence = b"ATNG";
1024        let mut encoded = vec![0u8; 10];
1025        let mut unknown_sequence = vec![0u8; 10];
1026        let mut masks = Vec::new();
1027
1028        let gc = encode_sequence(
1029            sequence,
1030            &mut encoded,
1031            &mut unknown_sequence,
1032            &mut masks,
1033            false,
1034        )
1035        .unwrap();
1036        assert!((gc - 0.25).abs() < 0.001); // 1 GC out of 4 = 25%
1037        assert!(crate::bitmap::test_bit(&unknown_sequence, 2)); // N should be marked in unknown_sequence
1038    }
1039
1040    #[test]
1041    fn test_simd_packed_encoding_matches_scalar_with_unknowns() {
1042        let sequence = b"ATCGNNNNGCATGCACTGACTNNATCGATCGXYZATCGATCGNNNNATCGATCGATCGATCG";
1043        let nucleotide_length = sequence.len();
1044        let encoded_len = (nucleotide_length * 2).div_ceil(8);
1045        let unknown_len = nucleotide_length.div_ceil(8);
1046
1047        let mut scalar_encoded = vec![0u8; encoded_len];
1048        let mut scalar_unknown = vec![0u8; unknown_len];
1049        let mut scalar_masks = Vec::new();
1050        let scalar_gc = encode_sequence(
1051            sequence,
1052            &mut scalar_encoded,
1053            &mut scalar_unknown,
1054            &mut scalar_masks,
1055            false,
1056        )
1057        .unwrap();
1058
1059        let mut simd_encoded = vec![0u8; encoded_len];
1060        let mut simd_unknown = vec![0u8; unknown_len];
1061        let simd_gc =
1062            encode_sequence_simd_wide_packed(sequence, &mut simd_encoded, &mut simd_unknown)
1063                .unwrap();
1064
1065        assert_eq!(simd_encoded, scalar_encoded);
1066        assert_eq!(simd_unknown, scalar_unknown);
1067        assert_eq!(scalar_masks.len(), 0);
1068        assert!((simd_gc - scalar_gc).abs() < f64::EPSILON);
1069    }
1070
1071    #[test]
1072    fn test_encode_sequence_masking() {
1073        // Create a sequence with 50+ N's to trigger masking (MASK_SIZE = 50)
1074        let mut sequence = b"ATC".to_vec();
1075        sequence.extend(vec![b'N'; 52]); // 52 N's should create a mask
1076        sequence.extend(b"GCG");
1077
1078        let mut encoded = vec![0u8; 60];
1079        let mut unknown_sequence = vec![0u8; 60];
1080        let mut masks = Vec::new();
1081
1082        let _gc = encode_sequence(
1083            &sequence,
1084            &mut encoded,
1085            &mut unknown_sequence,
1086            &mut masks,
1087            true,
1088        )
1089        .unwrap();
1090        assert!(!masks.is_empty()); // Should create at least one mask since we have 52 N's (> MASK_SIZE)
1091        assert_eq!(masks.len(), 1);
1092        assert_eq!(masks[0].begin, 3); // Start after "ATC"
1093        assert_eq!(masks[0].end, 54); // End at last N (3 + 52 - 1)
1094    }
1095
1096    #[test]
1097    fn test_is_gc() {
1098        let mut seq = vec![0u8; 10];
1099
1100        assert!(!is_gc(&seq, 0));
1101
1102        crate::bitmap::set_bit(&mut seq, 1);
1103        assert!(is_gc(&seq, 0));
1104
1105        let mut seq2 = vec![0u8; 10];
1106        crate::bitmap::set_bit(&mut seq2, 0);
1107        assert!(is_gc(&seq2, 0));
1108
1109        let mut seq3 = vec![0u8; 10];
1110        crate::bitmap::set_bit(&mut seq3, 0);
1111        crate::bitmap::set_bit(&mut seq3, 1);
1112        assert!(!is_gc(&seq3, 0));
1113    }
1114
1115    #[test]
1116    fn test_is_n() {
1117        let mut unknown_sequence = vec![0u8; 10];
1118
1119        assert!(!is_n(&unknown_sequence, 0));
1120        assert!(!is_n(&unknown_sequence, 100));
1121
1122        crate::bitmap::set_bit(&mut unknown_sequence, 5);
1123        assert!(is_n(&unknown_sequence, 5));
1124    }
1125
1126    #[test]
1127    fn test_calculate_background_mer_frequencies() {
1128        let seq = vec![0u8; 20]; // All A's
1129        let rseq = vec![0u8; 20]; // All A's
1130        let mut bg = vec![0.0; 16]; // 4^2 = 16 possible 2-mers
1131
1132        calculate_background_mer_frequencies(2, &seq, &rseq, 10, &mut bg);
1133
1134        // Should have high frequency for AA (index 0) and low for others
1135        assert!(bg[0] > 0.5); // AA should be common
1136    }
1137
1138    #[test]
1139    fn test_rcom_seq() {
1140        let seq = vec![0u8; 10];
1141        let unknown_sequence = vec![0u8; 10];
1142
1143        // Encode A at position 0
1144        // A = 00, complement = T = 11
1145
1146        let rseq = create_reverse_complement_sequence(&seq, &unknown_sequence, 2);
1147
1148        assert!(is_t(&rseq, 1));
1149    }
1150
1151    fn create_reverse_complement_sequence_reference(
1152        forward_sequence: &[u8],
1153        unknown_sequence: &[u8],
1154        nucleotide_length: usize,
1155    ) -> Vec<u8> {
1156        let mut reverse_complement_encoded_sequence = vec![0; forward_sequence.len()];
1157        let sequence_length = nucleotide_length * 2;
1158
1159        for i in 0..sequence_length {
1160            if !crate::bitmap::test_bit(forward_sequence, i) {
1161                let target_pos = if i % 2 == 0 {
1162                    sequence_length - i - 2
1163                } else {
1164                    sequence_length - i
1165                };
1166                if target_pos < sequence_length {
1167                    crate::bitmap::set_bit(&mut reverse_complement_encoded_sequence, target_pos);
1168                }
1169            }
1170        }
1171
1172        for i in 0..nucleotide_length {
1173            if crate::bitmap::test_bit(unknown_sequence, i) && sequence_length >= 2 + i * 2 {
1174                crate::bitmap::toggle_bit(
1175                    &mut reverse_complement_encoded_sequence,
1176                    sequence_length - 1 - i * 2,
1177                );
1178                crate::bitmap::toggle_bit(
1179                    &mut reverse_complement_encoded_sequence,
1180                    sequence_length - 2 - i * 2,
1181                );
1182            }
1183        }
1184
1185        reverse_complement_encoded_sequence
1186    }
1187
1188    #[test]
1189    fn test_reverse_complement_matches_reference_with_unknowns() {
1190        for sequence in [
1191            b"ATCGNNGCATCG".as_slice(),
1192            b"ATCGNNGCATCGA".as_slice(),
1193            b"NNNNATCGXYZATCG".as_slice(),
1194        ] {
1195            let nucleotide_length = sequence.len();
1196            let encoded_len = (nucleotide_length * 2).div_ceil(8);
1197            let unknown_len = nucleotide_length.div_ceil(8);
1198
1199            let mut encoded = vec![0u8; encoded_len];
1200            let mut unknown_sequence = vec![0u8; unknown_len];
1201            let mut masks = Vec::new();
1202            encode_sequence(
1203                sequence,
1204                &mut encoded,
1205                &mut unknown_sequence,
1206                &mut masks,
1207                false,
1208            )
1209            .unwrap();
1210
1211            let expected = create_reverse_complement_sequence_reference(
1212                &encoded,
1213                &unknown_sequence,
1214                nucleotide_length,
1215            );
1216            let actual =
1217                create_reverse_complement_sequence(&encoded, &unknown_sequence, nucleotide_length);
1218
1219            assert_eq!(actual, expected);
1220        }
1221    }
1222
1223    #[test]
1224    fn test_translation_table_functions() {
1225        assert!(uses_only_atg(6));
1226        assert!(uses_only_atg(10));
1227        assert!(!uses_only_atg(11));
1228
1229        assert!(gtg_not_start(1));
1230        assert!(gtg_not_start(22));
1231        assert!(!gtg_not_start(11));
1232
1233        assert!(ttg_not_start(1));
1234        assert!(ttg_not_start(9));
1235        assert!(!ttg_not_start(11));
1236    }
1237
1238    #[test]
1239    fn test_start_codon_with_training() {
1240        let mut training = Training {
1241            translation_table: 11,
1242            ..Training::default()
1243        };
1244
1245        let mut seq = vec![0u8; 20];
1246
1247        // Encode ATG at position 0
1248        crate::bitmap::set_bit(&mut seq, 2); // T
1249        crate::bitmap::set_bit(&mut seq, 3);
1250        crate::bitmap::set_bit(&mut seq, 4); // G
1251
1252        assert!(is_start(&seq, 0, &training));
1253
1254        // Test with table that only uses ATG
1255        training.translation_table = 6;
1256        assert!(is_start(&seq, 0, &training)); // ATG still works
1257
1258        // Encode GTG and test
1259        let mut seq2 = vec![0u8; 20];
1260        crate::bitmap::set_bit(&mut seq2, 0); // G
1261        crate::bitmap::set_bit(&mut seq2, 2); // T
1262        crate::bitmap::set_bit(&mut seq2, 3);
1263        crate::bitmap::set_bit(&mut seq2, 4); // G
1264
1265        assert!(!is_start(&seq2, 0, &training)); // GTG not allowed in table 6
1266    }
1267
1268    #[test]
1269    fn test_stop_codon_special_tables() {
1270        let mut training = Training::default();
1271        let mut seq = vec![0u8; 20];
1272
1273        // Test AGA stop in table 2
1274        training.translation_table = 2;
1275        // Encode AGA: A(00) G(10) A(00)
1276        crate::bitmap::set_bit(&mut seq, 2); // G
1277
1278        assert!(is_stop(&seq, 0, &training));
1279
1280        // Test TCA stop in table 22
1281        training.translation_table = 22;
1282        let mut seq2 = vec![0u8; 20];
1283        // Encode TCA: T(11) C(01) A(00)
1284        crate::bitmap::set_bit(&mut seq2, 0); // T
1285        crate::bitmap::set_bit(&mut seq2, 1);
1286        crate::bitmap::set_bit(&mut seq2, 3); // C
1287
1288        assert!(is_stop(&seq2, 0, &training));
1289    }
1290}