bixverse-rs 0.5.4

//! Contains the streaming engine for bixverse single cell application. This
//! has the structures for cells/spots and genes ready, plus various utilities,
//! writers and readers.

use bincode::{Decode, Encode, config, decode_from_slice, serde::encode_to_vec};
use half::f16;
use indexmap::IndexSet;
use lz4_flex::{compress_prepend_size, decompress_size_prepended};
use memmap2::MmapOptions;
use num_traits::FromPrimitive;
use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};
use serde::{Deserialize, Serialize};
use std::{
    fs::File,
    io::{BufWriter, Seek, SeekFrom, Write},
    marker::Sync,
    path::Path,
    sync::Arc,
};

use crate::prelude::*;

///////////////////
// Count wrapper //
///////////////////

/// Wrapper enum for raw counts that supports both u16 and u32 storage.
///
/// Most single cell data sets have counts well within the u16 range (0-65,535).
/// The u32 variant exists for edge cases such as mixed-omics data where counts
/// can exceed this limit. The variant is selected automatically at construction
/// time based on the data, and is encoded as a discriminant byte in the binary
/// chunk format.
#[derive(Encode, Decode, Serialize, Deserialize, Debug, Clone)]
pub enum RawCounts {
    /// Stores u16 values
    U16(Vec<u16>),
    /// Stores u32 values (for the rare case where a single gene exceeds the
    /// counts here)
    U32(Vec<u32>),
}

/// Raw count element size discriminant for the binary format.
/// Stored as a single byte in chunk headers.
const RAW_ELEM_U16: u8 = 2;
const RAW_ELEM_U32: u8 = 4;

/// Header length of a CSR cell chunk: 4+4+4+8+8+4 bytes. See
/// [`CsrCellChunk::write_to_bytes`] for the field layout.
const CSR_CHUNK_HEADER_LEN: usize = 32;

/// Header length of a CSC gene chunk: 4+4+4+2+2+8+8+4 bytes. See
/// [`CscGeneChunk::write_to_bytes`] for the field layout.
const CSC_CHUNK_HEADER_LEN: usize = 36;

/// Flush cadence for gene-based writers, in chunks.
///
/// Gene chunks are far fewer and far larger than cell chunks, so a tight
/// cadence keeps peak buffer occupancy bounded without costing many syscalls.
const GENE_FLUSH_FREQUENCY: usize = 1_000;

/// Flush cadence for cell-based writers, in chunks.
///
/// Cell chunks are small, so flushing this rarely effectively leaves the
/// 128 MiB [`BufWriter`] to manage itself. Kept as a backstop only.
const CELL_FLUSH_FREQUENCY: usize = 100_000;

impl RawCounts {
    /// Get a single value as u32
    ///
    /// ### Returns
    ///
    /// The count as u32
    pub fn get(&self, i: usize) -> u32 {
        match self {
            RawCounts::U16(v) => v[i] as u32,
            RawCounts::U32(v) => v[i],
        }
    }

    /// Total number of elements
    ///
    /// ### Returns
    ///
    /// The length of elements
    pub fn len(&self) -> usize {
        match self {
            RawCounts::U16(v) => v.len(),
            RawCounts::U32(v) => v.len(),
        }
    }

    /// Whether the vector is empty
    ///
    /// ### Returns
    ///
    /// Boolean if the data is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Iterator that yields all values as u32
    ///
    /// ### Returns
    ///
    /// An iterator
    pub fn iter(&self) -> RawCountsIter<'_> {
        RawCountsIter {
            inner: self,
            pos: 0,
        }
    }

    /// Element size in bytes (2 for u16, 4 for u32)
    ///
    /// ### Returns
    ///
    /// The element size
    pub fn elem_size(&self) -> u8 {
        match self {
            RawCounts::U16(_) => RAW_ELEM_U16,
            RawCounts::U32(_) => RAW_ELEM_U32,
        }
    }

    /// Write raw bytes to a writer
    ///
    /// ### Params
    ///
    /// * `writer` - The writer to use
    ///
    /// ### Returns
    ///
    /// The I/O result
    pub fn write_bytes(&self, writer: &mut impl Write) -> std::io::Result<()> {
        // Little-endian throughout, matching every header field in the format.
        match self {
            RawCounts::U16(v) => {
                let mut bytes = Vec::with_capacity(v.len() * 2);
                for x in v {
                    bytes.extend_from_slice(&x.to_le_bytes());
                }
                writer.write_all(&bytes)
            }
            RawCounts::U32(v) => {
                let mut bytes = Vec::with_capacity(v.len() * 4);
                for x in v {
                    bytes.extend_from_slice(&x.to_le_bytes());
                }
                writer.write_all(&bytes)
            }
        }
    }

    /// Read raw counts from a buffer given element size and count
    ///
    /// ### Params
    ///
    /// * `buffer` - The byte buffer to read from (must be positioned at the
    ///   start of the raw count data)
    /// * `count` - Number of elements to read
    /// * `elem_size` - Element size in bytes (2 for u16, 4 for u32)
    ///
    /// ### Returns
    ///
    /// The `RawCounts`, or [`BixverseErrors::RawElemSizeInvalid`] if the
    /// discriminant is unrecognised, or [`BixverseErrors::ChunkPayloadTruncated`]
    /// if the buffer is shorter than `count * elem_size`.
    pub fn read_from_buffer(
        buffer: &[u8],
        count: usize,
        elem_size: u8,
    ) -> Result<Self, BixverseErrors> {
        // Zero covers legacy files where the discriminant byte was padding.
        let width = match elem_size {
            RAW_ELEM_U32 => 4_usize,
            RAW_ELEM_U16 | 0 => 2_usize,
            other => return Err(BixverseErrors::RawElemSizeInvalid(other)),
        };

        let needed = count * width;
        if buffer.len() < needed {
            return Err(BixverseErrors::ChunkPayloadTruncated {
                expected: needed,
                found: buffer.len(),
            });
        }

        Ok(if width == 4 {
            RawCounts::U32(
                buffer[..needed]
                    .chunks_exact(4)
                    .map(|c| {
                        u32::from_le_bytes(c.try_into().expect("4-byte chunk by construction"))
                    })
                    .collect(),
            )
        } else {
            RawCounts::U16(
                buffer[..needed]
                    .chunks_exact(2)
                    .map(|c| {
                        u16::from_le_bytes(c.try_into().expect("2-byte chunk by construction"))
                    })
                    .collect(),
            )
        })
    }

    /// Build RawCounts from a u32 slice, auto-selecting the variant.
    ///
    /// If all values fit in u16, stores as U16 for better compression.
    /// Otherwise stores as U32.
    pub fn from_u32_auto(data: &[u32]) -> Self {
        if data.iter().all(|&v| v <= u16::MAX as u32) {
            RawCounts::U16(data.iter().map(|&v| v as u16).collect())
        } else {
            RawCounts::U32(data.to_vec())
        }
    }
}

/// Structure of the RawCountsIterator
pub struct RawCountsIter<'a> {
    /// The reference to the row counts
    inner: &'a RawCounts,
    /// The position
    pos: usize,
}

impl<'a> Iterator for RawCountsIter<'a> {
    type Item = u32;
    fn next(&mut self) -> Option<u32> {
        if self.pos < self.inner.len() {
            let val = self.inner.get(self.pos);
            self.pos += 1;
            Some(val)
        } else {
            None
        }
    }
}

//////////////////
// Cell Quality //
//////////////////

/// Structure to store QC information on cells
#[derive(Clone, Debug)]
pub struct CellQuality {
    /// Indices of which cells to keep.
    pub cell_indices: Vec<usize>,
    /// Indices of which genes to keep
    pub gene_indices: Vec<usize>,
    /// Library size of cells
    pub lib_size: Vec<usize>,
    /// Number of cells expressing this gene
    pub nnz: Vec<usize>,
}

impl CellQuality {
    /// Update the internal cell indices
    ///
    /// ### Params
    ///
    /// * `cell_indices` - Vector of cell indices to keep
    pub fn set_cell_indices(&mut self, cell_indices: &[usize]) {
        self.cell_indices = cell_indices.to_vec();
    }

    /// Update the internal gene indices
    ///
    /// ### Params
    ///
    /// * `gene_indices` - Vector of gene indices to keep
    pub fn set_gene_indices(&mut self, gene_indices: &[usize]) {
        self.gene_indices = gene_indices.to_vec();
    }
}

/// Structure that stores minimum QC thresholds/info for single cell
#[derive(Clone, Debug)]
pub struct MinCellQuality {
    /// Minimum number of unique genes per cell (or spot for Visium type
    /// technology).
    pub min_unique_genes: usize,
    /// Minimum library size per cell (or spot for Visium type technology).
    pub min_lib_size: usize,
    /// Minimum library size per cell (or spot for Visium type technology)/
    pub min_cells: usize,
    /// Target size for the library size normalisation (typical values 1e4 -
    /// 1e5 single cell.
    pub target_size: f32,
}

///////////////////////////
// Sparse data streaming //
///////////////////////////

///////////////////////
// CellOnFileQuality //
///////////////////////

/// CellOnFileQuality
///
/// This structure is being generate after a first scan of the file on disk and
/// defining which cells and genes to actually read in.
#[derive(Debug, Clone)]
pub struct CellOnFileQuality {
    /// Vector of indices of the cells to keep.
    pub cells_to_keep: Vec<usize>,
    /// Vector of indices of the genes to keep.
    pub genes_to_keep: Vec<usize>,
    /// HashSet of the indices to keep (for look-ups).
    pub cells_to_keep_set: FxHashSet<usize>,
    /// HashSet of the genes to keep (for look-ups).
    pub genes_to_keep_set: FxHashSet<usize>,
    /// Mapping of the old indices to the new indices for the cells.
    pub cell_old_to_new: FxHashMap<usize, usize>,
    /// Mapping of old indices to new indices for the genes.
    pub gene_old_to_new: FxHashMap<usize, usize>,
}

impl CellOnFileQuality {
    /// Create a new CellOnFileQuality structure
    ///
    /// ### Params
    ///
    /// * cells_to_keep - Index positions of the cells to keep
    /// * genes_to_keep - Index positions of the genes to keep
    ///
    /// ### Returns
    ///
    /// Initiliased self
    pub fn new(cells_to_keep: Vec<usize>, genes_to_keep: Vec<usize>) -> Self {
        Self {
            cells_to_keep,
            genes_to_keep,
            cells_to_keep_set: FxHashSet::default(),
            genes_to_keep_set: FxHashSet::default(),
            cell_old_to_new: FxHashMap::default(),
            gene_old_to_new: FxHashMap::default(),
        }
    }

    /// Generate internally the sets and maps
    pub fn generate_maps_sets(&mut self) {
        let cells_to_keep_set: FxHashSet<usize> = self.cells_to_keep.iter().cloned().collect();
        let genes_to_keep_set: FxHashSet<usize> = self.genes_to_keep.iter().cloned().collect();
        let cell_old_to_new: FxHashMap<usize, usize> = self
            .cells_to_keep
            .iter()
            .enumerate()
            .map(|(new_idx, &old_idx)| (old_idx, new_idx))
            .collect();
        let gene_old_to_new: FxHashMap<usize, usize> = self
            .genes_to_keep
            .iter()
            .enumerate()
            .map(|(new_idx, &old_idx)| (old_idx, new_idx))
            .collect();

        self.cells_to_keep_set = cells_to_keep_set;
        self.genes_to_keep_set = genes_to_keep_set;
        self.cell_old_to_new = cell_old_to_new;
        self.gene_old_to_new = gene_old_to_new
    }
}

//////////////////
// CsrCellChunk //
//////////////////

/// CsrCellChunk
///
/// This structure is designed to store the data of a single cell in a
/// CSR-like format optimised for rapid access on disk.
///
/// Raw counts are stored as a `RawCounts` enum which transparently handles
/// both u16 and u32 storage. Gene indices are stored as u32 to support
/// data sets with more than 65,535 features.
#[derive(Debug)]
pub struct CsrCellChunk {
    /// Raw counts for this cell, stored as either u16 or u32 depending on
    /// the data range
    pub data_raw: RawCounts,
    /// Vector of the norm counts of this cell. A lossy compression for f16 is
    /// applied.
    pub data_norm: Vec<F16>,
    /// Total library size/UMI counts of the cell.
    pub library_size: usize,
    /// Index positions of the genes (u32 to support >65k features)
    pub indices: Vec<u32>,
    /// Original index in the data
    pub original_index: usize,
    /// Flag if the cell should be kept. (Not used at the moment.)
    pub to_keep: bool,
}

impl CsrCellChunk {
    /// Function to generate the chunk from R data
    ///
    /// Assumes columns = genes, rows = cells.
    ///
    /// ### Params
    ///
    /// * `data` - The raw counts present in this cell
    /// * `col_idx` - The column indices where the gene is expressed.
    /// * `original_index` - Original row index in the matrix
    /// * `size_factor` - To which size to normalise to. 1e6 -> CPM normalisation.
    /// * `to_keep` - Whether the cell passes QC
    ///
    /// ### Returns
    ///
    /// The `CsrCellChunk` for this cell.
    pub fn from_data<T, U>(
        data: &[T],
        col_idx: &[U],
        original_index: usize,
        size_factor: f32,
        to_keep: bool,
    ) -> Self
    where
        T: FloatAndUInt,
        U: FloatAndUInt,
    {
        let data_f32 = data.iter().map(|&x| x.to_f32()).collect::<Vec<f32>>();
        let sum = data_f32.iter().sum::<f32>();
        let data_norm: Vec<F16> = data_f32
            .into_iter()
            .map(|x| {
                let norm = (x / sum * size_factor).ln_1p();
                F16::from(f16::from_f32(norm))
            })
            .collect();

        let raw_u32: Vec<u32> = data.iter().map(|&x| x.to_u32()).collect();
        let data_raw = RawCounts::from_u32_auto(&raw_u32);

        Self {
            data_raw,
            data_norm,
            library_size: sum as usize,
            indices: col_idx.iter().map(|&x| x.to_u32()).collect::<Vec<u32>>(),
            original_index,
            to_keep,
        }
    }

    /// Write directly to bytes on disk
    ///
    /// Binary layout (header = 32 bytes):
    ///
    /// ```text
    /// [0..4)   data_raw_len    u32
    /// [4..8)   data_norm_len   u32
    /// [8..12)  indices_len     u32
    /// [12..20) library_size    u64
    /// [20..28) original_index  u64
    /// [28]     to_keep         u8
    /// [29]     raw_elem_size   u8  (2 = u16, 4 = u32)
    /// [30..32) padding         2 bytes
    /// ```
    ///
    /// Followed by:
    /// - data_raw: `data_raw_len * raw_elem_size` bytes
    /// - data_norm: `data_norm_len * 2` bytes
    /// - indices: `indices_len * 4` bytes (u32)
    pub fn write_to_bytes(&self, writer: &mut impl Write) -> std::io::Result<()> {
        writer.write_all(&(self.data_raw.len() as u32).to_le_bytes())?;
        writer.write_all(&(self.data_norm.len() as u32).to_le_bytes())?;
        writer.write_all(&(self.indices.len() as u32).to_le_bytes())?;
        writer.write_all(&(self.library_size as u64).to_le_bytes())?;
        writer.write_all(&(self.original_index as u64).to_le_bytes())?;
        writer.write_all(&[self.to_keep as u8, self.data_raw.elem_size(), 0, 0])?;

        self.data_raw.write_bytes(writer)?;

        let mut data_norm_bytes = Vec::with_capacity(self.data_norm.len() * 2);
        for v in &self.data_norm {
            data_norm_bytes.extend_from_slice(&v.to_bits().to_le_bytes());
        }
        writer.write_all(&data_norm_bytes)?;

        let mut col_indices_bytes = Vec::with_capacity(self.indices.len() * 4);
        for i in &self.indices {
            col_indices_bytes.extend_from_slice(&i.to_le_bytes());
        }
        writer.write_all(&col_indices_bytes)?;

        Ok(())
    }

    /// Read data from buffer
    ///
    /// ### Params
    ///
    /// * `buffer` - A slice of u8's representing the buffer
    ///
    /// ### Return
    ///
    /// The `CsrCellChunk`
    pub fn read_from_buffer(buffer: &[u8]) -> Result<Self, BixverseErrors> {
        if buffer.len() < CSR_CHUNK_HEADER_LEN {
            return Err(BixverseErrors::ChunkBufferTooSmall {
                expected: CSR_CHUNK_HEADER_LEN,
                found: buffer.len(),
            });
        }

        let header = &buffer[0..CSR_CHUNK_HEADER_LEN];

        let data_raw_len =
            u32::from_le_bytes([header[0], header[1], header[2], header[3]]) as usize;
        let data_norm_len =
            u32::from_le_bytes([header[4], header[5], header[6], header[7]]) as usize;
        let col_indices_len =
            u32::from_le_bytes([header[8], header[9], header[10], header[11]]) as usize;
        let library_size = u64::from_le_bytes([
            header[12], header[13], header[14], header[15], header[16], header[17], header[18],
            header[19],
        ]) as usize;
        let original_index = u64::from_le_bytes([
            header[20], header[21], header[22], header[23], header[24], header[25], header[26],
            header[27],
        ]) as usize;
        let to_keep = header[28] != 0;
        let raw_elem_size = header[29];

        let data_start = CSR_CHUNK_HEADER_LEN;
        let elem_size = if raw_elem_size == RAW_ELEM_U32 {
            4usize
        } else {
            2usize
        };
        let data_end = data_start + data_raw_len * elem_size;
        let norm_end = data_end + data_norm_len * 2;
        let total_end = norm_end + col_indices_len * 4;

        // The three length fields are untrusted, so validate the total payload
        // before taking any slice.
        if buffer.len() < total_end {
            return Err(BixverseErrors::ChunkPayloadTruncated {
                expected: total_end,
                found: buffer.len(),
            });
        }

        let data_raw = RawCounts::read_from_buffer(
            &buffer[data_start..data_end],
            data_raw_len,
            raw_elem_size,
        )?;

        // Byte-wise reads rather than pointer casts: the buffer is an
        // lz4-decompressed `Vec<u8>`, so `norm_end` is not guaranteed to be
        // 4-byte aligned.
        let data_norm: Vec<F16> = buffer[data_end..norm_end]
            .chunks_exact(2)
            .map(|c| F16::from_le_bytes(c.try_into().expect("2-byte chunk by construction")))
            .collect();

        let indices: Vec<u32> = buffer[norm_end..total_end]
            .chunks_exact(4)
            .map(|c| u32::from_le_bytes(c.try_into().expect("4-byte chunk by construction")))
            .collect();

        Ok(Self {
            data_raw,
            data_norm,
            library_size,
            indices,
            original_index,
            to_keep,
        })
    }

    /// Generate a vector of Chunks from CompressedSparseData2
    ///
    /// ### Params
    ///
    /// * `sparse_data` - The `CompressedSparseData2` (in CSR format!)
    /// * `cell_qc` - Minimum QC thresholds
    ///
    /// ### Returns
    ///
    /// A tuple of the `Vec<CsrCellChunk>` and if the cell should be kept.
    pub fn generate_chunks_sparse_data<T, U>(
        sparse_data: CompressedSparseData2<T, U>,
        cell_qc: MinCellQuality,
    ) -> (Vec<CsrCellChunk>, CellQuality)
    where
        T: Clone + Default + Into<u32> + Sync,
        U: Clone + Default,
    {
        let n_cells = sparse_data.indptr.len() - 1;
        let n_genes = sparse_data.shape.1;

        // count how many cells express each gene
        let mut no_cells_exp_gene = vec![0usize; n_genes];
        for i in 0..n_cells {
            let start_i = sparse_data.indptr[i] as usize;
            let end_i = sparse_data.indptr[i + 1] as usize;
            for &gene_idx in &sparse_data.indices[start_i..end_i] {
                no_cells_exp_gene[gene_idx as usize] += 1;
            }
        }

        // filter genes
        let genes_to_keep: Vec<usize> = (0..n_genes)
            .filter(|&i| no_cells_exp_gene[i] >= cell_qc.min_cells)
            .collect();

        // create index mapping: old gene index -> new gene index
        let mut gene_index_map = vec![None; n_genes];
        for (new_idx, &old_idx) in genes_to_keep.iter().enumerate() {
            gene_index_map[old_idx] = Some(new_idx);
        }

        // Process cells
        let results: Vec<_> = (0..n_cells)
            .into_par_iter()
            .map(|i| {
                let start_i = sparse_data.indptr[i] as usize;
                let end_i = sparse_data.indptr[i + 1] as usize;

                let mut filtered_data = Vec::new();
                let mut filtered_indices = Vec::new();

                for idx in start_i..end_i {
                    let old_gene_idx = sparse_data.indices[idx] as usize;
                    if let Some(new_gene_idx) = gene_index_map[old_gene_idx] {
                        filtered_data.push(sparse_data.data[idx].clone().into());
                        filtered_indices.push(new_gene_idx);
                    }
                }

                let sum_data_i = filtered_data.iter().sum::<u32>() as usize;
                let nnz_i = filtered_indices.len();

                let to_keep_i =
                    (nnz_i >= cell_qc.min_unique_genes) && (sum_data_i >= cell_qc.min_lib_size);

                let chunk_i = CsrCellChunk::from_data(
                    &filtered_data,
                    &filtered_indices,
                    i,
                    cell_qc.target_size,
                    to_keep_i,
                );

                (chunk_i, sum_data_i, nnz_i, to_keep_i, i)
            })
            .collect();

        let mut res = Vec::with_capacity(n_cells);
        let mut lib_size = Vec::new();
        let mut nnz = Vec::new();
        let mut cells_to_keep = Vec::new();

        for (chunk, lib, n, keep, idx) in results {
            res.push(chunk);
            if keep {
                lib_size.push(lib);
                nnz.push(n);
                cells_to_keep.push(idx);
            }
        }

        let qc_data = CellQuality {
            cell_indices: cells_to_keep,
            gene_indices: genes_to_keep,
            lib_size,
            nnz,
        };

        (res, qc_data)
    }

    /// Helper function to get QC parameters for this cell
    ///
    /// ### Returns
    ///
    /// A tuple of `(no_genes, library_size)`
    pub fn get_qc_info(&self) -> (usize, usize) {
        (self.indices.len(), self.library_size)
    }

    /// Filter to keep only specified genes
    ///
    /// This will remove the genes that are not set to `true` in the boolean
    /// vector. This assumes that the indices of the genes and the booleans
    /// are the same!
    ///
    /// ### Params
    ///
    /// * `genes_to_keep` - Keep only the genes within this boolean vectors.
    pub fn filter_genes(&mut self, genes_to_keep: &[bool]) {
        let mut index_map = vec![None; genes_to_keep.len()];
        let mut new_idx = 0;
        for (old_idx, &keep) in genes_to_keep.iter().enumerate() {
            if keep {
                index_map[old_idx] = Some(new_idx);
                new_idx += 1;
            }
        }

        let mut new_indices = Vec::new();
        let mut new_norm = Vec::new();
        let mut keep_positions = Vec::new();

        for (i, &gene_idx) in self.indices.iter().enumerate() {
            if let Some(new_gene_idx) = index_map[gene_idx as usize] {
                new_indices.push(new_gene_idx);
                new_norm.push(self.data_norm[i]);
                keep_positions.push(i);
            }
        }

        let new_raw = match &self.data_raw {
            RawCounts::U16(v) => RawCounts::U16(keep_positions.iter().map(|&p| v[p]).collect()),
            RawCounts::U32(v) => RawCounts::U32(keep_positions.iter().map(|&p| v[p]).collect()),
        };

        self.indices = new_indices;
        self.data_raw = new_raw;
        self.data_norm = new_norm;
    }
}

//////////////////
// CscGeneChunk //
//////////////////

/// CscGeneChunk
///
/// This structure is designed to store the data of a single gene in a
/// CSC-like format optimised for rapid access on disk.
///
/// Raw counts are stored as a `RawCounts` enum (u16 or u32). Cell indices
/// are stored as u32.
#[derive(Encode, Decode, Serialize, Deserialize, Debug)]
pub struct CscGeneChunk {
    /// Raw counts per cell/spot for this gene, stored as either u16 or u32
    pub data_raw: RawCounts,
    /// Normalised counts per cell/spot for this gene. Lossy compression to
    /// f16 is applied.
    pub data_norm: Vec<F16>,
    /// Average expression of this gene in the data.
    pub avg_exp: F16,
    /// Number of cells expressing this gene
    pub nnz: usize,
    /// Indices of the cells expressing this gene
    pub indices: Vec<u32>,
    /// Original index from the data
    pub original_index: usize,
    /// Flag to indicate if gene shall be kept. Not in use at the moment.
    pub to_keep: bool,
}

impl CscGeneChunk {
    /// Helper function to generate the CscGeneChunk from converted data
    ///
    /// ### Params
    ///
    /// * `data_raw` - The raw counts for this gene
    /// * `data_norm` - The normalised counts for this gene.
    /// * `col_idx` - The column indices for which cells this gene is expressed.
    /// * `original_index` - Original row index
    /// * `to_keep` - Shall the gene be included in later analysis.
    ///
    /// ### Returns
    ///
    /// The `CscGeneChunk` for this gene.
    pub fn from_conversion(
        data_raw: RawCounts,
        data_norm: &[F16],
        col_idx: &[usize],
        original_index: usize,
        to_keep: bool,
    ) -> Self {
        // Accumulate in f32 and narrow once, as `pca.rs` already does. Summing
        // in f16 biases the total upwards: past 64 the f16 step is 0.0625, so
        // small addends round up every time and a thousand-cell gene overshoots
        // by roughly 5%. This value is written into the chunk header, so the
        // drift was being persisted.
        let avg_exp = F16::from_f32(data_norm.iter().map(|v| v.to_f32()).sum::<f32>());
        let nnz = data_raw.len();

        Self {
            data_raw,
            data_norm: data_norm.to_vec(),
            avg_exp,
            nnz,
            indices: col_idx.iter().map(|x| *x as u32).collect::<Vec<u32>>(),
            original_index,
            to_keep,
        }
    }

    /// Write directly to bytes on disk
    ///
    /// Binary layout (header = 36 bytes):
    ///
    /// ```text
    /// [0..4)   data_raw_len    u32
    /// [4..8)   data_norm_len   u32
    /// [8..12)  indices_len     u32
    /// [12..14) avg_exp         F16 (2 bytes)
    /// [14..16) padding         2 bytes
    /// [16..24) nnz             u64
    /// [24..32) original_index  u64
    /// [32]     to_keep         u8
    /// [33]     raw_elem_size   u8  (2 = u16, 4 = u32)
    /// [34..36) padding         2 bytes
    /// ```
    ///
    /// Followed by:
    /// - data_raw: `data_raw_len * raw_elem_size` bytes
    /// - data_norm: `data_norm_len * 2` bytes
    /// - indices: `indices_len * 4` bytes (u32)
    pub fn write_to_bytes(&self, writer: &mut impl Write) -> std::io::Result<()> {
        writer.write_all(&(self.data_raw.len() as u32).to_le_bytes())?;
        writer.write_all(&(self.data_norm.len() as u32).to_le_bytes())?;
        writer.write_all(&(self.indices.len() as u32).to_le_bytes())?;
        writer.write_all(&self.avg_exp.to_le_bytes())?;
        // first padding
        writer.write_all(&[0, 0])?;
        writer.write_all(&(self.nnz as u64).to_le_bytes())?;
        writer.write_all(&(self.original_index as u64).to_le_bytes())?;
        writer.write_all(&[self.to_keep as u8, self.data_raw.elem_size(), 0, 0])?;

        self.data_raw.write_bytes(writer)?;

        let mut data_norm_bytes = Vec::with_capacity(self.data_norm.len() * 2);
        for v in &self.data_norm {
            data_norm_bytes.extend_from_slice(&v.to_bits().to_le_bytes());
        }
        writer.write_all(&data_norm_bytes)?;

        let mut row_indices_bytes = Vec::with_capacity(self.indices.len() * 4);
        for i in &self.indices {
            row_indices_bytes.extend_from_slice(&i.to_le_bytes());
        }
        writer.write_all(&row_indices_bytes)?;

        Ok(())
    }

    /// Read data from buffer
    ///
    /// ### Params
    ///
    /// * `buffer` - A slice of u8's representing the buffer
    ///
    /// ### Return
    ///
    /// The `CscGeneChunk`
    pub fn read_from_buffer(buffer: &[u8]) -> Result<Self, BixverseErrors> {
        if buffer.len() < CSC_CHUNK_HEADER_LEN {
            return Err(BixverseErrors::ChunkBufferTooSmall {
                expected: CSC_CHUNK_HEADER_LEN,
                found: buffer.len(),
            });
        }

        let header = &buffer[0..CSC_CHUNK_HEADER_LEN];
        let data_raw_len =
            u32::from_le_bytes([header[0], header[1], header[2], header[3]]) as usize;
        let data_norm_len =
            u32::from_le_bytes([header[4], header[5], header[6], header[7]]) as usize;
        let row_indices_len =
            u32::from_le_bytes([header[8], header[9], header[10], header[11]]) as usize;
        let avg_exp = F16::from_le_bytes([header[12], header[13]]);
        // skip 2 bytes padding
        let nnz = u64::from_le_bytes([
            header[16], header[17], header[18], header[19], header[20], header[21], header[22],
            header[23],
        ]) as usize;
        let original_index = u64::from_le_bytes([
            header[24], header[25], header[26], header[27], header[28], header[29], header[30],
            header[31],
        ]) as usize;
        let to_keep = header[32] != 0;
        let raw_elem_size = header[33];

        let data_start = CSC_CHUNK_HEADER_LEN;
        let elem_size = if raw_elem_size == RAW_ELEM_U32 {
            4usize
        } else {
            2usize
        };
        let data_end = data_start + data_raw_len * elem_size;
        let norm_end = data_end + data_norm_len * 2;
        let total_end = norm_end + row_indices_len * 4;

        // The three length fields are untrusted, so validate the total payload
        // before taking any slice.
        if buffer.len() < total_end {
            return Err(BixverseErrors::ChunkPayloadTruncated {
                expected: total_end,
                found: buffer.len(),
            });
        }

        let data_raw = RawCounts::read_from_buffer(
            &buffer[data_start..data_end],
            data_raw_len,
            raw_elem_size,
        )?;

        // Byte-wise reads rather than pointer casts: the buffer is an
        // lz4-decompressed `Vec<u8>`, so `norm_end` is not guaranteed to be
        // 4-byte aligned.
        let data_norm: Vec<F16> = buffer[data_end..norm_end]
            .chunks_exact(2)
            .map(|c| F16::from_le_bytes(c.try_into().expect("2-byte chunk by construction")))
            .collect();

        let indices: Vec<u32> = buffer[norm_end..total_end]
            .chunks_exact(4)
            .map(|c| u32::from_le_bytes(c.try_into().expect("4-byte chunk by construction")))
            .collect();

        Ok(Self {
            data_raw,
            data_norm,
            avg_exp,
            nnz,
            indices,
            original_index,
            to_keep,
        })
    }

    /// Helper function that allows to filter out cells
    ///
    /// ### Params
    ///
    /// * `cells_to_keep` - IndexSet with cell index positions to keep.
    pub fn filter_selected_cells(&mut self, cells_to_keep: &IndexSet<u32>) {
        // build reverse mapping: cell_id -> position in this gene's data
        let cell_positions: FxHashMap<u32, usize> = self
            .indices
            .iter()
            .enumerate()
            .map(|(pos, &cell_id)| (cell_id, pos))
            .collect();

        let mut new_data_norm = Vec::with_capacity(cells_to_keep.len());
        let mut new_row_indices = Vec::with_capacity(cells_to_keep.len());
        let mut keep_positions = Vec::new();

        // iterate in cells_to_keep order (critical for PCA!)
        for (new_row_idx, &cell_index) in cells_to_keep.iter().enumerate() {
            if let Some(&pos) = cell_positions.get(&cell_index) {
                keep_positions.push(pos);
                new_data_norm.push(self.data_norm[pos]);
                new_row_indices.push(new_row_idx as u32);
            }
        }

        let new_raw = match &self.data_raw {
            RawCounts::U16(v) => RawCounts::U16(keep_positions.iter().map(|&p| v[p]).collect()),
            RawCounts::U32(v) => RawCounts::U32(keep_positions.iter().map(|&p| v[p]).collect()),
        };

        let nnz = new_row_indices.len();
        self.data_raw = new_raw;
        self.data_norm = new_data_norm;
        self.indices = new_row_indices;
        self.nnz = nnz;
    }

    /// Calculate the average gene expression given a set of cells
    ///
    /// ### Params
    ///
    /// * `cells_to_keep` - IndexSet with cell index positions to keep.
    ///
    /// ### Returns
    ///
    /// Tuple of `(gene_index, average expression)`.
    pub fn calculate_avg_exp(&self, cells_to_keep: &IndexSet<u32>) -> (usize, f32) {
        let sum: f32 = self
            .indices
            .iter()
            .zip(self.data_norm.iter())
            .filter(|(cell_idx, _)| cells_to_keep.contains(*cell_idx))
            .map(|(_, val)| val.to_f32())
            .sum();
        let avg = sum / cells_to_keep.len() as f32;

        (self.original_index, avg)
    }

    /// Transform the chunk to a sparse Axis of CSC type
    ///
    /// ### Params
    ///
    /// * `n_cells` - Number of cells represented in the data
    ///
    /// ### Returns
    ///
    /// `SparseAxis` with u32 in the main slot and f32 in the data_2 layer.
    pub fn to_sparse_axis(&self, n_cells: usize) -> SparseAxis<u32, f32> {
        let raw_u32: Vec<u32> = self.data_raw.iter().collect();
        SparseAxis::new_csc(
            self.indices.iter().map(|x| *x as usize).collect(),
            raw_u32,
            Some(self.data_norm.iter().map(|x| x.to_f32()).collect()),
            n_cells,
        )
    }
}

/// SparseDataHeader
///
/// Stores the information in terms of total cells, total genes, number of
/// chunks in terms of cells and genes and the offset vectors
#[derive(Encode, Decode, Serialize, Deserialize, Clone)]
pub struct SparseDataHeader {
    /// Total number of cells in the experiment.
    pub total_cells: usize,
    /// Total number of genes in the experiemnt.
    pub total_genes: usize,
    /// Is the file written in a way for fast cell data retrieval (set to true)
    /// or for fast gene retrieval (set to false)
    pub cell_based: bool,
    /// Number of chunks in this file storing either the cell or gene data
    pub no_chunks: usize,
    /// Offset vector for reading in the data
    pub chunk_offsets: Vec<u64>,
    /// FxHashMap with the original index -> chunk info
    pub index_map: FxHashMap<usize, usize>,
}

/// Leading bytes identifying a bixverse single-cell binary. Checked on open so
/// a foreign file is refused rather than parsed as one of ours.
const FILE_MAGIC: &[u8; 8] = b"SCRNASEQ";

/// Fixed-size file header that points to the main header location
#[repr(C)]
#[derive(Encode, Decode, Serialize, Deserialize)]
struct FileHeader {
    /// Magic string as bytes to recognise the file
    magic: [u8; 8],
    /// Version of the file
    version: u32,
    /// Offset of the main header, i.e., 64 bytes
    main_header_offset: u64,
    /// Is cell-based (CSR format)
    cell_based: bool,
    /// Library size the `data_norm` layer was normalised against. Carved out
    /// of the reserved block, so files predating this field decode it as
    /// `0.0`, which is the "unknown" sentinel (a real target size is never 0).
    target_size: f32,
    /// 28 additional reserved bytes for the future
    _reserved_1: [u8; 28],
    /// 3 additional reserved bytes for the future
    _reserved_2: [u8; 3],
}

impl FileHeader {
    /// Generate a new header
    ///
    /// ### Params
    ///
    /// * `cell_based` - Is the data stored for fast cell retrieval.
    /// * `target_size` - Library size the normalised layer was scaled to. Pass
    ///   `0.0` on raw-only paths where no normalisation was applied.
    ///
    /// ### Returns
    ///
    /// A new object of `FileHeader`
    fn new(cell_based: bool, target_size: f32) -> Self {
        Self {
            magic: *FILE_MAGIC,
            version: SC_FILE_VERSION,
            main_header_offset: 0,
            cell_based,
            target_size,
            _reserved_1: [0; 28],
            _reserved_2: [0; 3],
        }
    }
}

//////////////////////
// Streaming writer //
//////////////////////

/// CellGeneSparseWriter
///
/// Implementation of a structure for writing in a streamed manner two different
/// types of sparse stored data.
pub struct CellGeneSparseWriter {
    /// The header of the file.
    header: SparseDataHeader,
    /// Buffered writer to the file.
    writer: BufWriter<File>,
    /// The current position of the chunks.
    chunks_start_pos: u64,
    /// Boolean indicating if the writer is designed to write in an efficient
    /// manner for cells.
    cell_based: bool,
    /// The chunks written since the last flush
    chunks_since_flush: usize,
    /// The frequency of flushing, see [`CELL_FLUSH_FREQUENCY`] and
    /// [`GENE_FLUSH_FREQUENCY`].
    flush_frequency: usize,
    /// Library size the `data_norm` layer was normalised against. Written into
    /// the [`FileHeader`] on finalisation, which rebuilds the header from
    /// scratch, so the value has to survive on the writer.
    target_size: f32,
}

impl CellGeneSparseWriter {
    /// Create a new sparse writer instance
    ///
    /// This writer assumes that rows represent genes and columns represent
    /// cells.
    ///
    /// ### Params
    ///
    /// * `path_f` - Path to the .bin file to which to write to.
    /// * `cell_based` - Shall the writer be set up for writing cell-based
    ///   (`true`) or gene-based chunks.
    /// * `total_cells` - Total cells in the data.
    /// * `total_genes` - Total genes in the data.
    /// * `target_size` - Library size the `data_norm` layer was normalised
    ///   against. Pass `0.0` on raw-only paths where no normalisation applies;
    ///   readers report that as `None`.
    ///
    /// ### Returns
    ///
    /// The `CellGeneSparseWriter`.
    pub fn new<P: AsRef<Path>>(
        path_f: P,
        cell_based: bool,
        total_cells: usize,
        total_genes: usize,
        target_size: f32,
    ) -> Result<Self, BixverseErrors> {
        let file = File::create(path_f)?;
        let mut writer = BufWriter::with_capacity(128 * 1024 * 1024, file);

        let file_header = FileHeader::new(cell_based, target_size);
        let file_header_enc = encode_to_vec(&file_header, config::standard())
            .map_err(|_| BixverseErrors::HeaderEncodeFailed)?;
        if file_header_enc.len() < 64 {
            writer.write_all(&file_header_enc)?;
            writer.write_all(&vec![0u8; 64 - file_header_enc.len()])?;
        } else {
            writer.write_all(&file_header_enc[..64])?;
        }
        writer.flush()?;

        let chunks_start_pos = 64;

        let flush_frequency = if cell_based {
            CELL_FLUSH_FREQUENCY
        } else {
            GENE_FLUSH_FREQUENCY
        };

        let header = SparseDataHeader {
            total_cells,
            total_genes,
            cell_based,
            no_chunks: 0,
            chunk_offsets: Vec::new(),
            index_map: FxHashMap::default(),
        };

        Ok(Self {
            header,
            writer,
            chunks_start_pos,
            cell_based,
            chunks_since_flush: 0_usize,
            flush_frequency,
            target_size,
        })
    }

    /// Verify the writer was opened in the orientation the caller expects.
    ///
    /// ### Params
    ///
    /// * `want_cell_based` - `true` when a cell (CSR) chunk is being written.
    ///
    /// ### Returns
    ///
    /// `Ok(())`, or [`BixverseErrors::ReaderModeMismatch`] on a mismatch.
    fn check_mode(&self, want_cell_based: bool) -> Result<(), BixverseErrors> {
        if self.cell_based == want_cell_based {
            return Ok(());
        }
        let label = |b: bool| if b { "cell-based" } else { "gene-based" };
        Err(BixverseErrors::ReaderModeMismatch {
            actual: label(self.cell_based),
            requested: label(want_cell_based),
        })
    }

    /// Compress a serialised chunk, append it, and record its offset.
    ///
    /// Shared tail of [`Self::write_cell_chunk`] and [`Self::write_gene_chunk`].
    /// Flushes the underlying [`BufWriter`] every `flush_frequency` chunks.
    ///
    /// ### Params
    ///
    /// * `buffer` - The chunk serialised via its `write_to_bytes`.
    /// * `original_index` - Original index of the cell or gene in the data.
    ///
    /// ### Returns
    ///
    /// `Ok(())` on success, an I/O error otherwise.
    fn append_chunk(&mut self, buffer: &[u8], original_index: usize) -> Result<(), BixverseErrors> {
        let current_pos = self.writer.stream_position()?;
        let chunk_offset = current_pos - self.chunks_start_pos;
        self.header.chunk_offsets.push(chunk_offset);
        self.header
            .index_map
            .insert(original_index, self.header.no_chunks);

        let compressed = compress_prepend_size(buffer);

        self.writer
            .write_all(&(compressed.len() as u64).to_le_bytes())?;
        self.writer.write_all(&compressed)?;

        self.header.no_chunks += 1;
        self.chunks_since_flush += 1;

        if self.chunks_since_flush >= self.flush_frequency {
            self.writer.flush()?;
            self.chunks_since_flush = 0;
        }

        Ok(())
    }

    /// Write a Cell (Chunk) to the file
    ///
    /// The data is represented in a CSR-type format.
    ///
    /// ### Params
    ///
    /// * `cell_chunk` - The data representing that specific cell.
    ///
    /// ### Returns
    ///
    /// `Ok(())`, or [`BixverseErrors::ReaderModeMismatch`] if the writer was
    /// opened for gene-based chunks.
    pub fn write_cell_chunk(&mut self, cell_chunk: CsrCellChunk) -> Result<(), BixverseErrors> {
        self.check_mode(true)?;

        let mut buffer = Vec::new();
        cell_chunk.write_to_bytes(&mut buffer)?;

        self.append_chunk(&buffer, cell_chunk.original_index)
    }

    /// Write a Gene to the file
    ///
    /// ### Params
    ///
    /// * `gene_chunk` - The data representing that specific gene.
    ///
    /// ### Returns
    ///
    /// `Ok(())`, or [`BixverseErrors::ReaderModeMismatch`] if the writer was
    /// opened for cell-based chunks.
    pub fn write_gene_chunk(&mut self, gene_chunk: CscGeneChunk) -> Result<(), BixverseErrors> {
        self.check_mode(false)?;

        let mut buffer = Vec::new();
        gene_chunk.write_to_bytes(&mut buffer)?;

        self.append_chunk(&buffer, gene_chunk.original_index)
    }

    /// Finalise the file
    ///
    /// Writes the [`SparseDataHeader`] tail, then rewrites the fixed 64-byte
    /// file header at offset zero with the resolved main header offset.
    ///
    /// ### Returns
    ///
    /// `Ok(())` on success, [`BixverseErrors::HeaderEncodeFailed`] if either
    /// header could not be encoded.
    pub fn finalise(mut self) -> Result<(), BixverseErrors> {
        // write header size and header
        let main_header_offset = self.writer.stream_position()?;

        let header_data = encode_to_vec(&self.header, config::standard())
            .map_err(|_| BixverseErrors::HeaderEncodeFailed)?;
        let header_size = header_data.len() as u64;

        self.writer.write_all(&header_size.to_le_bytes())?;
        self.writer.write_all(&header_data)?;

        self.writer.seek(SeekFrom::Start(0))?;
        let mut file_header = FileHeader::new(self.cell_based, self.target_size);
        file_header.main_header_offset = main_header_offset;
        let file_header_enc = encode_to_vec(&file_header, config::standard())
            .map_err(|_| BixverseErrors::HeaderEncodeFailed)?;

        // ensure it's exactly 64 bytes
        if file_header_enc.len() < 64 {
            self.writer.write_all(&file_header_enc)?;
            self.writer
                .write_all(&vec![0u8; 64 - file_header_enc.len()])?;
        } else {
            self.writer.write_all(&file_header_enc[..64])?;
        }

        self.writer.flush()?;

        Ok(())
    }

    /// Update the number of cells in the header
    ///
    /// Helper function to update the number of cells to the number that was
    /// finally written on disk.
    ///
    /// ### Params
    ///
    /// * `no_cells` - New number of cells
    pub fn update_header_no_cells(&mut self, no_cells: usize) {
        self.header.total_cells = no_cells;
    }

    /// Update the number of cells in the header
    ///
    /// Helper function to update the number of genes to the number that was
    /// finally written on disk.
    ///
    /// ### Params
    ///
    /// * `no_genes` - New number of genes
    pub fn update_header_no_genes(&mut self, no_genes: usize) {
        self.header.total_genes = no_genes;
    }
}

///////////////////////
// SingleCellReading //
///////////////////////

/// Read access to a chunked single cell count store.
///
/// Abstracts over where the data physically lives. The only implementor today
/// is [`ParallelSparseReader`], which memory-maps the bixverse binary format
/// from disk, but the trait exposes nothing file-specific so that other
/// backends can be added later without touching consumers.
///
/// Whether a store is cell-based (CSR, fast per-cell access) or gene-based
/// (CSC, fast per-gene access) is a runtime property of the underlying data,
/// not of the type. Asking a gene-based store for cells returns
/// [`BixverseErrors::ReaderModeMismatch`].
///
/// Implementors need only supply `read_cells_parallel`, `read_gene_parallel`,
/// `get_header` and `is_cell_based`. The rest have default bodies written in
/// terms of those four. Backends that can do better should override: the
/// defaults for `read_cell_library_sizes` and `read_gene_nnz` parse whole
/// chunks, whereas a backend able to peek at chunk headers can avoid that.
///
/// `Send + Sync` is required: readers are shared across `rayon` threads.
pub trait SingleCellReading: Send + Sync {
    /// Read cells by index in a multi-threaded manner
    ///
    /// ### Params
    ///
    /// * `indices` - Index positions of the cells to retrieve
    ///
    /// ### Returns
    ///
    /// The `CsrCellChunk`s in the order given by `indices`.
    fn read_cells_parallel(&self, indices: &[usize]) -> Result<Vec<CsrCellChunk>, BixverseErrors>;

    /// Read genes by index in a multi-threaded manner
    ///
    /// ### Params
    ///
    /// * `indices` - Index positions of the genes to retrieve
    ///
    /// ### Returns
    ///
    /// The `CscGeneChunk`s in the order given by `indices`.
    fn read_gene_parallel(&self, indices: &[usize]) -> Result<Vec<CscGeneChunk>, BixverseErrors>;

    /// Return the header of the underlying store
    ///
    /// ### Returns
    ///
    /// The `SparseDataHeader`.
    fn get_header(&self) -> SparseDataHeader;

    /// Is the store laid out for fast cell retrieval?
    ///
    /// ### Returns
    ///
    /// `true` for a cell-based (CSR) store.
    fn is_cell_based(&self) -> bool;

    /// Is the store laid out for fast gene retrieval?
    ///
    /// ### Returns
    ///
    /// `true` for a gene-based (CSC) store.
    fn is_gene_based(&self) -> bool {
        !self.is_cell_based()
    }

    /// Library size the `data_norm` layer of this store was scaled to.
    ///
    /// Needed by anything that has to undo the `ln1p(count / lib_size *
    /// target_size)` normalisation, such as `CscGeneChunk::transform_to_clr`.
    /// Defaults to `None` for stores that do not record it.
    ///
    /// ### Returns
    ///
    /// `Some(target_size)`, or `None` when the store cannot report one.
    fn target_size(&self) -> Option<f32> {
        None
    }

    /// Read a single cell by index
    ///
    /// ### Params
    ///
    /// * `index` - Cell index
    ///
    /// ### Returns
    ///
    /// The `CsrCellChunk` of this cell.
    fn read_cell(&self, index: usize) -> Result<CsrCellChunk, BixverseErrors> {
        self.read_cells_parallel(&[index])?
            .into_iter()
            .next()
            .ok_or(BixverseErrors::ChunkIndexNotFound(index))
    }

    /// Read a single gene by index
    ///
    /// ### Params
    ///
    /// * `index` - Gene index
    ///
    /// ### Returns
    ///
    /// The `CscGeneChunk` of this gene.
    fn read_gene(&self, index: usize) -> Result<CscGeneChunk, BixverseErrors> {
        self.read_gene_parallel(&[index])?
            .into_iter()
            .next()
            .ok_or(BixverseErrors::ChunkIndexNotFound(index))
    }

    /// Read a contiguous range of cells
    ///
    /// Helper for memory-bounded gene generation.
    ///
    /// ### Params
    ///
    /// * `start` - First cell index (inclusive)
    /// * `end` - Last cell index (exclusive)
    ///
    /// ### Returns
    ///
    /// The `CsrCellChunk`s in that range.
    fn read_cells_range(
        &self,
        start: usize,
        end: usize,
    ) -> Result<Vec<CsrCellChunk>, BixverseErrors> {
        let indices: Vec<usize> = (start..end).collect();
        self.read_cells_parallel(&indices)
    }

    /// Return all cells in the store
    ///
    /// ### Returns
    ///
    /// Every `CsrCellChunk`.
    fn get_all_cells(&self) -> Result<Vec<CsrCellChunk>, BixverseErrors> {
        let indices: Vec<usize> = (0..self.get_header().total_cells).collect();
        self.read_cells_parallel(&indices)
    }

    /// Return all genes in the store
    ///
    /// ### Returns
    ///
    /// Every `CscGeneChunk`.
    fn get_all_genes(&self) -> Result<Vec<CscGeneChunk>, BixverseErrors> {
        let indices: Vec<usize> = (0..self.get_header().total_genes).collect();
        self.read_gene_parallel(&indices)
    }

    /// Read genes by index and filter them to a set of cells
    ///
    /// ### Params
    ///
    /// * `indices` - Index positions of the genes to retrieve
    /// * `cells_to_keep` - Cell index positions to keep, in output order
    ///
    /// ### Returns
    ///
    /// The filtered `CscGeneChunk`s.
    fn read_gene_parallel_filtered(
        &self,
        indices: &[usize],
        cells_to_keep: &IndexSet<u32>,
    ) -> Result<Vec<CscGeneChunk>, BixverseErrors> {
        let mut chunks = self.read_gene_parallel(indices)?;
        chunks
            .par_iter_mut()
            .for_each(|chunk| chunk.filter_selected_cells(cells_to_keep));
        Ok(chunks)
    }

    /// Read only the library sizes for the given cells
    ///
    /// ### Params
    ///
    /// * `indices` - Cell indices
    ///
    /// ### Returns
    ///
    /// Vector of library sizes in the order given by `indices`.
    fn read_cell_library_sizes(&self, indices: &[usize]) -> Result<Vec<usize>, BixverseErrors> {
        Ok(self
            .read_cells_parallel(indices)?
            .iter()
            .map(|chunk| chunk.library_size)
            .collect())
    }

    /// Read the NNZ for the given genes
    ///
    /// ### Params
    ///
    /// * `indices` - Gene indices
    ///
    /// ### Returns
    ///
    /// Vector of NNZ counts in the order given by `indices`.
    fn read_gene_nnz(&self, indices: &[usize]) -> Result<Vec<usize>, BixverseErrors> {
        Ok(self
            .read_gene_parallel(indices)?
            .iter()
            .map(|chunk| chunk.nnz)
            .collect())
    }

    /// Return NNZ for all genes
    ///
    /// ### Returns
    ///
    /// Vector of NNZ values for all genes.
    fn get_all_gene_nnz(&self) -> Result<Vec<usize>, BixverseErrors> {
        let indices: Vec<usize> = (0..self.get_header().total_genes).collect();
        self.read_gene_nnz(&indices)
    }

    /// Compute the shifted CLR per-cell offsets freshly from raw counts.
    ///
    /// The offset for cell i is `m_i = (1/D) * sum_j log1p(x_ij / s_i)`, where
    /// D is the total gene count and s_i is the cell's library size. Size
    /// factor is fixed at 1, independent of any size factor used during the
    /// original normalisation. Cells are read in batches to bound peak memory.
    ///
    /// ### Params
    ///
    /// * `indices` - Cell indices to compute offsets for
    /// * `batch_size` - Number of cells to load per batch (defaults to
    ///   [`CELL_BATCH_SIZE`])
    ///
    /// ### Returns
    ///
    /// Vector of CLR offsets in the same order as `indices`.
    fn get_clr_offsets(
        &self,
        indices: &[usize],
        batch_size: Option<usize>,
    ) -> Result<Vec<f64>, BixverseErrors> {
        if !self.is_cell_based() {
            return Err(BixverseErrors::ReaderModeMismatch {
                actual: "gene-based",
                requested: "cell-based",
            });
        }
        let batch_size = batch_size.unwrap_or(CELL_BATCH_SIZE);

        let n_genes = self.get_header().total_genes as f64;
        let mut offsets = Vec::with_capacity(indices.len());

        for batch in indices.chunks(batch_size) {
            let cells = self.read_cells_parallel(batch)?;
            let batch_offsets: Vec<f64> = cells
                .par_iter()
                .map(|chunk| {
                    let inv_s = 1.0_f64 / chunk.library_size as f64;
                    let sum: f64 = chunk
                        .data_raw
                        .iter()
                        .map(|c| (c as f64 * inv_s).ln_1p())
                        .sum();
                    sum / n_genes
                })
                .collect();
            offsets.extend_from_slice(&batch_offsets);
        }

        Ok(offsets)
    }
}

//////////////////////
// Streaming reader //
//////////////////////

/// Read the `target_size` out of a bin file without mapping the whole file.
///
/// Only the fixed 64-byte file header is touched, so this is cheap enough to
/// run over every input of a merge before deciding what to write.
///
/// ### Params
///
/// * `f_path` - Path to the .bin file.
///
/// ### Returns
///
/// `Some(target_size)`, or `None` for raw-only files and files written before
/// the field entered the header.
pub fn peek_target_size<P: AsRef<Path>>(f_path: P) -> Result<Option<f32>, BixverseErrors> {
    use std::io::Read;

    let mut buffer = [0u8; 64];
    File::open(f_path)?.read_exact(&mut buffer)?;

    let (file_header, _) = decode_from_slice::<FileHeader, _>(&buffer, config::standard())
        .map_err(|_| BixverseErrors::HeaderDecodeFailed)?;

    Ok((file_header.target_size != 0.0).then_some(file_header.target_size))
}

/// ParallelSparseReader
pub struct ParallelSparseReader {
    /// The file header
    header: SparseDataHeader,
    /// Reference to the memory map for safe sharing across threads
    mmap: Arc<memmap2::Mmap>,
    /// Start position of the chunks after the header file
    chunks_start: u64,
    /// Library size the `data_norm` layer was normalised against, as recorded
    /// in the [`FileHeader`]. `0.0` for files written before the field existed.
    target_size: f32,
}

impl ParallelSparseReader {
    /// Generate a new parallelised streaming reader
    ///
    /// ### Params
    ///
    /// * `f_path` - Path to the file
    ///
    /// ### Returns
    ///
    /// Initialised `ParallelSparseReader`.
    pub fn new(f_path: &str) -> Result<Self, BixverseErrors> {
        let file = File::open(f_path)?;
        let file_size = file.metadata()?.len();

        let mmap = unsafe {
            let mut opts = MmapOptions::new();
            if file_size <= 8 * 1024 * 1024 * 1024 {
                opts.populate();
            }
            opts.map(&file)?
        };

        #[cfg(unix)]
        mmap.advise(memmap2::Advice::Random)?;

        // Everything below indexes the mmap directly, so each offset is checked
        // before it is used. A corrupt or foreign file must return an error,
        // not panic on a slice bound or parse as if it were ours.
        if mmap.len() < 64 {
            return Err(BixverseErrors::HeaderDecodeFailed);
        }

        let file_header_bytes = &mmap[0..64];
        let (file_header, _) =
            decode_from_slice::<FileHeader, _>(file_header_bytes, config::standard())
                .map_err(|_| BixverseErrors::HeaderDecodeFailed)?;

        if &file_header.magic != FILE_MAGIC {
            return Err(BixverseErrors::HeaderDecodeFailed);
        }

        if file_header.version != SC_FILE_VERSION {
            return Err(BixverseErrors::FileVersionMismatch {
                expected: SC_FILE_VERSION,
                found: file_header.version,
            });
        }

        let main_header_offset = file_header.main_header_offset as usize;
        if main_header_offset.saturating_add(8) > mmap.len() {
            return Err(BixverseErrors::HeaderDecodeFailed);
        }

        let header_size = u64::from_le_bytes(
            mmap[main_header_offset..main_header_offset + 8]
                .try_into()
                .expect("8-byte slice, bounds checked above"),
        ) as usize;

        let header_end = main_header_offset
            .saturating_add(8)
            .saturating_add(header_size);
        if header_end > mmap.len() {
            return Err(BixverseErrors::HeaderDecodeFailed);
        }

        let header_bytes = &mmap[main_header_offset + 8..header_end];
        let (header, _) =
            decode_from_slice::<SparseDataHeader, _>(header_bytes, config::standard())
                .map_err(|_| BixverseErrors::HeaderDecodeFailed)?;

        Ok(Self {
            header,
            mmap: Arc::new(mmap),
            chunks_start: 64,
            target_size: file_header.target_size,
        })
    }

    /// Locate, decompress and return the raw bytes of a single chunk
    ///
    /// Shared by every read path: resolve the original index to a chunk offset
    /// via the header index map, read the length prefix, then lz4 decompress.
    ///
    /// ### Params
    ///
    /// * `original_index` - Original index of the cell or gene in the data
    ///
    /// ### Returns
    ///
    /// The decompressed chunk bytes.
    fn decompress_chunk(&self, original_index: usize) -> Result<Vec<u8>, BixverseErrors> {
        let chunk_index = *self
            .header
            .index_map
            .get(&original_index)
            .ok_or(BixverseErrors::ChunkIndexNotFound(original_index))?;
        let chunk_offset = (self.chunks_start + self.header.chunk_offsets[chunk_index]) as usize;

        let compressed_size = u64::from_le_bytes(
            self.mmap[chunk_offset..chunk_offset + 8]
                .try_into()
                .expect("8-byte slice by construction"),
        ) as usize;

        let compressed = &self.mmap[chunk_offset + 8..chunk_offset + 8 + compressed_size];

        decompress_size_prepended(compressed)
            .map_err(|_| BixverseErrors::ChunkDecompressionFailed(chunk_offset as u64))
    }
}

//////////////////////////
// Trait implementation //
//////////////////////////

impl SingleCellReading for ParallelSparseReader {
    /// Read in cells by indices in a multi-threaded manner
    ///
    /// ### Params
    ///
    /// * `indices` - Slice of index positions of the cells to retrieve
    ///
    /// ### Returns
    ///
    /// Returns an array of `CsrCellChunk`.
    fn read_cells_parallel(&self, indices: &[usize]) -> Result<Vec<CsrCellChunk>, BixverseErrors> {
        if !self.is_cell_based() {
            return Err(BixverseErrors::ReaderModeMismatch {
                actual: "gene-based",
                requested: "cell-based",
            });
        }

        indices
            .par_iter()
            .map(|&original_index| {
                let decompressed = self.decompress_chunk(original_index)?;

                CsrCellChunk::read_from_buffer(&decompressed)
            })
            .collect()
    }

    /// Read in genes by indices in a multi-threaded manner
    ///
    /// ### Params
    ///
    /// * `indices` - Slice of index positions of the genes to retrieve
    ///
    /// ### Returns
    ///
    /// Returns an array of `CscGeneChunk`.
    fn read_gene_parallel(&self, indices: &[usize]) -> Result<Vec<CscGeneChunk>, BixverseErrors> {
        if !self.is_gene_based() {
            return Err(BixverseErrors::ReaderModeMismatch {
                actual: "cell-based",
                requested: "gene-based",
            });
        }

        indices
            .par_iter()
            .map(|&original_index| {
                let decompressed = self.decompress_chunk(original_index)?;

                CscGeneChunk::read_from_buffer(&decompressed)
            })
            .collect()
    }

    /// Read in genes by indices in a multi-threaded manner
    ///
    /// And filter the genes by `cells_to_keep` in parallel. Overrides the
    /// default so that reading and filtering are fused into a single pass.
    ///
    /// ### Params
    ///
    /// * `indices` - Slice of index positions of the genes to retrieve
    /// * `cells_to_keep` - Cell index positions to keep, in output order
    ///
    /// ### Returns
    ///
    /// Returns an array of `CscGeneChunk`.
    fn read_gene_parallel_filtered(
        &self,
        indices: &[usize],
        cells_to_keep: &IndexSet<u32>,
    ) -> Result<Vec<CscGeneChunk>, BixverseErrors> {
        if !self.is_gene_based() {
            return Err(BixverseErrors::ReaderModeMismatch {
                actual: "cell-based",
                requested: "gene-based",
            });
        }

        indices
            .par_iter()
            .map(|&original_index| {
                let decompressed = self.decompress_chunk(original_index)?;

                let mut chunk = CscGeneChunk::read_from_buffer(&decompressed)?;

                chunk.filter_selected_cells(cells_to_keep);

                Ok(chunk)
            })
            .collect()
    }

    /// Helper to return the header
    ///
    /// ### Returns
    ///
    /// Returns the header file
    fn get_header(&self) -> SparseDataHeader {
        self.header.clone()
    }

    /// Read only library sizes for specified cells
    ///
    /// More efficient than the default, which parses whole chunks. Reads
    /// directly from the chunk header bytes (offset 12-19) without
    /// deserialising the rest.
    ///
    /// ### Params
    ///
    /// * `indices` - Cell indices
    ///
    /// ### Returns
    ///
    /// Vector of library sizes
    fn read_cell_library_sizes(&self, indices: &[usize]) -> Result<Vec<usize>, BixverseErrors> {
        if !self.is_cell_based() {
            return Err(BixverseErrors::ReaderModeMismatch {
                actual: "gene-based",
                requested: "cell-based",
            });
        }

        indices
            .par_iter()
            .map(|&original_index| {
                let decompressed = self.decompress_chunk(original_index)?;

                Ok(u64::from_le_bytes([
                    decompressed[12],
                    decompressed[13],
                    decompressed[14],
                    decompressed[15],
                    decompressed[16],
                    decompressed[17],
                    decompressed[18],
                    decompressed[19],
                ]) as usize)
            })
            .collect()
    }

    /// Read the NNZ for specific genes
    ///
    /// Overrides the default, which parses whole chunks. Reads directly from
    /// the chunk header bytes (offset 16-23) instead.
    ///
    /// ### Params
    ///
    /// * `indices` - Gene indices
    ///
    /// ### Returns
    ///
    /// Vector of number of NNZ genes.
    fn read_gene_nnz(&self, indices: &[usize]) -> Result<Vec<usize>, BixverseErrors> {
        if !self.is_gene_based() {
            return Err(BixverseErrors::ReaderModeMismatch {
                actual: "cell-based",
                requested: "gene-based",
            });
        }

        indices
            .par_iter()
            .map(|&original_index| {
                let decompressed = self.decompress_chunk(original_index)?;

                Ok(u64::from_le_bytes([
                    decompressed[16],
                    decompressed[17],
                    decompressed[18],
                    decompressed[19],
                    decompressed[20],
                    decompressed[21],
                    decompressed[22],
                    decompressed[23],
                ]) as usize)
            })
            .collect()
    }

    /// Is the reader connected to a cell-based file?
    ///
    /// ### Returns
    ///
    /// Bool
    fn is_cell_based(&self) -> bool {
        self.header.cell_based
    }

    /// Library size the `data_norm` layer of this file was scaled to.
    ///
    /// ### Returns
    ///
    /// `Some(target_size)`, or `None` for raw-only files and for files written
    /// before the field entered the header.
    fn target_size(&self) -> Option<f32> {
        (self.target_size != 0.0).then_some(self.target_size)
    }
}

//////////////////////
// Chunks to sparse //
//////////////////////

#[derive(Clone, Debug, Copy)]
/// Which of the data layers to return from a given chunk when transforming
/// into a [CompressedSparseData2]
pub enum DataLayerReturn {
    /// Return both layers, i.e., raw and
    BothLayers,
    /// Return only the raw layer
    Raw,
    /// Return only the norm layer
    Norm,
}

/// Converts a slice of gene chunks into a CSC sparse matrix
///
/// Constructs a cells x genes compressed sparse column matrix from individual
/// gene chunks. Depending on `data_layer`, the raw counts and/or the normalised
/// counts (f32) are populated.
///
/// Raw counts are stored on disk as `u32` and are converted into `T` exactly:
/// a count that does not fit returns [`BixverseErrors::RawCountOverflow`]
/// rather than saturating. Note that `T = f32` starts rounding above 2^24,
/// which is far beyond any plausible per-gene count.
///
/// ### Params
///
/// * `chunks` - Vector of gene chunks to convert
/// * `data_layer` - Which data layer(s) to populate in the output
/// * `n_cells` - Total number of cells in the dataset
///
/// ### Returns
///
/// A CSC-formatted sparse matrix with the requested layer(s) populated.
pub fn from_gene_chunks<T>(
    chunks: Vec<CscGeneChunk>,
    data_layer: &DataLayerReturn,
    n_cells: usize,
) -> Result<CompressedSparseData2<T, f32>, BixverseErrors>
where
    T: BixverseNumeric + FromPrimitive,
{
    let n_genes = chunks.len();
    let keep_raw = matches!(
        data_layer,
        DataLayerReturn::BothLayers | DataLayerReturn::Raw
    );
    let keep_norm = matches!(
        data_layer,
        DataLayerReturn::BothLayers | DataLayerReturn::Norm
    );

    let mut data: Vec<T> = Vec::new();
    let mut data_2: Vec<f32> = Vec::new();
    let mut indices: Vec<usize> = Vec::new();
    let mut indptr = Vec::with_capacity(n_genes + 1);
    indptr.push(0);

    for chunk in chunks {
        if keep_raw {
            for val in chunk.data_raw.iter() {
                data.push(T::from_u32(val).ok_or(BixverseErrors::RawCountOverflow {
                    value: val,
                    target_type: std::any::type_name::<T>(),
                })?);
            }
        }
        if keep_norm {
            for &val in &chunk.data_norm {
                data_2.push(val.to_f32());
            }
        }
        for &idx in &chunk.indices {
            indices.push(idx as usize);
        }
        indptr.push(indices.len());
    }

    Ok(CompressedSparseData2 {
        data,
        indices: indices.index_cast(),
        indptr: indptr.index_cast(),
        cs_type: CompressedSparseFormat::Csc,
        data_2: if keep_norm { Some(data_2) } else { None },
        shape: (n_cells, n_genes),
    })
}

/// Converts a slice of cell chunks into a CSR sparse matrix
///
/// Constructs a cells x genes compressed sparse row matrix from individual
/// cell chunks. Same exact-conversion behaviour as [`from_gene_chunks`]: a raw
/// count that does not fit `T` returns [`BixverseErrors::RawCountOverflow`]
/// rather than saturating.
///
/// ### Params
///
/// * `chunks` - Slice of cell chunks to convert
/// * `data_layer` - Which data layer(s) to populate in the output
/// * `n_genes` - Total number of genes in the dataset
///
/// ### Returns
///
/// A CSR-formatted sparse matrix with the requested layer(s) populated.
pub fn from_cell_chunks<T>(
    chunks: &[CsrCellChunk],
    data_layer: &DataLayerReturn,
    n_genes: usize,
) -> Result<CompressedSparseData2<T, f32>, BixverseErrors>
where
    T: BixverseNumeric + FromPrimitive,
{
    let n_cells = chunks.len();
    let keep_raw = matches!(
        data_layer,
        DataLayerReturn::BothLayers | DataLayerReturn::Raw
    );
    let keep_norm = matches!(
        data_layer,
        DataLayerReturn::BothLayers | DataLayerReturn::Norm
    );

    let mut data: Vec<T> = Vec::new();
    let mut data_2: Vec<f32> = Vec::new();
    let mut indices: Vec<usize> = Vec::new();
    let mut indptr = Vec::with_capacity(n_cells + 1);
    indptr.push(0);

    for chunk in chunks {
        if keep_raw {
            for val in chunk.data_raw.iter() {
                data.push(T::from_u32(val).ok_or(BixverseErrors::RawCountOverflow {
                    value: val,
                    target_type: std::any::type_name::<T>(),
                })?);
            }
        }
        if keep_norm {
            for &val in &chunk.data_norm {
                data_2.push(val.to_f32());
            }
        }
        for &idx in &chunk.indices {
            indices.push(idx as usize);
        }
        indptr.push(indices.len());
    }

    Ok(CompressedSparseData2 {
        data,
        indices: indices.index_cast(),
        indptr: indptr.index_cast(),
        cs_type: CompressedSparseFormat::Csr,
        data_2: if keep_norm { Some(data_2) } else { None },
        shape: (n_cells, n_genes),
    })
}

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

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

    /// RAII guard that removes a test's temp file even if an assert fails.
    struct TempBin(std::path::PathBuf);

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

    impl TempBin {
        /// Reserve a uniquely named scratch file in the system temp directory.
        ///
        /// ### Params
        ///
        /// * `name` - Test-unique suffix.
        ///
        /// ### Returns
        ///
        /// The guard; the path is available via [`Self::path`].
        fn new(name: &str) -> Self {
            Self(std::env::temp_dir().join(format!("bixverse_data_io_{name}.bin")))
        }

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

    /// Build a cell chunk from raw counts and gene indices.
    ///
    /// ### Params
    ///
    /// * `raw` - Raw counts, converted via [`RawCounts::from_u32_auto`].
    /// * `indices` - Gene indices, same length as `raw`.
    /// * `original_index` - Index the chunk is stored under.
    ///
    /// ### Returns
    ///
    /// The `CsrCellChunk`.
    fn cell_chunk(raw: &[u32], indices: &[u32], original_index: usize) -> CsrCellChunk {
        CsrCellChunk {
            data_raw: RawCounts::from_u32_auto(raw),
            data_norm: raw
                .iter()
                .map(|&v| F16::from(f16::from_f32(v as f32)))
                .collect(),
            library_size: raw.iter().map(|&v| v as usize).sum(),
            indices: indices.to_vec(),
            original_index,
            to_keep: true,
        }
    }

    /// Build a gene chunk from raw counts and cell indices.
    ///
    /// ### Params
    ///
    /// * `raw` - Raw counts, converted via [`RawCounts::from_u32_auto`].
    /// * `indices` - Cell indices, same length as `raw`.
    /// * `original_index` - Index the chunk is stored under.
    ///
    /// ### Returns
    ///
    /// The `CscGeneChunk`.
    fn gene_chunk(raw: &[u32], indices: &[u32], original_index: usize) -> CscGeneChunk {
        CscGeneChunk {
            data_raw: RawCounts::from_u32_auto(raw),
            data_norm: raw
                .iter()
                .map(|&v| F16::from(f16::from_f32(v as f32)))
                .collect(),
            avg_exp: F16::from(f16::from_f32(1.0)),
            nnz: raw.len(),
            indices: indices.to_vec(),
            original_index,
            to_keep: true,
        }
    }

    ////////////////////
    // Chunk codec //
    ////////////////////

    /// Pins the CSR chunk codec: every field must return bit-identical.
    #[test]
    fn test_cell_chunk_round_trip_preserves_fields() {
        let raw = [1_u32, 7, 42, 65_535, 3];
        let indices = [0_u32, 5, 9, 100_000, 300_000];
        let chunk = cell_chunk(&raw, &indices, 17);

        let mut buffer = Vec::new();
        chunk.write_to_bytes(&mut buffer).expect("serialise");
        let back = CsrCellChunk::read_from_buffer(&buffer).expect("deserialise");

        assert_eq!(back.data_raw.iter().collect::<Vec<_>>(), raw.to_vec());
        assert_eq!(back.indices, indices.to_vec());
        assert_eq!(back.library_size, chunk.library_size);
        assert_eq!(back.original_index, 17);
        assert!(back.to_keep);
        let bits: Vec<u16> = back.data_norm.iter().map(|v| v.to_bits()).collect();
        let want: Vec<u16> = chunk.data_norm.iter().map(|v| v.to_bits()).collect();
        assert_eq!(bits, want);
    }

    /// Same pin for the CSC codec, which carries `nnz` and `avg_exp` too.
    #[test]
    fn test_gene_chunk_round_trip_preserves_fields() {
        let raw = [2_u32, 0, 99, 70_000];
        let indices = [1_u32, 4, 8, 12];
        let chunk = gene_chunk(&raw, &indices, 3);

        let mut buffer = Vec::new();
        chunk.write_to_bytes(&mut buffer).expect("serialise");
        let back = CscGeneChunk::read_from_buffer(&buffer).expect("deserialise");

        assert_eq!(back.data_raw.iter().collect::<Vec<_>>(), raw.to_vec());
        assert_eq!(back.indices, indices.to_vec());
        assert_eq!(back.nnz, 4);
        assert_eq!(back.original_index, 3);
        assert_eq!(back.avg_exp.to_bits(), chunk.avg_exp.to_bits());
    }

    /// Regression: the CSC guard used to check 32 bytes then slice 36.
    #[test]
    fn test_gene_chunk_short_buffer_errors_instead_of_panicking() {
        for len in CSR_CHUNK_HEADER_LEN..CSC_CHUNK_HEADER_LEN {
            let buffer = vec![0_u8; len];
            match CscGeneChunk::read_from_buffer(&buffer) {
                Err(BixverseErrors::ChunkBufferTooSmall { expected, found }) => {
                    assert_eq!(expected, CSC_CHUNK_HEADER_LEN);
                    assert_eq!(found, len);
                }
                other => panic!("expected ChunkBufferTooSmall for len {len}, got {other:?}"),
            }
        }
    }

    /// Regression: payload lengths came from untrusted header fields and were
    /// never validated against the buffer.
    #[test]
    fn test_chunk_payload_truncation_is_caught() {
        let chunk = cell_chunk(&[1, 2, 3], &[0, 1, 2], 0);
        let mut buffer = Vec::new();
        chunk.write_to_bytes(&mut buffer).expect("serialise");

        // Claim four times as many raw counts as were written.
        buffer[0..4].copy_from_slice(&12_u32.to_le_bytes());

        assert!(matches!(
            CsrCellChunk::read_from_buffer(&buffer),
            Err(BixverseErrors::ChunkPayloadTruncated { .. })
        ));

        // Same for a gene chunk, via the indices length.
        let gene = gene_chunk(&[1, 2, 3], &[0, 1, 2], 0);
        let mut buffer = Vec::new();
        gene.write_to_bytes(&mut buffer).expect("serialise");
        buffer[8..12].copy_from_slice(&64_u32.to_le_bytes());

        assert!(matches!(
            CscGeneChunk::read_from_buffer(&buffer),
            Err(BixverseErrors::ChunkPayloadTruncated { .. })
        ));
    }

    /// Regression: with `elem_size == 2` and an odd raw/norm element count the
    /// old `*const u32` read landed on a 2-byte-aligned address.
    #[test]
    fn test_chunk_round_trip_with_unaligned_index_offset() {
        // 3 raw u16 + 3 norm u16 puts the indices at offset 32 + 6 + 6 = 44,
        // which is 4-byte misaligned.
        let raw = [1_u32, 2, 3];
        let indices = [7_u32, 8, 9];
        let chunk = cell_chunk(&raw, &indices, 0);
        assert_eq!(chunk.data_raw.elem_size(), RAW_ELEM_U16);

        let mut buffer = Vec::new();
        chunk.write_to_bytes(&mut buffer).expect("serialise");
        assert_eq!(buffer.len() % 4, 0);

        let back = CsrCellChunk::read_from_buffer(&buffer).expect("deserialise");
        assert_eq!(back.indices, indices.to_vec());
        assert_eq!(back.data_raw.iter().collect::<Vec<_>>(), raw.to_vec());
    }

    /// A zero discriminant still means u16; an unknown width must error.
    #[test]
    fn test_raw_counts_element_size_discriminant() {
        let payload = [1_u8, 0, 2, 0];

        // Legacy files left the discriminant byte as zero padding.
        let legacy = RawCounts::read_from_buffer(&payload, 2, 0).expect("legacy parses as u16");
        assert_eq!(legacy.iter().collect::<Vec<_>>(), vec![1, 2]);

        let explicit =
            RawCounts::read_from_buffer(&payload, 2, RAW_ELEM_U16).expect("explicit u16 parses");
        assert_eq!(explicit.iter().collect::<Vec<_>>(), vec![1, 2]);

        assert!(matches!(
            RawCounts::read_from_buffer(&payload, 2, 3),
            Err(BixverseErrors::RawElemSizeInvalid(3))
        ));
    }

    /////////////////////
    // Count widening //
    /////////////////////

    /// Regression: counts above `u16::MAX` used to saturate silently.
    #[test]
    fn test_large_raw_count_survives_round_trip() {
        let raw = [70_000_u32, 1];
        let chunk = cell_chunk(&raw, &[0, 1], 0);
        assert_eq!(chunk.data_raw.elem_size(), RAW_ELEM_U32);

        let mut buffer = Vec::new();
        chunk.write_to_bytes(&mut buffer).expect("serialise");
        let back = CsrCellChunk::read_from_buffer(&buffer).expect("deserialise");

        let csr = from_cell_chunks::<u32>(&[back], &DataLayerReturn::Raw, 2).expect("assemble");
        assert_eq!(csr.data, vec![70_000_u32, 1]);
    }

    /// A count above `u16::MAX` must be reported, not wrapped or saturated.
    #[test]
    fn test_narrow_target_type_reports_overflow() {
        let chunk = gene_chunk(&[70_000, 1], &[0, 1], 0);

        match from_gene_chunks::<u16>(vec![chunk], &DataLayerReturn::Raw, 2) {
            Err(BixverseErrors::RawCountOverflow { value, .. }) => assert_eq!(value, 70_000),
            other => panic!("expected RawCountOverflow, got {:?}", other.map(|_| ())),
        }
    }

    ///////////////////////
    // Writer and reader //
    ///////////////////////

    /// Writer to reader round trip: orientation, target size and both chunks.
    #[test]
    fn test_cell_based_file_round_trip() {
        let temp = TempBin::new("cell_round_trip");
        let chunks = vec![
            cell_chunk(&[5, 70_000], &[0, 2], 0),
            cell_chunk(&[1, 2, 3], &[0, 1, 2], 1),
        ];

        let mut writer =
            CellGeneSparseWriter::new(temp.path(), true, 2, 3, 1e4).expect("writer opens");
        for chunk in chunks {
            writer.write_cell_chunk(chunk).expect("write");
        }
        writer.finalise().expect("finalise");

        let reader = ParallelSparseReader::new(temp.path()).expect("reader opens");
        assert_eq!(reader.target_size(), Some(1e4));
        assert!(reader.is_cell_based());

        let back = reader.read_cells_parallel(&[0, 1]).expect("read");
        assert_eq!(back.len(), 2);
        assert_eq!(back[0].data_raw.iter().collect::<Vec<_>>(), vec![5, 70_000]);
        assert_eq!(back[1].indices, vec![0, 1, 2]);
    }

    /// Rewrite the 64-byte header slot of an existing file, applying `edit` to
    /// the decoded [`FileHeader`]. Lets a test corrupt one field of an otherwise
    /// valid file without hand-assembling the whole layout.
    fn patch_file_header(path: &str, edit: impl FnOnce(&mut FileHeader)) {
        let mut bytes = std::fs::read(path).expect("read back");
        let (mut header, _) = decode_from_slice::<FileHeader, _>(&bytes[0..64], config::standard())
            .expect("header decodes");

        edit(&mut header);

        let encoded = encode_to_vec(&header, config::standard()).expect("header encodes");
        let mut slot = [0_u8; 64];
        slot[..encoded.len()].copy_from_slice(&encoded);
        bytes[0..64].copy_from_slice(&slot);

        std::fs::write(path, bytes).expect("write back");
    }

    /// Build a minimal but valid cell-based file at `path`.
    fn write_valid_file(path: &str) {
        let mut writer = CellGeneSparseWriter::new(path, true, 1, 3, 1e4).expect("writer opens");
        writer
            .write_cell_chunk(cell_chunk(&[1, 2], &[0, 1], 0))
            .expect("write");
        writer.finalise().expect("finalise");
    }

    /// Regression: `ParallelSparseReader::new` rejects any file whose version is
    /// not `SC_FILE_VERSION`, and nothing exercised that guard.
    #[test]
    fn test_reader_rejects_a_version_mismatch() {
        let temp = TempBin::new("version_mismatch");
        write_valid_file(temp.path());
        patch_file_header(temp.path(), |h| h.version = SC_FILE_VERSION + 1);

        match ParallelSparseReader::new(temp.path()) {
            Err(BixverseErrors::FileVersionMismatch { expected, found }) => {
                assert_eq!(expected, SC_FILE_VERSION);
                assert_eq!(found, SC_FILE_VERSION + 1);
            }
            Err(other) => panic!("expected a version mismatch, got {other:?}"),
            Ok(_) => panic!("a mismatched version must not open"),
        }
    }

    /// A file that is not ours at all must be refused rather than parsed.
    #[test]
    fn test_reader_rejects_a_bad_magic() {
        let temp = TempBin::new("bad_magic");
        write_valid_file(temp.path());
        patch_file_header(temp.path(), |h| h.magic = *b"NOTOURS!");

        assert!(
            ParallelSparseReader::new(temp.path()).is_err(),
            "a wrong magic string must not open"
        );
    }

    /// A header pointing past the end of the file must error rather than read
    /// out of bounds off the mmap.
    #[test]
    fn test_reader_rejects_a_main_header_offset_past_eof() {
        let temp = TempBin::new("offset_past_eof");
        write_valid_file(temp.path());
        patch_file_header(temp.path(), |h| h.main_header_offset = u64::MAX / 2);

        assert!(
            ParallelSparseReader::new(temp.path()).is_err(),
            "an out-of-range main header offset must not open"
        );
    }

    /// A file truncated mid-payload must error rather than panic.
    #[test]
    fn test_reader_rejects_a_truncated_file() {
        let temp = TempBin::new("truncated");
        write_valid_file(temp.path());

        let bytes = std::fs::read(temp.path()).expect("read back");
        std::fs::write(temp.path(), &bytes[..bytes.len() / 2]).expect("truncate");

        assert!(
            ParallelSparseReader::new(temp.path()).is_err(),
            "a truncated file must not open"
        );
    }

    /// The gene-based half of the same pin, with one chunk forced to u32 raw.
    #[test]
    fn test_gene_based_file_round_trip() {
        let temp = TempBin::new("gene_round_trip");

        let mut writer =
            CellGeneSparseWriter::new(temp.path(), false, 4, 2, 1e5).expect("writer opens");
        writer
            .write_gene_chunk(gene_chunk(&[9, 8], &[0, 3], 0))
            .expect("write");
        writer
            .write_gene_chunk(gene_chunk(&[100_000], &[2], 1))
            .expect("write");
        writer.finalise().expect("finalise");

        let reader = ParallelSparseReader::new(temp.path()).expect("reader opens");
        assert_eq!(reader.target_size(), Some(1e5));
        assert!(reader.is_gene_based());

        let back = reader.read_gene_parallel(&[0, 1]).expect("read");
        assert_eq!(back[0].data_raw.iter().collect::<Vec<_>>(), vec![9, 8]);
        assert_eq!(back[1].data_raw.iter().collect::<Vec<_>>(), vec![100_000]);
    }

    /// Regression: the writer used to `assert!` on the orientation.
    #[test]
    fn test_writer_rejects_wrong_chunk_orientation() {
        let temp = TempBin::new("mode_mismatch");
        let mut writer =
            CellGeneSparseWriter::new(temp.path(), true, 1, 1, 0.0).expect("writer opens");

        match writer.write_gene_chunk(gene_chunk(&[1], &[0], 0)) {
            Err(BixverseErrors::ReaderModeMismatch { actual, requested }) => {
                assert_eq!(actual, "cell-based");
                assert_eq!(requested, "gene-based");
            }
            other => panic!("expected ReaderModeMismatch, got {other:?}"),
        }
    }

    /// A target size of 0 means raw only: reader and peek must both say `None`.
    #[test]
    fn test_raw_only_file_reports_no_target_size() {
        let temp = TempBin::new("no_target_size");
        let mut writer =
            CellGeneSparseWriter::new(temp.path(), true, 1, 2, 0.0).expect("writer opens");
        writer
            .write_cell_chunk(cell_chunk(&[1, 2], &[0, 1], 0))
            .expect("write");
        writer.finalise().expect("finalise");

        let reader = ParallelSparseReader::new(temp.path()).expect("reader opens");
        assert_eq!(reader.target_size(), None);
        assert_eq!(peek_target_size(temp.path()).expect("peek"), None);
    }

    /////////////////
    // File header //
    /////////////////

    /// The v3 header layout before `target_size` was carved out of the
    /// reserved block. Used to prove the swap is byte-neutral in both
    /// directions.
    #[derive(Encode, Decode, Serialize, Deserialize)]
    struct LegacyFileHeader {
        /// Magic string as bytes to recognise the file
        magic: [u8; 8],
        /// Version of the file
        version: u32,
        /// Offset of the main header
        main_header_offset: u64,
        /// Is cell-based (CSR format)
        cell_based: bool,
        /// 32 reserved bytes, all zero in files written by the old writer
        _reserved_1: [u8; 32],
        /// 3 further reserved bytes
        _reserved_2: [u8; 3],
    }

    /// Encode a header into the fixed, zero-padded 64-byte slot.
    ///
    /// ### Params
    ///
    /// * `header` - Anything the bincode standard config can encode.
    ///
    /// ### Returns
    ///
    /// Exactly 64 bytes.
    fn encode_slot<T: Serialize>(header: &T) -> [u8; 64] {
        let encoded = encode_to_vec(header, config::standard()).expect("encode");
        assert!(encoded.len() <= 64, "header must fit the fixed slot");
        let mut slot = [0_u8; 64];
        slot[..encoded.len()].copy_from_slice(&encoded);
        slot
    }

    /// Backwards: pre-`target_size` header bytes must still decode, reading 0.
    #[test]
    fn test_legacy_header_decodes_with_zero_target_size() {
        let legacy = LegacyFileHeader {
            magic: *b"SCRNASEQ",
            version: SC_FILE_VERSION,
            main_header_offset: 123_456,
            cell_based: true,
            _reserved_1: [0; 32],
            _reserved_2: [0; 3],
        };

        let slot = encode_slot(&legacy);
        let (decoded, _) =
            decode_from_slice::<FileHeader, _>(&slot, config::standard()).expect("decode as v3");

        assert_eq!(decoded.magic, *b"SCRNASEQ");
        assert_eq!(decoded.version, SC_FILE_VERSION);
        assert_eq!(decoded.main_header_offset, 123_456);
        assert!(decoded.cell_based);
        assert_eq!(decoded.target_size, 0.0);
    }

    /// Forwards: today's header must still decode under the old layout.
    #[test]
    fn test_new_header_still_decodes_on_the_old_layout() {
        let mut header = FileHeader::new(false, 1e4);
        header.main_header_offset = 987_654_321;

        let slot = encode_slot(&header);
        let (decoded, _) = decode_from_slice::<LegacyFileHeader, _>(&slot, config::standard())
            .expect("decode as legacy");

        assert_eq!(decoded.magic, *b"SCRNASEQ");
        assert_eq!(decoded.version, SC_FILE_VERSION);
        assert_eq!(decoded.main_header_offset, 987_654_321);
        assert!(!decoded.cell_based);
    }

    /// The new field survives the fixed, zero-padded 64-byte slot.
    #[test]
    fn test_target_size_survives_header_round_trip() {
        let slot = encode_slot(&FileHeader::new(true, 1e5));
        let (decoded, _) =
            decode_from_slice::<FileHeader, _>(&slot, config::standard()).expect("decode");
        assert_eq!(decoded.target_size, 1e5);
    }
}