bixverse-rs 0.4.3

//! Contains the MTX file-related parts to do read in data from mtx files and
//! transform them into the binarised files for usage in bixverse-rs

use rayon::prelude::*;
use rustc_hash::FxHashSet;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Read, Result as IoResult, Seek, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use thousands::Separable;

use crate::prelude::*;
use crate::single_cell::sc_data::data_io::{CellGeneSparseWriter, CellOnFileQuality};

/////////
// MTX //
/////////

/// Byte width of one record in a temporary bucket spill file:
/// `cell_idx: u32`, `gene_idx: u32`, `count: u32`.
///
/// These files live and die inside a single
/// [`MtxReader::process_mtx_and_write_bin_streaming`] call, so the layout
/// carries no on-disk compatibility obligation.
const BUCKET_RECORD_LEN: usize = 12;

/// MTX file metadata
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct MtxHeader {
    /// Number of cells identified in the .mtx header.
    pub total_cells: usize,
    /// Number of genes identified in the .mtx header.
    pub total_genes: usize,
    /// Number of entries identified in the .mtx header.
    pub total_entries: usize,
}

/// MTX final data
///
/// Structure to store final results after reading in the .mtx file
#[derive(Debug, Clone)]
pub struct MtxFinalData {
    /// Structure containing the information on which cells/genes to keep and
    /// library size and NNZ for cells.
    pub cell_qc: CellQuality,
    /// No of genes that were read in
    pub no_genes: usize,
    /// No of cells that were read in
    pub no_cells: usize,
}

/// RAII guard that removes temporary bucket files on drop.
///
/// Ensures temp files are cleaned up even if the bucketing or writing pass
/// returns early due to an I/O error or panic.
struct TempFileGuard(Vec<PathBuf>);

/// Drop implementation for `TempFileGuard`
///
/// Attempts to remove each tracked temp file. Errors are silently ignored since
/// this runs during unwind and the files may already have been removed by the
/// writing pass.
impl Drop for TempFileGuard {
    fn drop(&mut self) {
        for p in &self.0 {
            let _ = std::fs::remove_file(p);
        }
    }
}

/// MTX Reader for bixverse
pub struct MtxReader {
    /// Path to the mtx file
    path: PathBuf,
    /// Buffered reader of the mtx file
    reader: BufReader<File>,
    /// Header of the MtxFile
    header: MtxHeader,
    /// Cell and gene quality parameters to apply
    qc_params: MinCellQuality,
    /// Are the cells stored as rows
    cells_as_rows: bool,
}

impl MtxReader {
    /// Generate a new instance of the reader
    ///
    /// ### Params
    ///
    /// * `path` - Path to the mtx file.
    /// * `qc_params` - The min quality parameters that genes and cells have to
    ///   reach and the target library size
    /// * `cells_as_rows` - Boolean. Are the cells the rows (= true) or columns in
    ///   the mtx file.
    pub fn new<P: AsRef<Path>>(
        path: P,
        qc_params: MinCellQuality,
        cells_as_rows: bool,
    ) -> Result<Self, BixverseErrors> {
        let path = path.as_ref().to_path_buf();
        let file = File::open(&path)?;
        let mut reader = BufReader::with_capacity(1024 * 1024, file);

        let header = Self::parse_header(&mut reader, cells_as_rows)?;

        Ok(Self {
            path,
            reader,
            header,
            qc_params,
            cells_as_rows,
        })
    }

    /// Parse the header of the mtx file
    ///
    /// ### Returns
    ///
    /// The `MtxHeader`
    fn parse_header(
        reader: &mut BufReader<File>,
        cells_as_rows: bool,
    ) -> Result<MtxHeader, BixverseErrors> {
        let mut line = String::new();

        loop {
            line.clear();
            reader.read_line(&mut line)?;
            if !line.starts_with('%') {
                break;
            }
        }

        let parts: Vec<&str> = line.split_whitespace().collect();

        if parts.len() != 3 {
            return Err(BixverseErrors::MtxHeaderInvalid(
                "expected three whitespace-separated fields on the shape line",
            ));
        }

        let parse_field = |s: &str, field: &'static str| -> Result<usize, BixverseErrors> {
            s.parse()
                .map_err(|_| BixverseErrors::MtxParseError { field })
        };

        // Header is either `cells genes entries` or `genes cells entries`.
        let (cell_pos, gene_pos) = if cells_as_rows { (0, 1) } else { (1, 0) };

        Ok(MtxHeader {
            total_cells: parse_field(parts[cell_pos], "cell count")?,
            total_genes: parse_field(parts[gene_pos], "gene count")?,
            total_entries: parse_field(parts[2], "entry count")?,
        })
    }

    /// Helper to parse the file to understand which cells to keep
    ///
    /// ### Params
    ///
    /// * `verbose` - Controls verbosity of the function.
    ///
    /// ### Returns
    ///
    /// The CellOnFileQuality file containing the indices and mappings
    /// for the cells/genes to keep.
    pub fn parse_mtx_quality(&mut self, verbose: bool) -> IoResult<CellOnFileQuality> {
        const CHUNK_SIZE: u64 = 64 * 1024 * 1024;
        let file_size = self.reader.get_ref().metadata()?.len();
        let num_chunks = ((file_size / CHUNK_SIZE) as usize).max(1);

        if verbose {
            println!("First file pass - getting gene statistics:");
        }

        let first_scan_time = Instant::now();

        let boundaries = self.find_chunk_boundaries(num_chunks)?;
        let completed_chunks = Arc::new(AtomicUsize::new(0));
        let report_interval = (num_chunks / 10).max(1);

        // per-thread gene counts. MTX coordinate format guarantees each
        // (row, col) pair appears at most once, so counts can be summed across
        // chunks without double-counting.
        let results: Vec<Vec<u32>> = boundaries
            .par_iter()
            .map(|&(start, end)| {
                let mut local_gene_counts = vec![0u32; self.header.total_genes];

                if let Ok(file) = File::open(&self.path) {
                    let mut reader = BufReader::with_capacity(256 * 1024, file);
                    if reader.seek(std::io::SeekFrom::Start(start)).is_ok() {
                        let mut line_buffer = Vec::with_capacity(64);
                        let mut bytes_read = 0u64;

                        while bytes_read < (end - start) {
                            line_buffer.clear();
                            if let Ok(n) = reader.read_until(b'\n', &mut line_buffer) {
                                if n == 0 {
                                    break;
                                }
                                bytes_read += n as u64;

                                let len = line_buffer.len();
                                if len < 3 {
                                    continue;
                                }
                                let trim_end = if line_buffer[len - 1] == b'\n' {
                                    if len > 1 && line_buffer[len - 2] == b'\r' {
                                        len - 2
                                    } else {
                                        len - 1
                                    }
                                } else {
                                    len
                                };

                                if let Some((row, col, _)) =
                                    parse_mtx_line(&line_buffer[..trim_end])
                                {
                                    let gene_idx = if self.cells_as_rows {
                                        (col - 1) as usize
                                    } else {
                                        (row - 1) as usize
                                    };

                                    if gene_idx < self.header.total_genes {
                                        local_gene_counts[gene_idx] += 1;
                                    }
                                }
                            }
                        }
                    }
                }

                if verbose {
                    let completed = completed_chunks.fetch_add(1, Ordering::Relaxed) + 1;
                    if completed.is_multiple_of(report_interval) || completed == num_chunks {
                        let progress = (completed as f64 / num_chunks as f64 * 100.0) as usize;
                        println!(
                            "  Processed {}% of chunks ({}/{})",
                            progress, completed, num_chunks
                        );
                    }
                }

                local_gene_counts
            })
            .collect();

        // Sum counts across threads, then filter genes.
        let genes_to_keep_set: FxHashSet<usize> = (0..self.header.total_genes)
            .into_par_iter()
            .filter(|&i| {
                let total: u32 = results.iter().map(|local| local[i]).sum();
                total as usize >= self.qc_params.min_cells
            })
            .collect();

        drop(results);

        let first_scan_end = first_scan_time.elapsed();

        if verbose {
            println!("First pass done: {:.2?}", first_scan_end);
            println!("Second pass - cell statistics:");
        }

        let second_scan_time = Instant::now();
        let completed_chunks = Arc::new(AtomicUsize::new(0));

        // Parallel second pass - cell stats with filtered genes
        let results: Vec<_> = boundaries
            .par_iter()
            .map(|&(start, end)| {
                let mut local_cell_stats = vec![(0u32, 0u32); self.header.total_cells];

                if let Ok(file) = File::open(&self.path) {
                    let mut reader = BufReader::with_capacity(256 * 1024, file);
                    if reader.seek(std::io::SeekFrom::Start(start)).is_ok() {
                        let mut line_buffer = Vec::with_capacity(64);
                        let mut bytes_read = 0u64;

                        while bytes_read < (end - start) {
                            line_buffer.clear();
                            if let Ok(n) = reader.read_until(b'\n', &mut line_buffer) {
                                if n == 0 {
                                    break;
                                }
                                bytes_read += n as u64;

                                let len = line_buffer.len();
                                if len < 3 {
                                    continue;
                                }
                                let trim_end = if line_buffer[len - 1] == b'\n' {
                                    if len > 1 && line_buffer[len - 2] == b'\r' {
                                        len - 2
                                    } else {
                                        len - 1
                                    }
                                } else {
                                    len
                                };

                                if let Some((row, col, value)) =
                                    parse_mtx_line(&line_buffer[..trim_end])
                                {
                                    let (cell_idx, gene_idx) = if self.cells_as_rows {
                                        ((row - 1) as usize, (col - 1) as usize)
                                    } else {
                                        ((col - 1) as usize, (row - 1) as usize)
                                    };

                                    if genes_to_keep_set.contains(&gene_idx)
                                        && cell_idx < self.header.total_cells
                                    {
                                        local_cell_stats[cell_idx].0 += 1;
                                        local_cell_stats[cell_idx].1 += value;
                                    }
                                }
                            }
                        }
                    }
                }

                if verbose {
                    let completed = completed_chunks.fetch_add(1, Ordering::Relaxed) + 1;
                    if completed.is_multiple_of(report_interval) || completed == num_chunks {
                        let progress = (completed as f64 / num_chunks as f64 * 100.0) as usize;
                        println!(
                            "  Processed {}% of chunks ({}/{})",
                            progress, completed, num_chunks
                        );
                    }
                }

                local_cell_stats
            })
            .collect();

        // Merge cell results
        let mut cell_gene_count = vec![0u32; self.header.total_cells];
        let mut cell_lib_size = vec![0u32; self.header.total_cells];

        for local_cells in results {
            for (i, (count, size)) in local_cells.into_iter().enumerate() {
                cell_gene_count[i] += count;
                cell_lib_size[i] += size;
            }
        }

        // Filter cells
        let cells_to_keep: Vec<usize> = (0..self.header.total_cells)
            .filter(|&i| {
                cell_gene_count[i] as usize >= self.qc_params.min_unique_genes
                    && cell_lib_size[i] as f32 >= self.qc_params.min_lib_size as f32
            })
            .collect();

        let mut genes_to_keep: Vec<usize> = genes_to_keep_set.into_iter().collect();
        genes_to_keep.sort_unstable();

        let mut quality = CellOnFileQuality::new(cells_to_keep, genes_to_keep);
        quality.generate_maps_sets();

        let second_scan_end = second_scan_time.elapsed();

        if verbose {
            println!("Second pass done: {:.2?}", second_scan_end);
            println!(
                "Genes passing QC: {} / {}",
                quality.genes_to_keep.len().separate_with_underscores(),
                self.header.total_genes.separate_with_underscores()
            );
            println!(
                "Cells passing QC: {} / {}",
                quality.cells_to_keep.len().separate_with_underscores(),
                self.header.total_cells.separate_with_underscores()
            );
        }

        Ok(quality)
    }

    /// Process the mtx file and write to binarised Rust file
    ///
    /// ### Params
    ///
    /// * `bin_path` - Where to save the binarised file.
    /// * `quality` - Structure indicating which cells and genes to keep from
    ///   the mtx file.
    /// * `verbose` - Controls verbosity of the function.
    ///
    /// ### Returns
    ///
    /// The `MtxFinalData` with information how many cells were written to
    /// file, how many genes were included in which cells did not parse
    /// the thresholds.
    pub fn process_mtx_and_write_bin(
        mut self,
        bin_path: &str,
        quality: &CellOnFileQuality,
        verbose: bool,
    ) -> Result<MtxFinalData, BixverseErrors> {
        let mut writer = CellGeneSparseWriter::new(
            bin_path,
            true,
            quality.cells_to_keep.len(),
            quality.genes_to_keep.len(),
            self.qc_params.target_size,
        )?;

        // (gene_index, raw_count), both u32: gene index to support >65k
        // features, count to avoid saturating high-expression genes.
        let mut cell_data: Vec<Vec<(u32, u32)>> = vec![Vec::new(); quality.cells_to_keep.len()];
        let mut line_buffer = Vec::with_capacity(64);

        let start_read = Instant::now();

        self.reader.rewind()?;
        Self::skip_header(&mut self.reader)?;

        if verbose {
            println!(
                "Starting to write cells passing quality thresholds in a cell I/O-friendly format to disk."
            )
        }

        let mut lines_read = 0usize;
        let report_interval = (self.header.total_entries / 10).max(1);

        while {
            line_buffer.clear();
            self.reader.read_until(b'\n', &mut line_buffer)? > 0
        } {
            if line_buffer.last() == Some(&b'\n') {
                line_buffer.pop();
            }
            if line_buffer.last() == Some(&b'\r') {
                line_buffer.pop();
            }
            if line_buffer.is_empty() {
                continue;
            }

            let (row, col, value) = match parse_mtx_line(&line_buffer) {
                Some(parsed) => parsed,
                None => continue,
            };

            let (old_cell_idx, old_gene_idx) = if self.cells_as_rows {
                ((row - 1) as usize, (col - 1) as usize)
            } else {
                ((col - 1) as usize, (row - 1) as usize)
            };

            if !quality.genes_to_keep_set.contains(&old_gene_idx)
                || !quality.cells_to_keep_set.contains(&old_cell_idx)
            {
                continue;
            }

            let new_cell_idx = quality.cell_old_to_new[&old_cell_idx];
            let new_gene_idx = quality.gene_old_to_new[&old_gene_idx] as u32;

            cell_data[new_cell_idx].push((new_gene_idx, value));

            lines_read += 1;
            if verbose && lines_read.is_multiple_of(report_interval) {
                let progress =
                    (lines_read as f64 / self.header.total_entries as f64 * 100.0) as usize;
                println!("  Processed {}% of entries", progress);
            }
        }

        let mut lib_size = Vec::with_capacity(quality.cells_to_keep.len());
        let mut nnz = Vec::with_capacity(quality.cells_to_keep.len());

        for (cell_idx, data) in cell_data.iter_mut().enumerate() {
            if data.is_empty() {
                continue;
            }

            data.sort_by_key(|(gene_idx, _)| *gene_idx);

            let gene_indices: Vec<u32> = data.iter().map(|(g, _)| *g).collect();
            let gene_counts: Vec<u32> = data.iter().map(|(_, c)| *c).collect();

            let total_umi: u64 = gene_counts.iter().map(|&x| x as u64).sum();
            let n_genes = gene_counts.len();

            lib_size.push(total_umi as usize);
            nnz.push(n_genes);

            let cell_chunk = CsrCellChunk::from_data(
                &gene_counts,
                &gene_indices,
                cell_idx,
                self.qc_params.target_size,
                true,
            );
            writer.write_cell_chunk(cell_chunk)?;
        }

        writer.finalise()?;

        let end_read = start_read.elapsed();

        if verbose {
            println!("Reading in cell data done: {:.2?}", end_read);
        }

        let cell_quality = CellQuality {
            cell_indices: quality.cells_to_keep.to_vec(),
            gene_indices: quality.genes_to_keep.to_vec(),
            lib_size,
            nnz,
        };

        Ok(MtxFinalData {
            cell_qc: cell_quality,
            no_genes: quality.genes_to_keep.len(),
            no_cells: quality.cells_to_keep.len(),
        })
    }

    /// Process the mtx file and write to binarised Rust file
    ///
    /// Streams the mtx file in two passes via temp file bucketing, supporting
    /// both `cells_as_rows = true` and `cells_as_rows = false` without requiring
    /// the input to be sorted. Memory usage is bounded by the size of a single
    /// bucket rather than the total kept entries.
    ///
    /// ### Params
    ///
    /// * `bin_path` - Where to save the binarised file.
    /// * `quality` - Structure indicating which cells and genes to keep from
    ///   the mtx file.
    /// * `verbose` - Controls verbosity of the function.
    ///
    /// ### Returns
    ///
    /// The `MtxFinalData` with information how many cells were written to
    /// file, how many genes were included and the per-cell library size and
    /// NNZ statistics.
    pub fn process_mtx_and_write_bin_streaming(
        mut self,
        bin_path: &str,
        quality: &CellOnFileQuality,
        verbose: bool,
    ) -> Result<MtxFinalData, BixverseErrors> {
        let n_kept_cells = quality.cells_to_keep.len();
        let n_kept_genes = quality.genes_to_keep.len();

        if n_kept_cells == 0 {
            let writer = CellGeneSparseWriter::new(
                bin_path,
                true,
                0,
                n_kept_genes,
                self.qc_params.target_size,
            )?;
            writer.finalise()?;
            return Ok(MtxFinalData {
                cell_qc: CellQuality {
                    cell_indices: vec![],
                    gene_indices: quality.genes_to_keep.to_vec(),
                    lib_size: vec![],
                    nnz: vec![],
                },
                no_genes: n_kept_genes,
                no_cells: 0,
            });
        }

        // Bucket layout. Cap at 128 buckets to stay well below macOS fd limits.
        let n_buckets = n_kept_cells.div_ceil(5_000).clamp(1, 128);
        let cells_per_bucket = n_kept_cells.div_ceil(n_buckets);

        let bin_path_buf = PathBuf::from(bin_path);
        let temp_dir = bin_path_buf
            .parent()
            .filter(|p| !p.as_os_str().is_empty())
            .map(|p| p.to_path_buf())
            .unwrap_or_else(|| PathBuf::from("."));
        let stem = bin_path_buf
            .file_name()
            .map(|s| s.to_string_lossy().to_string())
            .unwrap_or_else(|| "mtx".into());

        let temp_paths: Vec<PathBuf> = (0..n_buckets)
            .map(|i| temp_dir.join(format!(".{}.bucket_{:04}.tmp", stem, i)))
            .collect();
        let _guard = TempFileGuard(temp_paths.clone());

        let pass1 = Instant::now();
        if verbose {
            println!(
                "Bucketing pass: {} buckets, ~{} cells/bucket",
                n_buckets, cells_per_bucket
            );
        }

        let mut bucket_writers: Vec<BufWriter<File>> = temp_paths
            .iter()
            .map(|p| File::create(p).map(|f| BufWriter::with_capacity(256 * 1024, f)))
            .collect::<IoResult<Vec<_>>>()?;

        self.reader.rewind()?;
        Self::skip_header(&mut self.reader)?;

        let mut line_buffer = Vec::with_capacity(64);
        let mut lines_read = 0usize;
        let report_interval = (self.header.total_entries / 10).max(1);

        while {
            line_buffer.clear();
            self.reader.read_until(b'\n', &mut line_buffer)? > 0
        } {
            if line_buffer.last() == Some(&b'\n') {
                line_buffer.pop();
            }
            if line_buffer.last() == Some(&b'\r') {
                line_buffer.pop();
            }
            if line_buffer.is_empty() {
                continue;
            }

            let (row, col, value) = match parse_mtx_line(&line_buffer) {
                Some(parsed) => parsed,
                None => continue,
            };

            let (old_cell_idx, old_gene_idx) = if self.cells_as_rows {
                ((row - 1) as usize, (col - 1) as usize)
            } else {
                ((col - 1) as usize, (row - 1) as usize)
            };

            if !quality.genes_to_keep_set.contains(&old_gene_idx)
                || !quality.cells_to_keep_set.contains(&old_cell_idx)
            {
                continue;
            }

            let new_cell_idx = quality.cell_old_to_new[&old_cell_idx] as u32;
            let new_gene_idx = quality.gene_old_to_new[&old_gene_idx] as u32;
            let bucket = (new_cell_idx as usize) / cells_per_bucket;

            let mut buf = [0u8; BUCKET_RECORD_LEN];
            buf[0..4].copy_from_slice(&new_cell_idx.to_le_bytes());
            buf[4..8].copy_from_slice(&new_gene_idx.to_le_bytes());
            buf[8..12].copy_from_slice(&value.to_le_bytes());
            bucket_writers[bucket].write_all(&buf)?;

            lines_read += 1;
            if verbose && lines_read.is_multiple_of(report_interval) {
                let progress =
                    (lines_read as f64 / self.header.total_entries as f64 * 100.0) as usize;
                println!("  Bucketed {}% of entries", progress);
            }
        }

        for w in bucket_writers.iter_mut() {
            w.flush()?;
        }
        drop(bucket_writers);

        if verbose {
            println!("Bucketing done: {:.2?}", pass1.elapsed());
        }

        let mut writer = CellGeneSparseWriter::new(
            bin_path,
            true,
            n_kept_cells,
            n_kept_genes,
            self.qc_params.target_size,
        )?;
        let mut lib_size = Vec::with_capacity(n_kept_cells);
        let mut nnz = Vec::with_capacity(n_kept_cells);

        let pass2 = Instant::now();
        if verbose {
            println!("Writing pass: streaming buckets to output");
        }

        for (bucket_idx, temp_path) in temp_paths.iter().enumerate() {
            let bucket_bytes = std::fs::read(temp_path)?;
            let n_entries = bucket_bytes.len() / BUCKET_RECORD_LEN;

            let mut entries: Vec<(u32, u32, u32)> = Vec::with_capacity(n_entries);
            for i in 0..n_entries {
                let off = i * BUCKET_RECORD_LEN;
                let cell = u32::from_le_bytes([
                    bucket_bytes[off],
                    bucket_bytes[off + 1],
                    bucket_bytes[off + 2],
                    bucket_bytes[off + 3],
                ]);
                let gene = u32::from_le_bytes([
                    bucket_bytes[off + 4],
                    bucket_bytes[off + 5],
                    bucket_bytes[off + 6],
                    bucket_bytes[off + 7],
                ]);
                let value = u32::from_le_bytes([
                    bucket_bytes[off + 8],
                    bucket_bytes[off + 9],
                    bucket_bytes[off + 10],
                    bucket_bytes[off + 11],
                ]);
                entries.push((cell, gene, value));
            }
            drop(bucket_bytes);

            entries.sort_unstable_by_key(|&(c, g, _)| (c, g));

            let mut i = 0;
            while i < entries.len() {
                let cell = entries[i].0;
                let mut j = i;
                while j < entries.len() && entries[j].0 == cell {
                    j += 1;
                }

                let gene_indices: Vec<u32> = entries[i..j].iter().map(|&(_, g, _)| g).collect();
                let gene_counts: Vec<u32> = entries[i..j].iter().map(|&(_, _, v)| v).collect();
                let total_umi: u64 = gene_counts.iter().map(|&x| x as u64).sum();

                lib_size.push(total_umi as usize);
                nnz.push(gene_counts.len());

                let cell_chunk = CsrCellChunk::from_data(
                    &gene_counts,
                    &gene_indices,
                    cell as usize,
                    self.qc_params.target_size,
                    true,
                );
                writer.write_cell_chunk(cell_chunk)?;

                i = j;
            }

            let _ = std::fs::remove_file(temp_path);

            if verbose
                && ((bucket_idx + 1) % (n_buckets / 10).max(1) == 0 || bucket_idx + 1 == n_buckets)
            {
                let progress = ((bucket_idx + 1) as f64 / n_buckets as f64 * 100.0) as usize;
                println!(
                    "  Wrote bucket {}/{} ({}%)",
                    bucket_idx + 1,
                    n_buckets,
                    progress
                );
            }
        }

        writer.finalise()?;

        if verbose {
            println!("Writing pass done: {:.2?}", pass2.elapsed());
        }

        let cell_quality = CellQuality {
            cell_indices: quality.cells_to_keep.to_vec(),
            gene_indices: quality.genes_to_keep.to_vec(),
            lib_size,
            nnz,
        };

        Ok(MtxFinalData {
            cell_qc: cell_quality,
            no_genes: n_kept_genes,
            no_cells: n_kept_cells,
        })
    }

    /// Helper function to skip the header
    fn skip_header(reader: &mut BufReader<File>) -> IoResult<()> {
        let mut line = String::new();
        loop {
            line.clear();
            reader.read_line(&mut line)?;
            if !line.starts_with('%') {
                break;
            }
        }
        Ok(())
    }

    /// Generate chunk boundaries for parallel processing
    ///
    /// ### Params
    ///
    /// * `num_chunks` - Number of desired chunks
    ///
    /// ### Results
    ///
    /// A vector of tuples indicating the chunk boundaries.
    fn find_chunk_boundaries(&mut self, num_chunks: usize) -> IoResult<Vec<(u64, u64)>> {
        let file_size = self.reader.get_ref().metadata()?.len();

        self.reader.rewind()?;
        let mut line = String::new();
        loop {
            line.clear();
            self.reader.read_line(&mut line)?;
            if !line.starts_with('%') {
                break;
            }
        }

        let data_start = self.reader.stream_position()?;

        let data_size = file_size - data_start;
        let chunk_size = data_size / num_chunks as u64;

        let mut boundaries = vec![(data_start, data_start)];

        for i in 1..num_chunks {
            let target_pos = data_start + (chunk_size * i as u64);
            if target_pos >= file_size {
                break;
            }

            self.reader.seek(std::io::SeekFrom::Start(target_pos))?;

            let mut byte = [0u8; 1];
            while self.reader.read(&mut byte)? > 0 {
                if byte[0] == b'\n' {
                    break;
                }
            }

            let boundary = self.reader.stream_position()?;
            boundaries.push((boundary, boundary));
        }

        boundaries.push((file_size, file_size));

        for i in 0..boundaries.len() - 1 {
            boundaries[i].1 = boundaries[i + 1].0;
        }
        boundaries.pop();

        Ok(boundaries)
    }
}

/////////////
// Helpers //
/////////////

/// Parse mtx line from bytes
///
/// ### Params
///
/// * `line` - The file line as bytes
///
/// ### Return
///
/// Returns an Option of a tuple representing
/// `<row_index, col_index, raw_count>`. Counts are kept at full `u32` width;
/// the on-disk chunk format narrows to u16 only when every value fits.
#[inline]
fn parse_mtx_line(line: &[u8]) -> Option<(u32, u32, u32)> {
    let mut i = 0;
    let len = line.len();

    // Parse first number (row)
    let mut row = 0u32;
    while i < len && line[i].is_ascii_digit() {
        row = row * 10 + (line[i] - b'0') as u32;
        i += 1;
    }
    if i == 0 {
        return None;
    }

    // Skip whitespace
    while i < len && (line[i] == b' ' || line[i] == b'\t') {
        i += 1;
    }
    if i >= len {
        return None;
    }

    // Parse second number (col)
    let mut col = 0u32;
    while i < len && line[i].is_ascii_digit() {
        col = col * 10 + (line[i] - b'0') as u32;
        i += 1;
    }

    // Skip whitespace
    while i < len && (line[i] == b' ' || line[i] == b'\t') {
        i += 1;
    }
    if i >= len {
        return None;
    }

    // Parse third number (value)
    let mut val = 0u32;
    while i < len && line[i].is_ascii_digit() {
        val = val * 10 + (line[i] - b'0') as u32;
        i += 1;
    }

    Some((row, col, val))
}

///////////
// Tests //
///////////

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

    ///////////////
    // Fixtures //
    ///////////////

    /// Three cells x four genes, cells as rows, five entries. All indices are
    /// base-1 as the Matrix Market spec demands.
    ///
    /// ```text
    ///        g0 g1 g2 g3
    ///   c0 [  5  .  7  . ]
    ///   c1 [  .  3  .  . ]
    ///   c2 [  2  .  .  9 ]
    /// ```
    const CELLS_AS_ROWS_MTX: &str = "%%MatrixMarket matrix coordinate integer general\n\
%\n\
3 4 5\n\
1 1 5\n\
1 3 7\n\
2 2 3\n\
3 1 2\n\
3 4 9\n";

    /// The transpose of [`CELLS_AS_ROWS_MTX`], i.e. genes as rows.
    const GENES_AS_ROWS_MTX: &str = "%%MatrixMarket matrix coordinate integer general\n\
%\n\
4 3 5\n\
1 1 5\n\
3 1 7\n\
2 2 3\n\
1 3 2\n\
4 3 9\n";

    /// The `(gene indices, raw counts)` the two fixtures both encode.
    fn expected_cells() -> Vec<(Vec<u32>, Vec<u32>)> {
        vec![
            (vec![0, 2], vec![5, 7]),
            (vec![1], vec![3]),
            (vec![0, 3], vec![2, 9]),
        ]
    }

    /// RAII guard that removes a test's scratch file even if an assert fails.
    struct TempPath(PathBuf);

    /// Drop implementation for [`TempPath`]. Errors are ignored: the file may
    /// already be gone, and this runs during unwind.
    impl Drop for TempPath {
        fn drop(&mut self) {
            let _ = std::fs::remove_file(&self.0);
        }
    }

    impl TempPath {
        /// Reserve a scratch file named after the calling test. The test name
        /// is the uniqueness guarantee: `cargo test` runs the whole module in
        /// one process, so a PID alone would collide across threads.
        fn new(name: &str, ext: &str) -> Self {
            Self(std::env::temp_dir().join(format!("bixverse_mtx_io_{name}.{ext}")))
        }

        /// Path of the guarded file as a `&str`.
        fn path(&self) -> &str {
            self.0.to_str().expect("temp path is valid UTF-8")
        }
    }

    /// Write `body` to a scratch `.mtx` file and hand back its guard.
    fn write_mtx(name: &str, body: &str) -> TempPath {
        let temp = TempPath::new(name, "mtx");
        std::fs::write(&temp.0, body).expect("mtx file written");
        temp
    }

    /// QC thresholds that keep every cell and gene of the fixtures.
    fn keep_all_qc() -> MinCellQuality {
        MinCellQuality {
            min_unique_genes: 1,
            min_lib_size: 1,
            min_cells: 1,
            target_size: 1e4,
        }
    }

    /// Read the first `n_cells` cells of a written store back as
    /// `(gene indices, raw counts)`.
    fn read_back(bin_path: &str, n_cells: usize) -> Vec<(Vec<u32>, Vec<u32>)> {
        let reader = ParallelSparseReader::new(bin_path).expect("reader opens");
        let indices: Vec<usize> = (0..n_cells).collect();
        reader
            .read_cells_parallel(&indices)
            .expect("cells read back")
            .into_iter()
            .map(|c| (c.indices.clone(), c.data_raw.iter().collect()))
            .collect()
    }

    /////////////////
    // Line parser //
    /////////////////

    /// Regression: the parser used to saturate counts at `u16::MAX`.
    #[test]
    fn test_parse_mtx_line_keeps_full_u32_counts() {
        assert_eq!(parse_mtx_line(b"1 2 3"), Some((1, 2, 3)));
        assert_eq!(parse_mtx_line(b"7\t9\t65535"), Some((7, 9, 65_535)));
        assert_eq!(parse_mtx_line(b"4 5 70000"), Some((4, 5, 70_000)));
        assert_eq!(parse_mtx_line(b"1 2 4294967295"), Some((1, 2, u32::MAX)));
    }

    /// Empty, non-numeric and short lines give `None`, not a partial triplet.
    #[test]
    fn test_parse_mtx_line_rejects_malformed_input() {
        assert_eq!(parse_mtx_line(b""), None);
        assert_eq!(parse_mtx_line(b"abc"), None);
        assert_eq!(parse_mtx_line(b"1"), None);
        assert_eq!(parse_mtx_line(b"1 2"), None);
    }

    ////////////
    // Header //
    ////////////

    /// The shape line is split by orientation, so both readings of the same
    /// three numbers have to be pinned.
    #[test]
    fn mtx_header_shape_line_is_read_per_orientation() {
        let mtx = write_mtx("header_shape", CELLS_AS_ROWS_MTX);

        let cells_as_rows = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        assert_eq!(cells_as_rows.header.total_cells, 3);
        assert_eq!(cells_as_rows.header.total_genes, 4);
        assert_eq!(cells_as_rows.header.total_entries, 5);

        // Same file read the other way round: the first two fields swap roles.
        let genes_as_rows = MtxReader::new(mtx.path(), keep_all_qc(), false).expect("reader opens");
        assert_eq!(genes_as_rows.header.total_cells, 4);
        assert_eq!(genes_as_rows.header.total_genes, 3);
        assert_eq!(genes_as_rows.header.total_entries, 5);
    }

    /// A shape line that is not three fields must raise `MtxHeaderInvalid`
    /// rather than index past the end of the split.
    #[test]
    fn mtx_header_with_wrong_field_count_is_rejected() {
        let mtx = write_mtx(
            "header_two_fields",
            "%%MatrixMarket matrix coordinate integer general\n3 4\n1 1 5\n",
        );

        assert!(matches!(
            MtxReader::new(mtx.path(), keep_all_qc(), true),
            Err(BixverseErrors::MtxHeaderInvalid(_))
        ));
    }

    /// `MtxParseError` has to name the field that failed, and the name depends
    /// on the orientation because the first two fields swap.
    #[test]
    fn mtx_header_parse_error_names_the_failing_field() {
        let bad_second = write_mtx(
            "header_bad_second",
            "%%MatrixMarket matrix coordinate integer general\n3 x 5\n",
        );
        assert!(matches!(
            MtxReader::new(bad_second.path(), keep_all_qc(), true),
            Err(BixverseErrors::MtxParseError {
                field: "gene count"
            })
        ));
        // Cells are the columns now, so the same broken field is the cell count.
        assert!(matches!(
            MtxReader::new(bad_second.path(), keep_all_qc(), false),
            Err(BixverseErrors::MtxParseError {
                field: "cell count"
            })
        ));

        let bad_third = write_mtx(
            "header_bad_third",
            "%%MatrixMarket matrix coordinate integer general\n3 4 x\n",
        );
        assert!(matches!(
            MtxReader::new(bad_third.path(), keep_all_qc(), true),
            Err(BixverseErrors::MtxParseError {
                field: "entry count"
            })
        ));
    }

    /// A missing file must surface as an I/O error, not a panic on unwrap.
    #[test]
    fn mtx_reader_reports_a_missing_file() {
        let missing = std::env::temp_dir().join("bixverse_mtx_io_does_not_exist.mtx");
        assert!(matches!(
            MtxReader::new(&missing, keep_all_qc(), true),
            Err(BixverseErrors::BinaryIo(_))
        ));
    }

    /////////////////
    // Quality pass //
    /////////////////

    /// `parse_mtx_quality` drives every downstream index remap, so pin the
    /// gene-then-cell filter order and the old-to-new maps it builds.
    #[test]
    fn parse_mtx_quality_filters_genes_then_cells() {
        let mtx = write_mtx("quality_filters", CELLS_AS_ROWS_MTX);

        // Everything permissive: nothing is dropped and the maps are identities.
        let mut reader = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        let quality = reader.parse_mtx_quality(false).expect("quality pass");
        assert_eq!(quality.cells_to_keep, vec![0, 1, 2]);
        assert_eq!(quality.genes_to_keep, vec![0, 1, 2, 3]);

        // min_cells = 2 keeps only gene 0, which is seen in cells 0 and 2.
        // Cell 1 then has zero kept entries and fails min_unique_genes = 1.
        let qc = MinCellQuality {
            min_unique_genes: 1,
            min_lib_size: 1,
            min_cells: 2,
            target_size: 1e4,
        };
        let mut reader = MtxReader::new(mtx.path(), qc, true).expect("reader opens");
        let quality = reader.parse_mtx_quality(false).expect("quality pass");
        assert_eq!(quality.genes_to_keep, vec![0]);
        assert_eq!(quality.cells_to_keep, vec![0, 2]);
        assert_eq!(quality.cell_old_to_new[&0], 0);
        assert_eq!(quality.cell_old_to_new[&2], 1);
        assert_eq!(quality.gene_old_to_new[&0], 0);
    }

    /////////////////
    // Round trips //
    /////////////////

    /// The only mtx -> bin -> reader path in the crate. Counts and gene indices
    /// have to survive the conversion untouched.
    #[test]
    fn mtx_round_trip_preserves_counts_and_indices() {
        let mtx = write_mtx("round_trip", CELLS_AS_ROWS_MTX);
        let bin = TempPath::new("round_trip", "bin");

        let mut reader = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        let quality = reader.parse_mtx_quality(false).expect("quality pass");
        let final_data = reader
            .process_mtx_and_write_bin(bin.path(), &quality, false)
            .expect("conversion");

        assert_eq!(final_data.no_cells, 3);
        assert_eq!(final_data.no_genes, 4);
        // Library sizes: 5 + 7 = 12, 3, 2 + 9 = 11.
        assert_eq!(final_data.cell_qc.lib_size, vec![12, 3, 11]);
        assert_eq!(final_data.cell_qc.nnz, vec![2, 1, 2]);

        assert_eq!(read_back(bin.path(), 3), expected_cells());
    }

    /// Matrix Market is base-1. The lowest legal index has to land on 0 and the
    /// highest on `n - 1`, with nothing shifted in between.
    #[test]
    fn mtx_indices_shift_from_base_one_to_base_zero() {
        let mtx = write_mtx(
            "base_one",
            "%%MatrixMarket matrix coordinate integer general\n2 3 2\n1 1 42\n2 3 7\n",
        );
        let bin = TempPath::new("base_one", "bin");

        let mut reader = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        let quality = reader.parse_mtx_quality(false).expect("quality pass");
        reader
            .process_mtx_and_write_bin(bin.path(), &quality, false)
            .expect("conversion");

        // Only genes 1 and 3 (base-1) carry counts, so the kept genes are the
        // old 0 and 2, renumbered to 0 and 1.
        assert_eq!(quality.genes_to_keep, vec![0, 2]);
        assert_eq!(
            read_back(bin.path(), 2),
            vec![(vec![0], vec![42]), (vec![1], vec![7])]
        );
    }

    /// The two orientations parse different files but must produce the same
    /// store, which is the only cross-check on the row/column swap.
    #[test]
    fn mtx_gene_as_rows_matches_cell_as_rows() {
        let mtx = write_mtx("orientation", GENES_AS_ROWS_MTX);
        let bin = TempPath::new("orientation", "bin");

        let mut reader = MtxReader::new(mtx.path(), keep_all_qc(), false).expect("reader opens");
        let quality = reader.parse_mtx_quality(false).expect("quality pass");
        let final_data = reader
            .process_mtx_and_write_bin(bin.path(), &quality, false)
            .expect("conversion");

        assert_eq!(final_data.no_cells, 3);
        assert_eq!(final_data.no_genes, 4);
        assert_eq!(final_data.cell_qc.lib_size, vec![12, 3, 11]);
        assert_eq!(read_back(bin.path(), 3), expected_cells());
    }

    /// The streaming writer buckets to temp files and re-sorts; it must land on
    /// exactly the same store as the in-memory writer.
    #[test]
    fn mtx_streaming_matches_non_streaming() {
        let mtx = write_mtx("streaming", CELLS_AS_ROWS_MTX);
        let bin_plain = TempPath::new("streaming_plain", "bin");
        let bin_stream = TempPath::new("streaming_stream", "bin");

        let mut reader = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        let quality = reader.parse_mtx_quality(false).expect("quality pass");
        let plain = reader
            .process_mtx_and_write_bin(bin_plain.path(), &quality, false)
            .expect("conversion");

        let reader = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        let streamed = reader
            .process_mtx_and_write_bin_streaming(bin_stream.path(), &quality, false)
            .expect("streaming conversion");

        assert_eq!(streamed.no_cells, plain.no_cells);
        assert_eq!(streamed.no_genes, plain.no_genes);
        assert_eq!(streamed.cell_qc.lib_size, plain.cell_qc.lib_size);
        assert_eq!(streamed.cell_qc.nnz, plain.cell_qc.nnz);
        assert_eq!(
            read_back(bin_stream.path(), 3),
            read_back(bin_plain.path(), 3)
        );
    }

    /// An empty input has no cells to keep; the streaming path short-circuits
    /// there and still has to leave a readable, empty store behind.
    #[test]
    fn mtx_streaming_writes_a_readable_store_when_nothing_is_kept() {
        let bin = TempPath::new("streaming_empty", "bin");
        let mtx = write_mtx("streaming_empty", CELLS_AS_ROWS_MTX);

        let reader = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        let mut quality = CellOnFileQuality::new(vec![], vec![0, 1, 2, 3]);
        quality.generate_maps_sets();

        let final_data = reader
            .process_mtx_and_write_bin_streaming(bin.path(), &quality, false)
            .expect("streaming conversion");

        assert_eq!(final_data.no_cells, 0);
        assert_eq!(final_data.no_genes, 4);
        assert!(final_data.cell_qc.lib_size.is_empty());

        let reader = ParallelSparseReader::new(bin.path()).expect("reader opens");
        assert!(reader.is_cell_based());
        assert_eq!(reader.get_header().total_cells, 0);
    }

    /// Defect pin: a kept cell with no kept entries is never written, so
    /// `lib_size` / `nnz` come back shorter than `cell_indices` and the two
    /// stop lining up. Both writers agree on the flaw.
    #[test]
    fn mtx_kept_cell_without_entries_desyncs_the_quality_vectors() {
        let mtx = write_mtx("empty_cell", CELLS_AS_ROWS_MTX);
        let bin = TempPath::new("empty_cell", "bin");

        // Keep every cell but only gene 0, which cell 1 does not express.
        let reader = MtxReader::new(mtx.path(), keep_all_qc(), true).expect("reader opens");
        let mut quality = CellOnFileQuality::new(vec![0, 1, 2], vec![0]);
        quality.generate_maps_sets();

        let final_data = reader
            .process_mtx_and_write_bin(bin.path(), &quality, false)
            .expect("conversion");

        assert_eq!(final_data.cell_qc.cell_indices.len(), 3);
        assert_eq!(final_data.cell_qc.lib_size, vec![5, 2]);
        assert_eq!(final_data.cell_qc.nnz, vec![1, 1]);

        // Cell 1 has no chunk on disk at all.
        let store = ParallelSparseReader::new(bin.path()).expect("reader opens");
        assert!(matches!(
            store.read_cells_parallel(&[1]),
            Err(BixverseErrors::ChunkIndexNotFound(1))
        ));
    }
}