msbwt2 0.3.2

extern crate log;

use log::info;
use needletail::{parse_fastx_file,Sequence};
use std::convert::TryInto;
use std::io::prelude::*;
use std::fs;

use crate::msbwt_core::*;

use crate::run_block_av_flat::VC_LEN;
use crate::rle_bplus_tree::RLEBPlusTree;
use crate::string_util::convert_stoi;

/// The inital k-mer size used for short circuiting
const INITIAL_QUERY: usize = 10;
/// A factor that influences how fast the query size will change
const COST_FACTOR: f64 = 0.000001;

/// A BWT type built on top of a B+ tree that allows for new strings to be added via a function.
/// This is most useful for directly building a BWT from scratch or for adding to an existing BWT with additional logic.
/// It attempts to short circuit the calculation of the insertion point for sorted strings, but will fall back to a full query if that fails.
pub struct DynamicBWT {
    /// The actual B+ tree structure
    tree_bwt: RLEBPlusTree,
    /// The total number of each symbol captured by the BWT
    symbol_counts: [u64; VC_LEN],
    /// The start index of each symbol ($ is always 0)
    start_index: [u64; VC_LEN],
    /// The total number of strings in the BWT
    string_count: u64,
    /// The total number of symbols in the BWT
    total_count: u64,
    /// This determines how many bases to query when inserting sorted strings to attempt to short-circuit.
    /// This number is dynamically adjusted based on the success/failure of short-circuiting strings while inserting.
    sort_query_len: f64,
    /// A simple helper array for storing the number of [successful short circuits, impossible short circuits (i.e. identical strings), and failed short circuits].
    /// This is a semi-frequent output while inserting strings.
    short_circuits: [usize; 3],
}

impl Default for DynamicBWT {
    /// This will create an empty BWT that's ready to have strings added directly to it.
    fn default() -> Self {
        DynamicBWT {
            string_count: 0,
            total_count: 0,
            tree_bwt: Default::default(),
            symbol_counts: [0; VC_LEN],
            start_index: [0; VC_LEN],
            sort_query_len: INITIAL_QUERY as f64,
            short_circuits: [0; 3]
        }
    }
}

impl BWT for DynamicBWT {
    /// Initializes the BWT from a compressed BWT vector.
    /// # Arguments
    /// * `bwt` - the run-length encoded BWT stored in a Vec<u8> 
    /// # Examples
    /// ```rust
    /// use msbwt2::msbwt_core::BWT;
    /// use msbwt2::dynamic_bwt::DynamicBWT;
    /// use msbwt2::bwt_converter::convert_to_vec;
    /// //strings "ACGT" and "CCGG"
    /// let seq = "TG$$CAGCCG";
    /// let vec = convert_to_vec(seq.as_bytes());
    /// let mut bwt = DynamicBWT::new();
    /// bwt.load_vector(vec);
    /// ```
    fn load_vector(&mut self, bwt: Vec<u8>) {
        info!("Initializing BWT with {:?} compressed values...", bwt.len());

        //reset all of these 
        self.tree_bwt = Default::default();
        self.symbol_counts = [0; VC_LEN];
        self.sort_query_len = INITIAL_QUERY as f64;
        self.short_circuits = [0; 3];

        //general strategy here is to build up the count of a character and then insert them
        let mut prev_char: u8 = 255;
        let mut current_char: u8;
        let mut power_multiple: u64 = 1;
        let mut current_count: u64;

        //go through each compressed block in the RLE encoded vector to calculate total character counts
        let mut current_index: u64 = 0;
        for value in bwt {
            current_char = value & MASK;
            if current_char == prev_char {
                power_multiple *= NUM_POWER as u64;
            }
            else {
                power_multiple = 1;
            }
            prev_char = current_char;
            current_count = (value >> LETTER_BITS) as u64 * power_multiple;

            for _ in 0..current_count {
                let _dummy = self.tree_bwt.insert_and_count(current_index, current_char);
                current_index += 1;
            }

            //add this to the total counts
            self.symbol_counts[current_char as usize] += current_count;
        }

        self.string_count = self.symbol_counts[0];
        self.total_count = self.symbol_counts.iter().sum();
        self.start_index = self.symbol_counts.iter().scan(0, |sum, &count| {
            *sum += count;
            Some(*sum - count)
        }).collect::<Vec<u64>>().try_into().unwrap();
        info!("Loaded BWT with symbol counts: {:?}", self.symbol_counts);
        info!("Finished BWT initialization.")
    }

    /// Initializes the BWT from the numpy file format for compressed BWTs
    /// # Arguments
    /// * `filename` - the name of the file to load into memory
    /// # Examples
    /// ```rust
    /// use msbwt2::msbwt_core::BWT;
    /// use msbwt2::dynamic_bwt::DynamicBWT;
    /// use msbwt2::string_util;
    /// let mut bwt = DynamicBWT::new();
    /// let filename: String = "test_data/two_string.npy".to_string();
    /// bwt.load_numpy_file(&filename);
    /// assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"ACGT")), 1);
    /// ```
    fn load_numpy_file(&mut self, filename: &str) -> std::io::Result<()> {
        //read the numpy header: http://docs.scipy.org/doc/numpy-1.10.1/neps/npy-format.html
        //get the initial file size
        let file_metadata: fs::Metadata = fs::metadata(filename)?;
        let full_file_size: u64 = file_metadata.len();

        //read the initial fixed header
        let mut file = fs::File::open(filename)?;
        let mut init_header: Vec<u8> = vec![0; 10];
        let read_count: usize = file.read(&mut init_header[..])?;
        if read_count != 10 {
            panic!("Could not read initial 10 bytes of header for file {:?}", filename);
        }

        //read the dynamic header
        let header_len: usize = init_header[8] as usize + 256 * init_header[9] as usize;
        let mut skip_bytes: usize = 10+header_len;
        if skip_bytes % 16 != 0 {
            skip_bytes = ((skip_bytes / 16)+1)*16;
        }
        let mut skip_header: Vec<u8> = vec![0; skip_bytes-10];
        match file.read_exact(&mut skip_header[..]) {
            Ok(()) => {},
            Err(e) => {
                return Err(
                    std::io::Error::new(
                        e.kind(),
                        format!("Could not read bytes 10-{skip_bytes:?} of header for file {filename:?}, root-error {e:?}")
                    )
                );
            }
        }
        
        //parse the header string for the expected length, requires a lot of manipulation of the string because of numpy header styling
        let header_string = String::from_utf8(skip_header).unwrap()
            .replace('\'', "\"")
            .replace("False", "false")
            .replace('(', "[")
            .replace(')', "]")
            .replace(", }", "}")
            .replace(", ]", "]")
            .replace(",]", "]");
        let header_dict: serde_json::Value = serde_json::from_str(&header_string)
            .unwrap_or_else(|_| panic!("Error while parsing header string: {:?}", header_string));
        let expected_length: u64 = header_dict["shape"][0].as_u64().unwrap();
        
        //check that the disk size matches our expectation
        let bwt_disk_size: u64 = full_file_size - skip_bytes as u64;
        if expected_length != bwt_disk_size {
            return Err(
                std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    format!("Header indicates shape of {expected_length:?}, but remaining file size is {bwt_disk_size:?}")
                )
            );
        }

        //finally read in everything else
        info!("Loading BWT with {:?} compressed values from disk...", bwt_disk_size);
        let mut bwt_data: Vec<u8> = Vec::<u8>::with_capacity(bwt_disk_size as usize);
        let read_count: usize = file.read_to_end(&mut bwt_data)?;
        if read_count as u64 != bwt_disk_size {
            return Err(
                std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    format!("Only read {read_count:?} of {bwt_disk_size:?} bytes of BWT body for file {filename:?}")
                )
            );
        }
        
        //we loaded the file into memory, now just do the load from vec
        self.load_vector(bwt_data);

        Ok(())
    }

    /// Returns the total number of occurences of a given symbol
    /// # Arguments
    /// * `symbol` - the symbol in integer form
    /// # Examples
    /// ```rust
    /// # use msbwt2::msbwt_core::BWT;
    /// # use msbwt2::dynamic_bwt::DynamicBWT;
    /// # use msbwt2::bwt_converter::convert_to_vec;
    /// # let seq = "TG$$CAGCCG";
    /// # let vec = convert_to_vec(seq.as_bytes());
    /// # let mut bwt = DynamicBWT::new();
    /// # bwt.load_vector(vec);
    /// let string_count = bwt.get_symbol_count(0);
    /// assert_eq!(string_count, 2);
    /// ```
    #[inline]
    fn get_symbol_count(&self, symbol: u8) -> u64 {
        self.symbol_counts[symbol as usize]
    }

    /// This will return the total number of symbols contained by the BWT
    /// # Examples
    /// ```rust
    /// # use msbwt2::msbwt_core::BWT;
    /// # use msbwt2::dynamic_bwt::DynamicBWT;
    /// # use msbwt2::bwt_converter::convert_to_vec;
    /// # let seq = "TG$$CAGCCG";
    /// # let vec = convert_to_vec(seq.as_bytes());
    /// # let mut bwt = DynamicBWT::new();
    /// # bwt.load_vector(vec);
    /// let total_size = bwt.get_total_size();
    /// assert_eq!(total_size, 10);
    /// ```
    #[inline]
    fn get_total_size(&self) -> u64 {
        self.total_count
    }

    /// Performs a range constraint on a BWT range. This implicitly represents prepending a character `sym` to a k-mer
    /// represented by `input_range` to create a new range representing a (k+1)-mer.
    /// # Arguments
    /// * `sym` - the symbol to pre-pend in integer form
    /// * `input_range` - the range to pre-pend to
    /// # Safety
    /// This function is unsafe because there are no guarantees that the symbol or bounds will be checked by the implementing structure.
    unsafe fn constrain_range(&self, sym: u8, input_range: &BWTRange) -> BWTRange {
        BWTRange {
            l: (self.start_index[sym as usize]+self.tree_bwt.count(input_range.l, sym)),
            h: (self.start_index[sym as usize]+self.tree_bwt.count(input_range.h, sym))
        }
    }
}

impl DynamicBWT {
    /// Allocation function for the BWT, look at `load_vector(...)` for initialization.
    /// # Examples
    /// ```rust
    /// use msbwt2::dynamic_bwt::DynamicBWT;
    /// let mut bwt = DynamicBWT::new();
    /// ```
    pub fn new() -> Self {
        Default::default()
    }

    /// This will return the array of symbol counts contained within the BWT
    #[inline]
    pub fn get_symbol_counts(&self) -> [u64; VC_LEN] {
        self.symbol_counts
    }

    /// This will return the current height of the B+ tree storing the data
    #[inline]
    pub fn get_height(&self) -> usize {
        self.tree_bwt.get_height()
    }

    /// This will return the total number of data nodes in the B+ tree
    #[inline]
    pub fn get_node_count(&self) -> usize {
        self.tree_bwt.get_node_count()
    }
    
    /// This is the main function for adding strings to the DynamicBWT.
    /// # Arguments
    /// * val - the string to add
    /// * sorted - if true, this will add the string to it's sorted position, otherwise the end
    /// # Examples
    /// ```rust
    /// use msbwt2::dynamic_bwt::DynamicBWT;
    /// let data: String = "ACGNT".to_string();
    /// let bwt: Vec<u8> = vec![5, 0, 1, 2, 3, 4];
    /// let mut ubwt: DynamicBWT = Default::default();
    /// ubwt.insert_string(&data, false);
    /// assert_eq!(ubwt.to_vec(), bwt);
    /// ```
    #[inline]
    pub fn insert_string(&mut self, val: &str, sorted: bool) {
        let int_form: Vec<u8> = convert_stoi(val);

        //initial position is the total number of string
        let mut next_insert: u64;
        
        if sorted {
            let mut start_index: u64 = 0;
            next_insert = self.total_count;

            //attempt a short circuit
            let query_len = std::cmp::min(self.sort_query_len as usize, int_form.len());
            for pred_symbol in int_form[..query_len].iter().rev() {
                start_index = self.tree_bwt.count(start_index, *pred_symbol)+self.start_index[*pred_symbol as usize];
                next_insert = self.tree_bwt.count(next_insert, *pred_symbol)+self.start_index[*pred_symbol as usize];
            }
            start_index = self.tree_bwt.count(start_index, 0);
            next_insert = self.tree_bwt.count(next_insert, 0);

            if start_index != next_insert {
                let original_ni: u64 = next_insert;

                //short circuit failed
                for pred_symbol in int_form.iter().rev() {
                    next_insert = self.tree_bwt.count(next_insert, *pred_symbol)+self.start_index[*pred_symbol as usize];
                }
                next_insert = self.tree_bwt.count(next_insert, 0);
                
                if original_ni == next_insert {
                    //the full search did nothing this is a copy sequence
                    //cutting down the search will save on the initial k-mer query
                    self.sort_query_len -= 2.0 * COST_FACTOR * query_len as f64;
                    self.short_circuits[1] += 1;
                } else {
                    //the full search did make it more specific, so we should increase the query short circuit size
                    //this could save at most the length of the full sequence
                    self.sort_query_len += COST_FACTOR * int_form.len() as f64;
                    self.short_circuits[2] += 1;
                }
            } else {
                //short circuit success, making it smaller will save 2 bases queries, and we also have downweighted this
                self.sort_query_len -= 2.0 * COST_FACTOR;
                self.short_circuits[0] += 1;
            }
        } else {
            next_insert = self.string_count;
        }
        
        //go through the characters in reverse
        let mut symbol: u8 = 0; // $
        for pred_symbol in int_form.iter().rev() {
            next_insert = self.tree_bwt.insert_and_count(next_insert, *pred_symbol);
            self.symbol_counts[*pred_symbol as usize] += 1;
            for i in (symbol+1) as usize..VC_LEN {
                self.start_index[i] += 1;
            }

            //after any adjustments add in the new offset for the symbol we're currently at
            next_insert += self.start_index[*pred_symbol as usize];
            symbol = *pred_symbol;
        }

        //one final insert for the $
        self.tree_bwt.insert_and_count(next_insert, 0);
        self.symbol_counts[0] += 1;
        for i in (symbol+1) as usize..VC_LEN {
            self.start_index[i] += 1;
        }

        self.total_count += (int_form.len()+1) as u64;
        self.string_count += 1;

        if self.string_count % 10000 == 0 {
            info!("Strings: {}\tShort-k: {:.2}\t[pass, dup, fail]: {:?}\tHeight, nodes: {} {}", self.string_count, self.sort_query_len, self.short_circuits, self.get_height(), self.get_node_count());
            self.short_circuits = [0; 3];
        }
    }

    /// This will return the data in a plain Vector format, with one symbol per index.
    /// # Examples
    /// ```rust
    /// use msbwt2::dynamic_bwt::DynamicBWT;
    /// let data: String = "ACGNT".to_string();
    /// let bwt: Vec<u8> = vec![5, 0, 1, 2, 3, 4];
    /// let mut ubwt: DynamicBWT = Default::default();
    /// ubwt.insert_string(&data, false);
    /// assert_eq!(ubwt.to_vec(), bwt);
    /// ```
    pub fn to_vec(&self) -> Vec<u8> {
        self.tree_bwt.to_vec()
    }

    /// This will return an iterator over the symbols in the BWT in their number format (e.g. 0-5).
    /// # Examples
    /// ```rust
    /// use msbwt2::dynamic_bwt::{create_from_fastx,DynamicBWT};
    /// use msbwt2::msbwt_core::BWT;
    /// use msbwt2::string_util;
    /// let npy_result: String = "test_data/two_string.npy".to_string();
    /// let mut truth_bwt: DynamicBWT = Default::default();
    /// truth_bwt.load_numpy_file(&npy_result);
    /// let single_file = vec!["./test_data/two_string.fa"];
    /// let bwt: DynamicBWT = create_from_fastx(&single_file, true).unwrap();
    /// assert_eq!(truth_bwt.to_vec(), bwt.iter().collect::<Vec<u8>>());
    /// for sym in bwt.iter() {
    ///     print!("{}", sym);
    /// }
    /// ```
    pub fn iter(&self) -> impl Iterator<Item = u8> + '_ {
        self.tree_bwt.into_iter()
    }

    /// This will return an iterator over the runs in the BWT in format (symbol, length).
    /// # Examples
    /// ```rust
    /// use msbwt2::dynamic_bwt::DynamicBWT;
    /// let mut bwt: DynamicBWT = Default::default();
    /// bwt.insert_string("ACCC", true);
    /// let runs: Vec<(u8, u64)> = bwt.run_iter().collect::<Vec<(u8, u64)>>();
    /// //C$CCA
    /// let expected_runs: Vec<(u8, u64)> = vec![(2, 1), (0, 1), (2, 2), (1, 1)];
    /// assert_eq!(expected_runs, runs);
    /// ```
    pub fn run_iter(&self) -> impl Iterator<Item = (u8, u64)> + '_ {
        self.tree_bwt.run_iter()
    }
}


/// This will create a BWT from a list of FASTX files and return that BWT upon completion.
/// # Arguments 
/// * `filenames` - a list of filenames to be parsed, fastq, fasta, and gzipped versions of each are supported
/// * `sorted` - if True, strings will be added in lexicographic (e.g. sorted) order, otherwise they will be inserted chronologically (e.g. the order in the file)
/// # Examples
/// ```rust
/// use msbwt2::dynamic_bwt::{create_from_fastx,DynamicBWT};
/// use msbwt2::msbwt_core::BWT;
/// use msbwt2::string_util;
/// let npy_result: String = "test_data/two_string.npy".to_string();
/// let mut truth_bwt: DynamicBWT = Default::default();
/// truth_bwt.load_numpy_file(&npy_result);
/// let single_file = vec!["./test_data/two_string.fa"];
/// let bwt: DynamicBWT = create_from_fastx(&single_file, true).unwrap();
/// assert_eq!(truth_bwt.to_vec(), bwt.to_vec());
/// assert_eq!(truth_bwt.count_kmer(&string_util::convert_stoi(&"$")), 2);
/// assert_eq!(truth_bwt.count_kmer(&string_util::convert_stoi(&"ACGT")), 1);
/// assert_eq!(truth_bwt.count_kmer(&string_util::convert_stoi(&"TGCA")), 1);
/// ```
pub fn create_from_fastx<T: std::convert::AsRef<std::path::Path> + std::fmt::Display>(filenames: &[T], sorted: bool) -> Result<DynamicBWT, Box<dyn std::error::Error>> {
    let mut bwt: DynamicBWT = Default::default();
    info!("Creating BWT from FASTX files...");
    for filename in filenames {
        let mut reader = parse_fastx_file(filename)?;

        //go through all the records
        let initial_string_count = bwt.get_symbol_count(0);
        info!("Loading file \"{}\"...", filename);
        while let Some(record) = reader.next() {
            //all we care about is the sequence length
            let seq_rec = record?;
            let norm_seq = seq_rec.normalize(false);
            bwt.insert_string(std::str::from_utf8(norm_seq.as_ref()).unwrap(), sorted);
        }
        let count = bwt.get_symbol_count(0) - initial_string_count;
        info!("Finished loading file with {} sequences.", count);
    }
    info!("Finished creating BWT, symbol counts: {:?}", bwt.get_symbol_counts());
    Ok(bwt)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bwt_converter::*;
    use crate::bwt_util::naive_bwt;
    use crate::string_util;
    use tempfile::{Builder, NamedTempFile};

    #[test]
    fn test_init() {
        let ubwt: DynamicBWT = Default::default();
        assert_eq!(ubwt.to_vec(), Vec::<u8>::new());
    }

    #[test]
    fn test_single_string() {
        let data: String = "ACGNT".to_string();
        let bwt: Vec<u8> = vec![5, 0, 1, 2, 3, 4];
        let mut ubwt: DynamicBWT = Default::default();
        ubwt.insert_string(&data, false);
        assert_eq!(ubwt.to_vec(), bwt);
    }

    #[test]
    fn test_multi_string_unsorted() {
        let mut data: Vec<&str> = vec!["CCGT", "ACG", "N"];
        let bwt = convert_stoi(&"GTN$$ACCC$G");
        
        //to get an identical result to sorted rope, we have to sort
        data.sort();
        
        //insert the strings in the now sorted order
        let mut ubwt: DynamicBWT = Default::default();
        for s in data.iter() {
            ubwt.insert_string(s, false);
        }
        assert_eq!(ubwt.to_vec(), bwt);
    }

    #[test]
    fn test_multi_string_sorted() {
        let data: Vec<&str> = vec!["ACG", "N", "CCGT", "N", "ACG", "ACG", "CCGT", "N"];
        let bwt = string_util::convert_stoi(&naive_bwt(&data));
        
        //insert the strings in the now sorted order
        let mut ubwt: DynamicBWT = Default::default();
        for s in data.iter() {
            ubwt.insert_string(s, true);
        }
        assert_eq!(ubwt.to_vec(), bwt);
    }

    #[test]
    fn test_multi_length() {
        //getting bigger in order
        let data: Vec<&str> = vec!["A", "AA", "AAA", "AAAA", "AAAAA"];
        let bwt = string_util::convert_stoi(&naive_bwt(&data));
        
        let mut ubwt: DynamicBWT = Default::default();
        for s in data.iter() {
            ubwt.insert_string(s, true);
        }
        assert_eq!(ubwt.to_vec(), bwt);

        //getting smaller in order
        let data: Vec<&str> = vec!["AAAAA", "AAAA", "AAA", "AA", "A"];
        let bwt = string_util::convert_stoi(&naive_bwt(&data));
        
        let mut ubwt: DynamicBWT = Default::default();
        for s in data.iter() {
            ubwt.insert_string(s, true);
        }
        assert_eq!(ubwt.to_vec(), bwt);
    }

    #[test]
    fn test_sampled_bwt() {
        let genome: String = "ACCGTGTTGCCGTAGTGAAAAGTGACGACGTGAGATGGCCAAAGTGGGTCTCTGTG".to_string();
        let read_length: usize = 20;
        let coverage: usize = 32;//make sure we get some runs
        //let read_length: usize = 5;
        //let coverage: usize = 1;
        let mut data: Vec<&str> = vec![];
        for s in 0..genome.len()-read_length {//-43 {
            for _ in 0..coverage {
                data.push(&genome[s..s+read_length]);
            }
        }

        //let data: Vec<&str> = vec!["ACCGT", "TGCCGC"];
        //let data: Vec<&str> = vec!["ACCGT", "CCGTG"];
        //println!("data: {:?}", data);

        //use this function to make a bwt the naive (i.e. slow) way
        let naive = string_util::convert_stoi(&naive_bwt(&data));
        
        //insert the strings in the now sorted order
        let mut ubwt: DynamicBWT = Default::default();
        for s in data.iter() {
            ubwt.insert_string(s, true);
        }
        assert_eq!(ubwt.to_vec(), naive);
    }

    #[test]
    fn test_load_dynamicbwt_from_vec() {
        //strings - "CCGT\nACG\nN"
        //build the BWT
        let data: Vec<&str> = vec!["CCGT", "N", "ACG"];
        
        //stream and compress the BWT
        //let bwt_stream = stream_bwt_from_fastqs(&fastq_filenames).unwrap();
        let bwt_stream = naive_bwt(&data);
        let compressed_bwt = convert_to_vec(bwt_stream.as_bytes());
        
        let mut bwt = DynamicBWT::new();
        bwt.load_vector(compressed_bwt);

        let expected_totals = vec![3, 1, 3, 2, 1, 1];
        for i in 0..6 {
            //make sure the total counts are correct
            assert_eq!(bwt.get_symbol_count(i as u8), expected_totals[i]);
        }
    }

    #[test]
    fn test_load_dynamicbwt_from_npy() {
        //strings - "CCGT\nACG\nN"
        //build the BWT
        let data: Vec<&str> = vec!["CCGT", "N", "ACG"];
        
        //stream and compress the BWT
        //let bwt_stream = stream_bwt_from_fastqs(&fastq_filenames).unwrap();
        let bwt_stream = naive_bwt(&data);
        let compressed_bwt = convert_to_vec(bwt_stream.as_bytes());
        
        //save the output to a temporary numpy file
        let bwt_file: NamedTempFile = Builder::new().prefix("temp_data_").suffix(".npy").tempfile().unwrap();
        let filename: String = bwt_file.path().to_str().unwrap().to_string();
        save_bwt_numpy(&compressed_bwt[..], &filename).unwrap();
        
        //load it back in and verify counts
        let mut bwt = DynamicBWT::new();
        bwt.load_numpy_file(&filename).unwrap();

        let expected_totals = vec![3, 1, 3, 2, 1, 1];
        for i in 0..6 {
            //make sure the total counts are correct
            assert_eq!(bwt.get_symbol_count(i as u8), expected_totals[i]);
        }
    }

    #[test]
    fn test_constrain_range() {
        //strings - "CCGT\nACG\nN"
        //build the BWT
        let data: Vec<&str> = vec!["CCGT", "N", "ACG"];
        
        //stream and compress the BWT
        let bwt_stream = naive_bwt(&data);
        assert_eq!(bwt_stream, "GTN$$ACCC$G");
        let bwt_int_form = string_util::convert_stoi(&bwt_stream);
        let compressed_bwt = convert_to_vec(bwt_stream.as_bytes());
        //[G, T, N, 2$, A, 3C, $, G]
        assert_eq!(compressed_bwt.len(), 8);
        
        //load it back in and verify counts
        let mut bwt: DynamicBWT = Default::default();
        bwt.load_vector(compressed_bwt.clone());

        let initial_range = BWTRange {
            l: 0,
            h: bwt_stream.len() as u64
        };
        
        //this is verifying that all single-symbol queries get the start/end range
        for sym in 0..VC_LEN {
            let new_range = unsafe {
                bwt.constrain_range(sym as u8, &initial_range)
            };
            assert_eq!(new_range, BWTRange{
                l: bwt.start_index[sym] as u64, 
                h: (bwt.start_index[sym]+bwt.symbol_counts[sym]) as u64
            });
        }

        //now lets verify that we get all ascending symbols
        for sym in 0..VC_LEN {
            let mut sym_count = 0;
            for ind in 0..(bwt_stream.len()+1) {
                //test from 0 to the current index
                let initial_range = BWTRange {
                    l: 0,
                    h: ind as u64
                };

                let new_range = unsafe {
                    bwt.constrain_range(sym as u8, &initial_range)
                };
                assert_eq!(new_range, BWTRange {
                    l: bwt.start_index[sym] as u64,
                    h: (bwt.start_index[sym]+sym_count) as u64
                });

                //test from the current index to the high point
                let initial_range = BWTRange {
                    l: ind as u64,
                    h: bwt_stream.len() as u64
                };

                let new_range = unsafe {
                    bwt.constrain_range(sym as u8, &initial_range)
                };
                assert_eq!(new_range, BWTRange {
                    l: (bwt.start_index[sym]+sym_count) as u64,
                    h: (bwt.start_index[sym]+bwt.symbol_counts[sym]) as u64
                });

                //check if we need to adjust our expected values at all
                if ind < bwt_stream.len() && bwt_int_form[ind] == sym as u8 {
                    sym_count += 1;
                }
            }
        }
    }

    #[test]
    fn test_count_kmer() {
        //strings - "CCGT\nACG\nN"
        //build the BWT
        let data: Vec<&str> = vec!["CCGTACGTA", "GGTACAGTA", "ACGACGACG"];
        
        //stream and compress the BWT
        let bwt_stream = naive_bwt(&data);
        let compressed_bwt = convert_to_vec(bwt_stream.as_bytes());
        
        //load it back in and verify counts
        let mut bwt: DynamicBWT = Default::default();
        bwt.load_vector(compressed_bwt.clone());

        //simple sanity checks, make sure our single-character symbols matches the total count
        for c in 0..VC_LEN as u8 {
            let test_seq: Vec<u8> = vec![c];
            assert_eq!(bwt.get_symbol_count(c), bwt.count_kmer(&test_seq));
        }
        
        //check that each string shows up once
        for seq in data.iter() {
            let test_seq = string_util::convert_stoi(seq);
            assert_eq!(bwt.count_kmer(&test_seq), 1);
        }

        //now lets check some semi-arbitrary substrings
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"ACG")), 4);
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"CC")), 1);
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"TAC")), 2);
    }

    #[test]
    fn test_load_and_add() {
        //strings - "CCGT\nACG\nN"
        //build the BWT
        let mut data: Vec<&str> = vec!["CCGTACGTA", "GGTACAGTA", "ACGACGACG"];
        
        //stream and compress the BWT
        //let bwt_stream = stream_bwt_from_fastqs(&fastq_filenames).unwrap();
        let bwt_stream = naive_bwt(&data);
        let compressed_bwt = convert_to_vec(bwt_stream.as_bytes());
        
        //load it back in and verify counts
        let mut bwt: DynamicBWT = Default::default();
        bwt.load_vector(compressed_bwt.clone());

        // now lets add a new string
        let new_string = "AAGTCATAT";
        bwt.insert_string(&new_string, true);
        data.push(new_string);

        //simple sanity checks, make sure our single-character symbols matches the total count
        for c in 0..VC_LEN as u8 {
            let test_seq: Vec<u8> = vec![c];
            assert_eq!(bwt.get_symbol_count(c), bwt.count_kmer(&test_seq));
        }
        
        //check that each string shows up once
        for seq in data.iter() {
            let test_seq = string_util::convert_stoi(seq);
            assert_eq!(bwt.count_kmer(&test_seq), 1);
        }

        //now lets check some semi-arbitrary substrings
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"ACG")), 4);
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"CC")), 1);
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"TAC")), 2);

        //these should have changed with the new string
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"AA")), 1);
        assert_eq!(bwt.count_kmer(&string_util::convert_stoi(&"GT")), 5);
    }

    #[test]
    fn test_create_from_fastx() {
        //empty test
        let empty_file_list: Vec<String> = Default::default();
        let bwt: DynamicBWT = create_from_fastx(&empty_file_list, true).unwrap();
        assert_eq!(bwt.to_vec(), Vec::<u8>::new());

        //two string test
        let npy_result: String = "test_data/two_string.npy".to_string();
        let mut truth_bwt: DynamicBWT = Default::default();
        truth_bwt.load_numpy_file(&npy_result).unwrap();
        
        let single_file = vec!["./test_data/two_string.fa"];
        let bwt: DynamicBWT = create_from_fastx(&single_file, true).unwrap();
        assert_eq!(truth_bwt.to_vec(), bwt.to_vec());
        assert_eq!(truth_bwt.iter().collect::<Vec<u8>>(), bwt.iter().collect::<Vec<u8>>());
        assert_eq!(truth_bwt.count_kmer(&string_util::convert_stoi(&"$")), 2);
        assert_eq!(truth_bwt.count_kmer(&string_util::convert_stoi(&"ACGT")), 1);
        assert_eq!(truth_bwt.count_kmer(&string_util::convert_stoi(&"TGCA")), 1);
    }

    #[test]
    fn test_run_iter() {
        //let's just do a basic test, all the main stuff happens elsewhere
        let mut bwt: DynamicBWT = Default::default();
        let runs: Vec<(u8, u64)> = bwt.run_iter().collect::<Vec<(u8, u64)>>();
        let expected_runs: Vec<(u8, u64)> = vec![];
        assert_eq!(expected_runs, runs);

        //add one string
        bwt.insert_string("AAAA", true);
        let runs: Vec<(u8, u64)> = bwt.run_iter().collect::<Vec<(u8, u64)>>();
        //AAAA$
        let expected_runs: Vec<(u8, u64)> = vec![(1, 4), (0, 1)];
        assert_eq!(expected_runs, runs);

        //add another
        bwt.insert_string("ACCC", true);
        let runs: Vec<(u8, u64)> = bwt.run_iter().collect::<Vec<(u8, u64)>>();
        //ACAAA$$CCA
        let expected_runs: Vec<(u8, u64)> = vec![(1, 1), (2, 1), (1, 3), (0, 2), (2, 2), (1, 1)];
        assert_eq!(expected_runs, runs);
    }
}