bixverse-sc-data 0.9.4

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

use hdf5::{File, Result};
use ndarray::{Array2, s};
use num_traits::Float;
use rayon::prelude::*;
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Instant;
use thousands::Separable;

use crate::H5_CELL_SLICE_SIZE;
use crate::data_io::{CellGeneSparseWriter, CellOnFileQuality};
use crate::data_io::{
    INDEX_DROPPED, compress_cell_row, dense_index_map, write_cell_chunks_built, write_cell_rows,
};
use crate::h5_chunked::H5ChunkedReader;
use crate::h5_filters::{check_h5_filters, check_h5ad_filters};
use crate::prelude::*;

////////////
// Consts //
////////////

/// Number of data entries to check
const N_CHECK: usize = 100;

/////////////////
// h5ad format //
/////////////////

/// The format in which the H5AD counts were stored (because the field is
/// insane and stores them in so many ways...)
pub enum H5ADFormat {
    /// CSR type format.
    Csr,
    /// CSC type format
    Csc,
    /// Dense (row major)
    DenseRow,
    /// Dense (col major)
    DenseCol,
}

/// Parse the h5ad save format
///
/// ### Param
///
/// * `s` - The string to parse
///
/// ### Returns
///
/// The [H5ADFormat]
pub fn parse_h5ad_format(s: &str) -> Option<H5ADFormat> {
    match s.to_lowercase().as_str() {
        "csr" => Some(H5ADFormat::Csr),
        "csc" => Some(H5ADFormat::Csc),
        "dense_row" => Some(H5ADFormat::DenseRow),
        "dense_col" => Some(H5ADFormat::DenseCol),
        _ => None,
    }
}

///////////////////
// Raw data slot //
///////////////////

/// Enum describing in which of the slots the raw counts are stored
#[derive(Clone, Debug, Default, Copy)]
pub enum RawDataSlot {
    /// Default assumption that raw counts are stored `ad.X`
    #[default]
    DataX,
    /// Alternative assumption that raw counts are stored in `ad.raw.X`
    RawX,
    /// Alternative assumption that raw counts are stored in `ad.layers.counts`
    LayerCounts,
}

impl RawDataSlot {
    /// Returns the indptr path
    ///
    /// ### Returns
    ///
    /// `indptr` path from the h5ad file
    #[inline]
    pub fn get_indptr(&self) -> &str {
        match self {
            RawDataSlot::DataX => "X/indptr",
            RawDataSlot::RawX => "raw/X/indptr",
            RawDataSlot::LayerCounts => "layers/counts/indptr",
        }
    }

    /// Returns the indices path
    ///
    /// ### Returns
    ///
    /// `indices` path from the h5ad file
    #[inline]
    pub fn get_indices(&self) -> &str {
        match self {
            RawDataSlot::DataX => "X/indices",
            RawDataSlot::RawX => "raw/X/indices",
            RawDataSlot::LayerCounts => "layers/counts/indices",
        }
    }

    /// Returns the data path
    ///
    /// ### Returns
    ///
    /// `data` path from the h5ad file
    #[inline]
    pub fn get_data(&self) -> &str {
        match self {
            RawDataSlot::DataX => "X/data",
            RawDataSlot::RawX => "raw/X/data",
            RawDataSlot::LayerCounts => "layers/counts/data",
        }
    }

    /// Get the dense path
    ///
    /// ### Returns
    ///
    /// The path to the dense array
    #[inline]
    pub fn get_dense_path(&self) -> &str {
        match self {
            RawDataSlot::DataX => "X",
            RawDataSlot::RawX => "raw/X",
            RawDataSlot::LayerCounts => "layers/counts",
        }
    }
}

/// Parse where the raw counts are stored
///
/// ### Params
///
/// * `s` - String to parse
///
/// ### Returns
///
/// An option of the [RawDataSlot] to use subsequently.
pub fn parse_raw_slot(s: &str) -> Option<RawDataSlot> {
    match s.to_lowercase().as_str() {
        "x" | "X" => Some(RawDataSlot::DataX),
        "raw" | "raw.x" => Some(RawDataSlot::RawX),
        "layers.counts" => Some(RawDataSlot::LayerCounts),
        _ => None,
    }
}

/// Helper function to check if data is likely raw data
///
/// ### Params
///
/// * `data` - Numeric vector to test
/// * `n_check` - Number of values to check
///
/// ### Returns
///
/// Nothing if okay, Err if not.
pub fn check_is_raw<T>(data: &[T], n_check: Option<usize>) -> Result<(), ScDataErrors>
where
    T: Float,
{
    let n_check = n_check.unwrap_or(N_CHECK);

    let res = data.iter().take(n_check).all(|x| x.fract() == T::zero());

    if res {
        Ok(())
    } else {
        Err(ScDataErrors::NotRawCounts)
    }
}

/// Check raw data in h5ad
///
/// ### Params
///
/// * `file` - The h5ad file reference
/// * `raw_slot` - The slot for the raw data
/// * `n_check` - Optional number of entries to check. If not provided, defaults
///   to `N_CHECK` (100).
///
/// ### Returns
///
/// Nothing if okay, Err if not.
pub fn check_h5ad_is_raw(
    file: &File,
    raw_slot: &RawDataSlot,
    n_check: Option<usize>,
) -> Result<(), ScDataErrors> {
    let ds = file.dataset(raw_slot.get_data())?;
    let n_check = n_check.unwrap_or(N_CHECK).min(ds.size());

    let data: Vec<f32> = ds.read_slice(0..n_check)?.to_vec();

    check_is_raw(&data, Some(n_check))?;

    Ok(())
}

/// Sanity-check that a dense slot looks like raw counts.
///
/// ### Params
///
/// * `file` - The h5ad file reference
/// * `raw_slot` - The slot for the raw data
/// * `n_check` - Optional number of entries to check. If not provided, defaults
///   to `N_CHECK` (100).
///
/// ### Returns
///
/// Nothing if okay, Err if not.
pub fn check_h5ad_dense_is_raw(
    file: &File,
    raw_slot: &RawDataSlot,
    n_check: Option<usize>,
) -> Result<(), ScDataErrors> {
    let n_check = n_check.unwrap_or(100);

    let ds = file.dataset(raw_slot.get_dense_path())?;
    let shape = ds.shape();
    if shape.len() != 2 {
        return Ok(());
    }

    let rows = shape[0];
    let cols = shape[1];
    if rows == 0 || cols == 0 {
        return Ok(());
    }

    let take_cols = cols.min(n_check.max(1));
    let take_rows = n_check.div_ceil(take_cols).min(rows).max(1);

    let block = ds.read_slice_2d((0..take_rows, 0..take_cols))?;
    let nonzeros: Vec<f32> = block.iter().copied().filter(|&v| v != 0.0).collect();
    if nonzeros.is_empty() {
        return Ok(());
    }
    let n = nonzeros.len().min(n_check);
    check_is_raw(&nonzeros[..n], Some(n))
}

/////////////
// Writers //
/////////////

/// Writes h5ad data to disk in the binarised format
///
/// Function to take in h5 data and write it to disk (first the cells) in a
/// binarised format for fast retrieval of cells. Pending on how the file was
/// stored (CSC or CSR) different paths will be used to process the data.
///
/// ### Params
///
/// * `h5_path` - Path to the h5 object.
/// * `bin_path` - Path to the binarised object on disk to write to
/// * `h5ad_format` - Was the h5ad data stored in: "csc", "csr", "dense_row",
///   or "dense_col".
/// * `no_cells` - Total number of obversations in the data.
/// * `no_genes` - Total number of vars in the data.
/// * `cell_quality` - Structure containing information on the desired minimum
///   cell and gene quality + target size for library normalisation.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot]
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// A tuple with `(no_cells, no_genes, cell quality metrics)`
#[allow(clippy::too_many_arguments)]
pub fn write_h5_counts<P: AsRef<Path>>(
    h5_path: P,
    bin_path: P,
    h5ad_format: &str,
    no_cells: usize,
    no_genes: usize,
    cell_quality: MinCellQuality,
    raw_slot: RawDataSlot,
    verbose: bool,
) -> Result<(usize, usize, CellQuality), ScDataErrors> {
    if verbose {
        println!(
            "Step 1/4: Analysing data structure, calculating QC metrics and identifying cells/genes to take..."
        );
    }

    let file_format = parse_h5ad_format(h5ad_format)
        .ok_or_else(|| ScDataErrors::UnsupportH5ADFormat(h5ad_format.to_string()))?;

    check_h5ad_filters(&File::open(h5_path.as_ref())?, &raw_slot, &file_format)?;

    let file_quality = match file_format {
        H5ADFormat::Csr => parse_h5_csr_quality(
            &h5_path,
            (no_cells, no_genes),
            &raw_slot,
            &cell_quality,
            verbose,
        )?,
        H5ADFormat::Csc => parse_h5_csc_quality(
            &h5_path,
            (no_cells, no_genes),
            &cell_quality,
            &raw_slot,
            verbose,
        )?,
        H5ADFormat::DenseRow => parse_h5_dense_row_quality(
            &h5_path,
            (no_cells, no_genes),
            &raw_slot,
            &cell_quality,
            verbose,
        )?,
        H5ADFormat::DenseCol => return dense_col_not_supported(),
    };

    if verbose {
        println!("Step 2/4: QC Results:");
        println!(
            "  Genes passing QC (i.e., getting included): {} / {}",
            file_quality.genes_to_keep.len().separate_with_underscores(),
            no_genes.separate_with_underscores()
        );
        println!(
            "  Cells passing QC (i.e., getting included): {} / {}",
            file_quality.cells_to_keep.len().separate_with_underscores(),
            no_cells.separate_with_underscores()
        );
        println!("Step 3/4: Loading filtered data from h5...");
    }

    let file_data: CompressedSparseData2<u32> = match file_format {
        H5ADFormat::Csr => read_h5ad_x_data_csr(&h5_path, &file_quality, &raw_slot, verbose)?,
        H5ADFormat::Csc => {
            let data = read_h5ad_x_data_csc(&h5_path, &file_quality, &raw_slot, verbose)?;
            data.transpose_and_convert()
        }
        H5ADFormat::DenseRow => {
            read_h5ad_x_data_dense_row(&h5_path, &file_quality, &raw_slot, verbose)?
        }
        H5ADFormat::DenseCol => return dense_col_not_supported(),
    };

    if verbose {
        println!("Step 4/4: Writing to binary format...");
    }

    let n_cells = file_data.indptr.len() - 1;
    let mut writer =
        CellGeneSparseWriter::new(bin_path, true, no_cells, no_genes, cell_quality.target_size)?;

    let (nnz, lib_size) = write_cell_chunks_built(
        n_cells,
        |i| {
            let start_i = file_data.indptr[i] as usize;
            let end_i = file_data.indptr[i + 1] as usize;
            CsrCellChunk::from_data(
                &file_data.data[start_i..end_i],
                &file_data.indices[start_i..end_i],
                i,
                cell_quality.target_size,
                true,
            )
        },
        &mut writer,
    )?;

    if verbose {
        println!(
            "  Written {} / {} cells (complete).",
            n_cells.separate_with_underscores(),
            n_cells.separate_with_underscores()
        );
        println!("Finalising file...");
    }

    writer.update_header_no_cells(file_quality.cells_to_keep.len());
    writer.update_header_no_genes(file_quality.genes_to_keep.len());
    writer.finalise()?;

    let cell_qc = CellQuality {
        cell_indices: file_quality.cells_to_keep.clone(),
        gene_indices: file_quality.genes_to_keep.clone(),
        lib_size,
        nnz,
    };

    Ok((
        file_quality.cells_to_keep.len(),
        file_quality.genes_to_keep.len(),
        cell_qc,
    ))
}

/// Read an `indptr` dataset as `usize`.
///
/// AnnData switches `indptr` to int64 once a matrix passes 2^31 stored values,
/// and a single Tahoe-style plate can pass 2^32. Reading straight into `u32`
/// would not survive that, so every h5 path goes through here.
///
/// ### Params
///
/// * `ds` - The `indptr` dataset
///
/// ### Returns
///
/// The offsets as `usize`.
pub(crate) fn read_indptr(ds: &hdf5::Dataset) -> Result<Vec<usize>> {
    Ok(ds.read_1d::<usize>()?.to_vec())
}

/// Function that streams the h5 counts to disk
///
/// This function will avoid (as far as possible) loading in the data into
/// memory and leverage direct streaming to disk where possible. Works best
/// with data that is already stored in CSR on disk (assuming genes * cells).
///
/// ### Params
///
/// * `h5_path` - Path to the h5 object.
/// * `bin_path` - Path to the binarised object on disk to write to
/// * `h5ad_format` - Was the h5ad data stored in: "csc", "csr", "dense_row",
///   or "dense_col".
/// * `no_cells` - Total number of obversations in the data.
/// * `no_genes` - Total number of vars in the data.
/// * `cell_quality` - Structure containing information on the desired minimum
///   cell and gene quality + target size for library normalisation.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot]
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// A tuple with `(no_cells, no_genes, cell quality metrics)`
#[allow(clippy::too_many_arguments)]
pub fn stream_h5_counts<P: AsRef<Path>>(
    h5_path: P,
    bin_path: P,
    h5ad_format: &str,
    no_cells: usize,
    no_genes: usize,
    cell_quality: MinCellQuality,
    raw_slot: RawDataSlot,
    verbose: bool,
) -> Result<(usize, usize, CellQuality), ScDataErrors> {
    if verbose {
        println!("Step 1/3: Analysing data structure and calculating QC metrics...");
    }

    let file_format = parse_h5ad_format(h5ad_format)
        .ok_or_else(|| ScDataErrors::UnsupportH5ADFormat(h5ad_format.to_string()))?;

    check_h5ad_filters(&File::open(h5_path.as_ref())?, &raw_slot, &file_format)?;

    let file_quality = match file_format {
        H5ADFormat::Csr => parse_h5_csr_quality(
            &h5_path,
            (no_cells, no_genes),
            &raw_slot,
            &cell_quality,
            verbose,
        )?,
        H5ADFormat::Csc => parse_h5_csc_quality(
            &h5_path,
            (no_cells, no_genes),
            &cell_quality,
            &raw_slot,
            verbose,
        )?,
        H5ADFormat::DenseRow => parse_h5_dense_row_quality(
            &h5_path,
            (no_cells, no_genes),
            &raw_slot,
            &cell_quality,
            verbose,
        )?,
        H5ADFormat::DenseCol => return dense_col_not_supported(),
    };

    if verbose {
        println!("Step 2/3: QC Results:");
        println!(
            "  Genes passing QC: {} / {}",
            file_quality.genes_to_keep.len().separate_with_underscores(),
            no_genes.separate_with_underscores()
        );
        println!(
            "  Cells passing QC: {} / {}",
            file_quality.cells_to_keep.len().separate_with_underscores(),
            no_cells.separate_with_underscores()
        );
        println!("Step 3/3: Writing cells to CSR format...");
    }

    let mut cell_qc = match file_format {
        H5ADFormat::Csr => write_h5_csr_streaming(
            &h5_path,
            &bin_path,
            &file_quality,
            cell_quality,
            &raw_slot,
            verbose,
        ),
        H5ADFormat::Csc => write_h5_csc_to_csr_streaming(
            &h5_path,
            &bin_path,
            &file_quality,
            &raw_slot,
            cell_quality.target_size,
            verbose,
        ),
        H5ADFormat::DenseRow => write_h5_dense_row_streaming(
            &h5_path,
            &bin_path,
            &file_quality,
            cell_quality,
            &raw_slot,
            verbose,
        ),
        H5ADFormat::DenseCol => return dense_col_not_supported(),
    }?;

    cell_qc.set_cell_indices(&file_quality.cells_to_keep);
    cell_qc.set_gene_indices(&file_quality.genes_to_keep);

    Ok((
        file_quality.cells_to_keep.len(),
        file_quality.genes_to_keep.len(),
        cell_qc,
    ))
}

//////////////
// CSC data //
//////////////

/// Get the cell quality data from a CSC file
///
/// This file assumes that the rows are representing cells and the columns
/// represent genes (typical for h5ad).
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file
/// * `shape` - Tuple with `(no_cells, no_genes)`.
/// * `cell_quality` - Structure defining the minimum quality values that are
///   expected here.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot]
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// `CellOnFileQuality` structure that contains all of the information about
/// which cells and genes to include.
pub fn parse_h5_csc_quality<P: AsRef<Path>>(
    file_path: P,
    shape: (usize, usize),
    cell_quality: &MinCellQuality,
    raw_slot: &RawDataSlot,
    verbose: bool,
) -> Result<CellOnFileQuality, ScDataErrors> {
    let file_path = file_path.as_ref();

    if verbose {
        println!(
            "  Reading CSC matrix structure (shape: {} x {})...",
            shape.0.separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
    }

    let file = File::open(file_path)?;
    let indptr = read_indptr(&file.dataset(raw_slot.get_indptr())?)?;

    check_h5ad_is_raw(&file, raw_slot, None)?;

    let mut no_cells_exp_gene: Vec<usize> = Vec::with_capacity(shape.1);
    for i in 0..shape.1 {
        no_cells_exp_gene.push(indptr[i + 1] - indptr[i]);
    }

    if verbose {
        let max_expr = no_cells_exp_gene.iter().max().unwrap_or(&0);
        let min_expr = no_cells_exp_gene.iter().min().unwrap_or(&0);
        let avg_expr = no_cells_exp_gene
            .iter()
            .sum::<usize>()
            .checked_div(shape.1)
            .unwrap_or(0);
        println!(
            "  Gene expression stats: min = {} | max = {} | avg = {} cells per gene",
            min_expr.separate_with_underscores(),
            max_expr.separate_with_underscores(),
            avg_expr.separate_with_underscores()
        );
    }

    let genes_to_keep: Vec<usize> = (0..shape.1)
        .filter(|&i| no_cells_exp_gene[i] >= cell_quality.min_cells)
        .collect();

    if verbose {
        println!(
            "  Genes passing filter: {} / {}",
            genes_to_keep.len().separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
        println!("Calculating cell metrics in parallel...");
    }

    const GENE_CHUNK_SIZE: usize = 10000;
    let gene_chunks: Vec<&[usize]> = genes_to_keep.chunks(GENE_CHUNK_SIZE).collect();
    let num_chunks = gene_chunks.len();

    let calc_time = Instant::now();
    let completed_chunks = Arc::new(AtomicUsize::new(0));
    let report_interval = (num_chunks / 10).max(1);

    let data_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_data())?);
    let indices_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_indices())?);

    let cell_stats: Vec<(Vec<usize>, Vec<f32>)> = gene_chunks
        .par_iter()
        .map(|gene_chunk| {
            let mut local_unique = vec![0usize; shape.0];
            let mut local_lib_size = vec![0.0f32; shape.0];

            let chunk_start_gene = gene_chunk[0];
            let chunk_end_gene = gene_chunk[gene_chunk.len() - 1];
            let data_start = indptr[chunk_start_gene];
            let data_end = indptr[chunk_end_gene + 1];

            if data_start >= data_end {
                return (local_unique, local_lib_size);
            }

            let (Ok(chunk_data), Ok(chunk_indices)) = (
                data_rd.read::<f32>(data_start..data_end),
                indices_rd.read::<u32>(data_start..data_end),
            ) else {
                return (local_unique, local_lib_size);
            };

            for &gene_idx in gene_chunk.iter() {
                let gene_data_start = indptr[gene_idx] - data_start;
                let gene_data_end = indptr[gene_idx + 1] - data_start;

                for idx in gene_data_start..gene_data_end {
                    let cell_idx = chunk_indices[idx] as usize;
                    if cell_idx < shape.0 {
                        local_unique[cell_idx] += 1;
                        local_lib_size[cell_idx] += chunk_data[idx];
                    }
                }
            }

            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 * 10.0).round() as usize) * 10;
                    println!(
                        "  Processed {}% of chunks ({}/{})",
                        progress, completed, num_chunks
                    );
                }
            }

            (local_unique, local_lib_size)
        })
        .collect();

    let mut cell_unique_genes = vec![0usize; shape.0];
    let mut cell_lib_size = vec![0.0f32; shape.0];

    for (local_unique, local_lib) in cell_stats {
        for i in 0..shape.0 {
            cell_unique_genes[i] += local_unique[i];
            cell_lib_size[i] += local_lib[i];
        }
    }

    let calc_elapsed = calc_time.elapsed();

    if verbose {
        println!("Cell metrics calculation done: {:.2?}", calc_elapsed);

        let max_genes = cell_unique_genes.iter().max().unwrap_or(&0);
        let min_genes = cell_unique_genes.iter().min().unwrap_or(&0);
        let max_lib = cell_lib_size.iter().fold(0.0f32, |a, &b| a.max(b));
        let min_lib = cell_lib_size.iter().fold(f32::INFINITY, |a, &b| a.min(b));

        println!(
            "  Cell stats: genes per cell: min = {} | max={}",
            min_genes.separate_with_underscores(),
            max_genes.separate_with_underscores()
        );
        println!(
            "  Cell stats: library size: min = {:.1} | max={:.1}",
            min_lib.separate_with_underscores(),
            max_lib.separate_with_underscores()
        );
    }

    let cells_to_keep: Vec<usize> = (0..shape.0)
        .filter(|&i| {
            cell_unique_genes[i] >= cell_quality.min_unique_genes
                && cell_lib_size[i] >= cell_quality.min_lib_size as f32
        })
        .collect();

    if verbose {
        println!(
            "  Cells passing filter: {} / {}",
            cells_to_keep.len().separate_with_underscores(),
            shape.0.separate_with_underscores()
        );
    }

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

    Ok(file_quality_data)
}

/// Helper function to scan a CSC file and get the lib sizes
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `quality` - Information on the which cells and genes to include after a
///   first pass of the file.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot]
///
/// ### Returns
///
/// A vector of library sizes in the cells.
pub fn scan_h5_csc_library_sizes<P: AsRef<Path>>(
    file_path: P,
    quality: &CellOnFileQuality,
    raw_slot: &RawDataSlot,
) -> Result<Vec<u32>, ScDataErrors> {
    let file = File::open(file_path)?;

    check_h5ad_is_raw(&file, raw_slot, None)?;

    let data_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_data())?);
    let indices_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_indices())?);
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;
    let indptr_raw = read_indptr(&indptr_ds)?;

    // Only track library sizes - very light in terms of memory
    let mut cell_lib_sizes = vec![0u32; quality.cells_to_keep.len()];
    let cell_map = dense_index_map(&quality.cells_to_keep);

    const GENE_CHUNK_SIZE: usize = 5000;

    for gene_chunk in quality.genes_to_keep.chunks(GENE_CHUNK_SIZE) {
        let start_pos = gene_chunk.iter().map(|&g| indptr_raw[g]).min().unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1])
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

        // only read data values, not indices - saves memory
        let chunk_data: Vec<f32> = data_rd.read(start_pos..end_pos)?;
        let chunk_indices: Vec<u32> = indices_rd.read(start_pos..end_pos)?;

        for &gene_idx in gene_chunk {
            let gene_start = indptr_raw[gene_idx];
            let gene_end = indptr_raw[gene_idx + 1];

            for idx in gene_start..gene_end {
                let local_idx = idx - start_pos;
                let old_cell_idx = chunk_indices[local_idx] as usize;

                let new_cell_idx = cell_map.get(old_cell_idx).copied().unwrap_or(INDEX_DROPPED);
                if new_cell_idx != INDEX_DROPPED {
                    cell_lib_sizes[new_cell_idx as usize] += chunk_data[local_idx] as u32;
                }
            }
        }
    }

    Ok(cell_lib_sizes)
}

/// Helper function that reads in full CSC data from an h5 file
///
/// The function assumes that the data is stored as cells x genes.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `quality` - Information on the which cells and genes to include after a
///   first pass of the file.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot]
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// The `CompressedSparseData2` in CSR format with the counts stored as u32.
pub fn read_h5ad_x_data_csc<P: AsRef<Path>>(
    file_path: P,
    quality: &CellOnFileQuality,
    raw_slot: &RawDataSlot,
    verbose: bool,
) -> Result<CompressedSparseData2<u32>, ScDataErrors> {
    let file = File::open(file_path)?;

    check_h5ad_is_raw(&file, raw_slot, None)?;

    let data_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_data())?);
    let indices_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_indices())?);
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;

    // Read indptr first (small array)
    let indptr_raw = read_indptr(&indptr_ds)?;

    let mut new_data: Vec<u32> = Vec::new();
    let mut new_indices: Vec<usize> = Vec::new();
    let mut new_indptr: Vec<usize> = Vec::with_capacity(quality.genes_to_keep.len() + 1);
    new_indptr.push(0);
    let cell_map = dense_index_map(&quality.cells_to_keep);

    let total_genes = quality.genes_to_keep.len();

    if verbose {
        println!(
            "  Processing {} genes in chunks...",
            total_genes.separate_with_underscores()
        );
    }

    // Process genes in chunks to reduce memory usage
    const CHUNK_SIZE: usize = 1000;

    for (chunk_idx, gene_chunk) in quality.genes_to_keep.chunks(CHUNK_SIZE).enumerate() {
        if verbose && chunk_idx % 10 == 0 {
            let processed = chunk_idx * CHUNK_SIZE;
            println!(
                "   Processed {} / {} genes",
                processed.min(total_genes).separate_with_underscores(),
                total_genes.separate_with_underscores()
            );
        }

        // calculate range for this chunk
        let start_pos = gene_chunk.iter().map(|&g| indptr_raw[g]).min().unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1])
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

        // read only the data range needed for this chunk
        let chunk_data: Vec<f32> = data_rd.read(start_pos..end_pos)?;
        let chunk_indices: Vec<u32> = indices_rd.read(start_pos..end_pos)?;

        // process each gene in the chunk
        for &gene_idx in gene_chunk {
            let gene_start = indptr_raw[gene_idx];
            let gene_end = indptr_raw[gene_idx + 1];

            for idx in gene_start..gene_end {
                let local_idx = idx - start_pos;
                let old_cell_idx = chunk_indices[local_idx] as usize;

                let new_cell_idx = cell_map.get(old_cell_idx).copied().unwrap_or(INDEX_DROPPED);
                if new_cell_idx != INDEX_DROPPED {
                    new_data.push(chunk_data[local_idx] as u32);
                    new_indices.push(new_cell_idx as usize);
                }
            }
            new_indptr.push(new_data.len());
        }
    }

    if verbose {
        println!(
            "   Processed {} / {} genes (complete)",
            total_genes.separate_with_underscores(),
            total_genes.separate_with_underscores()
        );
    }

    let shape = (quality.genes_to_keep.len(), quality.cells_to_keep.len());

    Ok(CompressedSparseData2 {
        data: new_data,
        indices: new_indices.index_cast(),
        indptr: new_indptr.index_cast(),
        cs_type: CompressedSparseFormat::Csr,
        data_2: None::<Vec<u32>>,
        shape,
    })
}

/// Helper function that streams CSC data to CSR format on disk
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `bin_path` - Path to the binary file on disk to write to.
/// * `quality` - Information on the which cells and genes to include after a
///   first pass of the file.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot]
/// * `target_size` - Target size for the library normalisation.
///
/// ### Returns
///
/// After writing, information on the CellQuality with cell and gene indices,
/// NNZ and lib size per gene.
pub fn write_h5_csc_to_csr_streaming<P: AsRef<Path>>(
    file_path: P,
    bin_path: P,
    quality: &CellOnFileQuality,
    raw_slot: &RawDataSlot,
    target_size: f32,
    verbose: bool,
) -> Result<CellQuality, ScDataErrors> {
    let file = File::open(file_path)?;

    check_h5ad_is_raw(&file, raw_slot, None)?;

    let n_cells = quality.cells_to_keep.len();
    let data_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_data())?);
    let indices_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_indices())?);
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;
    let indptr_raw = read_indptr(&indptr_ds)?;

    // accumulate cells in memory (necessary for CSR)
    // (gene_index, raw_count) - gene index is u32 to support >65k features
    let mut cell_data: Vec<Vec<(u32, u32)>> = vec![Vec::new(); quality.cells_to_keep.len()];
    let cell_map = dense_index_map(&quality.cells_to_keep);

    const GENE_CHUNK_SIZE: usize = 5000;
    let total_genes = quality.genes_to_keep.len();

    if verbose {
        println!(
            "  Processing {} genes in chunks...",
            total_genes.separate_with_underscores()
        );
    }

    for (chunk_idx, gene_chunk) in quality.genes_to_keep.chunks(GENE_CHUNK_SIZE).enumerate() {
        if verbose && chunk_idx % 10 == 0 {
            let processed = chunk_idx * GENE_CHUNK_SIZE;
            println!(
                "   Processed {} / {} genes",
                processed.min(total_genes).separate_with_underscores(),
                total_genes.separate_with_underscores()
            );
        }

        let start_pos = gene_chunk.iter().map(|&g| indptr_raw[g]).min().unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1])
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

        let chunk_data: Vec<f32> = data_rd.read(start_pos..end_pos)?;
        let chunk_indices: Vec<u32> = indices_rd.read(start_pos..end_pos)?;

        for &gene_idx in gene_chunk {
            let new_gene_idx = quality.gene_old_to_new[&gene_idx] as u32;
            let gene_start = indptr_raw[gene_idx];
            let gene_end = indptr_raw[gene_idx + 1];

            for idx in gene_start..gene_end {
                let local_idx = idx - start_pos;
                let old_cell_idx = chunk_indices[local_idx] as usize;

                let new_cell_idx = cell_map.get(old_cell_idx).copied().unwrap_or(INDEX_DROPPED);
                if new_cell_idx != INDEX_DROPPED {
                    let raw_count = chunk_data[local_idx] as u32;
                    cell_data[new_cell_idx as usize].push((new_gene_idx, raw_count));
                }
            }
        }
    }

    if verbose {
        println!(
            "   Processed {} / {} genes (complete)",
            total_genes.separate_with_underscores(),
            total_genes.separate_with_underscores()
        );
        println!(
            "  Writing {} cells to disk...",
            n_cells.separate_with_underscores()
        );
    }

    // now write cells with correct normalisation
    let mut writer = CellGeneSparseWriter::new(
        bin_path,
        true,
        n_cells,
        quality.genes_to_keep.len(),
        target_size,
    )?;
    let (nnz, lib_size) = write_cell_rows(&mut cell_data, 0, target_size, &mut writer)?;

    writer.finalise()?;

    if verbose {
        println!(
            "   Written {} / {} cells (complete).",
            n_cells.separate_with_underscores(),
            n_cells.separate_with_underscores()
        );
    }

    Ok(CellQuality {
        cell_indices: Vec::new(),
        gene_indices: Vec::new(),
        lib_size,
        nnz,
    })
}

/////////
// CSR //
/////////

/// Get the cell quality data from a CSR file
///
/// This file assumes that the rows are representing the cells and the columns
/// the genes and the data was stored in CSR type format.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file
/// * `shape` - Tuple with `(no_cells, no_genes)`.
/// * `cell_quality` - Structure defining the minimum quality values that are
///   expected here.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot].
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// `CellOnFileQuality` structure that contains all of the information about
/// which cells and genes to include.
pub fn parse_h5_csr_quality<P: AsRef<Path>>(
    file_path: P,
    shape: (usize, usize),
    raw_slot: &RawDataSlot,
    cell_quality: &MinCellQuality,
    verbose: bool,
) -> Result<CellOnFileQuality, ScDataErrors> {
    let file_path = file_path.as_ref();

    if verbose {
        println!(
            "  Reading CSR matrix structure (shape: {} x {} )...",
            shape.0.separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
    }

    let file = File::open(file_path)?;

    check_h5ad_is_raw(&file, raw_slot, None)?;

    let indptr = read_indptr(&file.dataset(raw_slot.get_indptr())?)?;

    const CELL_CHUNK_SIZE: usize = 10000;
    let chunks: Vec<usize> = (0..shape.0).step_by(CELL_CHUNK_SIZE).collect();
    let num_chunks = chunks.len();

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

    let first_pass_time = Instant::now();
    let completed_chunks = Arc::new(AtomicUsize::new(0));
    let report_interval = (num_chunks / 10).max(1);

    let indices_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_indices())?);

    let gene_counts: Vec<Vec<usize>> = chunks
        .par_iter()
        .map(|&chunk_start_cell| {
            let mut local_counts = vec![0usize; shape.1];

            let chunk_end_cell = (chunk_start_cell + CELL_CHUNK_SIZE).min(shape.0) - 1;
            let data_start = indptr[chunk_start_cell];
            let data_end = indptr[chunk_end_cell + 1];

            if data_start >= data_end {
                return local_counts;
            }

            let Ok(chunk_indices) = indices_rd.read::<u32>(data_start..data_end) else {
                return local_counts;
            };

            for cell_idx in chunk_start_cell..=chunk_end_cell {
                let cell_data_start = indptr[cell_idx] - data_start;
                let cell_data_end = indptr[cell_idx + 1] - data_start;

                for local_idx in cell_data_start..cell_data_end {
                    let gene_idx = chunk_indices[local_idx] as usize;
                    if gene_idx < shape.1 {
                        local_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 * 10.0).round() as usize) * 10;
                    println!(
                        "  Processed {}% of chunks ({}/{})",
                        progress, completed, num_chunks
                    );
                }
            }

            local_counts
        })
        .collect();

    let mut no_cells_exp_gene = vec![0usize; shape.1];
    for local_counts in gene_counts {
        for (i, count) in local_counts.into_iter().enumerate() {
            no_cells_exp_gene[i] += count;
        }
    }

    if verbose {
        let max_expr = no_cells_exp_gene.iter().max().unwrap_or(&0);
        let min_expr = no_cells_exp_gene.iter().min().unwrap_or(&0);
        let avg_expr = no_cells_exp_gene
            .iter()
            .sum::<usize>()
            .checked_div(shape.1)
            .unwrap_or(0);

        println!(
            "  Gene expression stats: min = {} | max = {} | avg = {} cells per gene",
            min_expr.separate_with_underscores(),
            max_expr.separate_with_underscores(),
            avg_expr.separate_with_underscores()
        );
    }

    let genes_to_keep: Vec<usize> = (0..shape.1)
        .filter(|&i| no_cells_exp_gene[i] >= cell_quality.min_cells)
        .collect();

    let first_pass_elapsed = first_pass_time.elapsed();

    if verbose {
        println!("First pass done: {:.2?}", first_pass_elapsed);
        println!(
            "  Genes passing filter: {} / {}",
            genes_to_keep.len().separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
        println!("Second pass - cell statistics:");
    }

    let mut genes_to_keep_lookup = vec![false; shape.1];
    for &gene_idx in &genes_to_keep {
        genes_to_keep_lookup[gene_idx] = true;
    }

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

    let data_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_data())?);

    let cell_stats: Vec<(Vec<usize>, Vec<f32>)> = chunks
        .par_iter()
        .map(|&chunk_start_cell| {
            let chunk_end_cell = (chunk_start_cell + CELL_CHUNK_SIZE).min(shape.0) - 1;
            let mut local_unique = vec![0usize; chunk_end_cell - chunk_start_cell + 1];
            let mut local_lib_size = vec![0.0f32; chunk_end_cell - chunk_start_cell + 1];

            let data_start = indptr[chunk_start_cell];
            let data_end = indptr[chunk_end_cell + 1];

            if data_start >= data_end {
                return (local_unique, local_lib_size);
            }

            let (Ok(chunk_data), Ok(chunk_indices)) = (
                data_rd.read::<f32>(data_start..data_end),
                indices_rd.read::<u32>(data_start..data_end),
            ) else {
                return (local_unique, local_lib_size);
            };

            for cell_idx in chunk_start_cell..=chunk_end_cell {
                let cell_data_start = indptr[cell_idx] - data_start;
                let cell_data_end = indptr[cell_idx + 1] - data_start;
                let local_cell_idx = cell_idx - chunk_start_cell;

                for local_idx in cell_data_start..cell_data_end {
                    let gene_idx = chunk_indices[local_idx] as usize;
                    if genes_to_keep_lookup[gene_idx] {
                        local_unique[local_cell_idx] += 1;
                        local_lib_size[local_cell_idx] += chunk_data[local_idx];
                    }
                }
            }

            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 * 10.0).round() as usize) * 10;
                    println!(
                        "  Processed {}% of chunks ({}/{})",
                        progress, completed, num_chunks
                    );
                }
            }

            (local_unique, local_lib_size)
        })
        .collect();

    let mut cell_unique_genes = vec![0usize; shape.0];
    let mut cell_lib_size = vec![0.0f32; shape.0];

    for (chunk_idx, (local_unique, local_lib)) in cell_stats.into_iter().enumerate() {
        let chunk_start = chunks[chunk_idx];
        for (i, (unique, lib)) in local_unique.into_iter().zip(local_lib).enumerate() {
            cell_unique_genes[chunk_start + i] = unique;
            cell_lib_size[chunk_start + i] = lib;
        }
    }

    if verbose {
        let max_genes = cell_unique_genes.iter().max().unwrap_or(&0);
        let min_genes = cell_unique_genes.iter().min().unwrap_or(&0);
        let max_lib = cell_lib_size.iter().fold(0.0f32, |a, &b| a.max(b));
        let min_lib = cell_lib_size.iter().fold(f32::INFINITY, |a, &b| a.min(b));

        println!(
            "  Cell stats: genes per cell: min = {} | max = {}",
            min_genes.separate_with_underscores(),
            max_genes.separate_with_underscores()
        );
        println!(
            "  Cell stats: library size: min = {:.1} | max = {:.1}",
            min_lib, max_lib
        );
    }

    let cells_to_keep: Vec<usize> = (0..shape.0)
        .filter(|&i| {
            cell_unique_genes[i] >= cell_quality.min_unique_genes
                && cell_lib_size[i] >= cell_quality.min_lib_size as f32
        })
        .collect();

    let second_pass_elapsed = second_pass_time.elapsed();

    if verbose {
        println!("Second pass done: {:.2?}", second_pass_elapsed);
        println!(
            "  Cells passing filter: {} / {}",
            cells_to_keep.len().separate_with_underscores(),
            shape.0.separate_with_underscores()
        );
    }

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

    Ok(file_quality_data)
}

/// Helper function that reads in full CSR data from an h5 file
///
/// The function assumes that the data is stored as cells x genes.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `quality` - Information on the which cells and genes to include after a
///   first pass of the file.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot].
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// The `CompressedSparseData2` in CSR format with the counts stored as u32.
pub fn read_h5ad_x_data_csr<P: AsRef<Path>>(
    file_path: P,
    quality: &CellOnFileQuality,
    raw_slot: &RawDataSlot,
    verbose: bool,
) -> Result<CompressedSparseData2<u32>, ScDataErrors> {
    let file = File::open(file_path)?;

    let data_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_data())?);
    let indices_rd = H5ChunkedReader::new(file.dataset(raw_slot.get_indices())?);
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;

    let indptr_raw = read_indptr(&indptr_ds)?;

    // Build CSC format directly: indptr = cells, indices = genes. Sized from
    // the kept cells' stored values, an upper bound before gene filtering, so
    // neither vector reallocates.
    let max_nnz: usize = quality
        .cells_to_keep
        .iter()
        .map(|&c| indptr_raw[c + 1] - indptr_raw[c])
        .sum();
    let mut new_data: Vec<u32> = Vec::with_capacity(max_nnz);
    let mut new_indices: Vec<u32> = Vec::with_capacity(max_nnz);
    let mut cell_data: Vec<(usize, u32)> = Vec::new();
    let mut new_indptr: Vec<usize> = Vec::with_capacity(quality.cells_to_keep.len() + 1);
    new_indptr.push(0);
    let gene_map = dense_index_map(&quality.genes_to_keep);

    let total_cells = quality.cells_to_keep.len();

    let start_write = Instant::now();

    if verbose {
        println!(
            "  Processing {} cells in chunks (CSC format)...",
            total_cells.separate_with_underscores()
        );
    }

    const CHUNK_SIZE: usize = 1000;

    for (chunk_idx, cell_chunk) in quality.cells_to_keep.chunks(CHUNK_SIZE).enumerate() {
        if verbose && chunk_idx % 100 == 0 {
            let processed = chunk_idx * CHUNK_SIZE;
            println!(
                "   Processed {} / {} cells",
                processed.min(total_cells).separate_with_underscores(),
                total_cells.separate_with_underscores()
            );
        }

        let start_pos = cell_chunk.iter().map(|&c| indptr_raw[c]).min().unwrap_or(0);
        let end_pos = cell_chunk
            .iter()
            .map(|&c| indptr_raw[c + 1])
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            // Add empty cells to maintain indptr structure
            for _ in cell_chunk {
                new_indptr.push(new_data.len());
            }
            continue;
        }

        let chunk_data: Vec<f32> = data_rd.read(start_pos..end_pos)?;
        let chunk_indices: Vec<u32> = indices_rd.read(start_pos..end_pos)?;

        for &old_cell_idx in cell_chunk {
            let cell_start = indptr_raw[old_cell_idx];
            let cell_end = indptr_raw[old_cell_idx + 1];

            // Collect this cell's data with gene filtering
            cell_data.clear();

            for idx in cell_start..cell_end {
                let local_idx = idx - start_pos;
                let old_gene_idx = chunk_indices[local_idx] as usize;

                let new_gene_idx = gene_map.get(old_gene_idx).copied().unwrap_or(INDEX_DROPPED);
                if new_gene_idx != INDEX_DROPPED {
                    cell_data.push((new_gene_idx as usize, chunk_data[local_idx] as u32));
                }
            }

            // Sort by gene index to maintain consistent ordering
            cell_data.sort_by_key(|&(gene_idx, _)| gene_idx);

            // Add to CSC arrays
            for &(gene_idx, value) in &cell_data {
                new_data.push(value);
                new_indices.push(gene_idx as u32);
            }

            new_indptr.push(new_data.len());
        }
    }

    let end_write = start_write.elapsed();

    if verbose {
        println!(
            "   Processed {} / {} cells (complete) in {:.2?}.",
            total_cells.separate_with_underscores(),
            total_cells.separate_with_underscores(),
            end_write
        );
    }

    let shape = (quality.genes_to_keep.len(), quality.cells_to_keep.len());

    Ok(CompressedSparseData2 {
        data: new_data,
        indices: new_indices,
        indptr: new_indptr.index_cast(),
        cs_type: CompressedSparseFormat::Csc, // Return CSC format
        data_2: None::<Vec<u32>>,
        shape,
    })
}

/// Slices the h5 reader thread may run ahead of the cell builder.
///
/// Two keeps the reader busy while one slice is being built and one is queued,
/// and caps the queued raw data at two slices.
const H5_READ_AHEAD: usize = 2;

/// Write cells from cell-major compressed h5 data into the cell-based binary.
///
/// Shared by the h5ad CSR and the 10x paths, which store the same layout: one
/// `indptr` run per cell over `data` and `indices`. A dedicated thread reads
/// the next [`H5_CELL_SLICE_SIZE`] cells from the h5 file (HDF5 decompression
/// is single-threaded) while the current slice is filtered, normalised,
/// serialised and lz4-compressed in parallel. Payloads are then appended in
/// cell order.
///
/// ### Params
///
/// * `file_path` - Path to the h5 file
/// * `data_path` - Dataset path of the stored values
/// * `indices_path` - Dataset path of the gene indices
/// * `indptr` - Per-cell offsets into `data` and `indices`
/// * `cells_to_keep` - File-local cells to write, in output order
/// * `gene_map` - Dense file-local to output gene map, see [`dense_index_map`]
/// * `cell_offset` - Output index of the first written cell
/// * `target_size` - Target size for the library normalisation
/// * `writer` - Cell-based writer to append to
/// * `verbose` - Controls verbosity
///
/// ### Returns
///
/// `(nnz, lib_size)` per written cell.
#[allow(clippy::too_many_arguments)]
pub(crate) fn write_h5_cell_major_cells(
    file_path: &Path,
    data_path: &str,
    indices_path: &str,
    indptr: &[usize],
    cells_to_keep: &[usize],
    gene_map: &[u32],
    cell_offset: usize,
    target_size: f32,
    writer: &mut CellGeneSparseWriter,
    verbose: bool,
) -> Result<(Vec<usize>, Vec<usize>), ScDataErrors> {
    let total_cells = cells_to_keep.len();
    let batches: Vec<&[usize]> = cells_to_keep.chunks(H5_CELL_SLICE_SIZE).collect();
    let num_batches = batches.len();

    let mut nnz = Vec::with_capacity(total_cells);
    let mut lib_size = Vec::with_capacity(total_cells);

    if verbose {
        println!(
            "  Processing {} cells in batches of {}...",
            total_cells.separate_with_underscores(),
            H5_CELL_SLICE_SIZE.separate_with_underscores()
        );
    }

    type Slice = (usize, Vec<f32>, Vec<u32>);

    std::thread::scope(|scope| -> Result<(), ScDataErrors> {
        let (tx, rx) = std::sync::mpsc::sync_channel::<Result<Slice, ScDataErrors>>(H5_READ_AHEAD);
        let batches_ref = &batches;

        scope.spawn(move || {
            let read_all = || -> Result<(), ScDataErrors> {
                let file = File::open(file_path)?;
                let data_rd = H5ChunkedReader::new(file.dataset(data_path)?);
                let indices_rd = H5ChunkedReader::new(file.dataset(indices_path)?);
                for batch in batches_ref {
                    let start = batch.iter().map(|&c| indptr[c]).min().unwrap_or(0);
                    let end = batch.iter().map(|&c| indptr[c + 1]).max().unwrap_or(0);
                    let slice = if start < end {
                        (
                            start,
                            data_rd.read(start..end)?,
                            indices_rd.read(start..end)?,
                        )
                    } else {
                        (start, Vec::new(), Vec::new())
                    };
                    // the builder only hangs up after an error of its own
                    if tx.send(Ok(slice)).is_err() {
                        return Ok(());
                    }
                }
                Ok(())
            };
            if let Err(e) = read_all() {
                let _ = tx.send(Err(e));
            }
        });

        let mut offset = 0;
        for (batch_idx, batch) in batches.iter().enumerate() {
            let (start, data, indices) = rx
                .recv()
                .expect("h5 reader thread sends one slice per batch or an error")?;

            let built = batch
                .par_iter()
                .enumerate()
                .map_init(Vec::<(u32, u32)>::new, |buf, (k, &old)| {
                    buf.clear();
                    for idx in indptr[old] - start..indptr[old + 1] - start {
                        let new = gene_map
                            .get(indices[idx] as usize)
                            .copied()
                            .unwrap_or(INDEX_DROPPED);
                        if new != INDEX_DROPPED {
                            buf.push((new, data[idx] as u32));
                        }
                    }
                    compress_cell_row(buf, cell_offset + offset + k, target_size, writer.format())
                })
                .collect::<Result<Vec<_>, ScDataErrors>>()?;

            let mut payloads = Vec::with_capacity(built.len());
            for (nnz_i, payload) in built {
                nnz.push(nnz_i);
                lib_size.push(payload.library_size);
                payloads.push(payload);
            }
            writer.write_compressed_cell_chunks(&payloads)?;
            offset += batch.len();

            if verbose
                && (batch_idx % (num_batches / 10).max(1) == 0 || batch_idx == num_batches - 1)
            {
                println!(
                    "  Processed {} / {} cells",
                    offset.separate_with_underscores(),
                    total_cells.separate_with_underscores()
                );
            }
        }
        Ok(())
    })?;

    Ok((nnz, lib_size))
}

/// Stream h5 CSR data directly to disk with batched reading
///
/// Memory stays bounded by a few slices of [`H5_CELL_SLICE_SIZE`] cells, see
/// [`write_h5_cell_major_cells`].
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file
/// * `bin_path` - Path to the to-be-written binary file for the cell data
/// * `quality` - Information on the which cells and genes to include after a
///   first pass of the file.
/// * `cell_qc` - Structure containing the information on which minimum criteria
///   cells and genes need to pass.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot].
/// * `verbose` - Controls verbosity of the function
///
/// ### Returns
///
/// After writing, information on the CellQuality with cell and gene indices,
/// NNZ and lib size per gene.
pub fn write_h5_csr_streaming<P: AsRef<Path>>(
    file_path: P,
    bin_path: P,
    quality: &CellOnFileQuality,
    cell_qc: MinCellQuality,
    raw_slot: &RawDataSlot,
    verbose: bool,
) -> Result<CellQuality, ScDataErrors> {
    let indptr_raw = {
        let file = File::open(&file_path)?;
        read_indptr(&file.dataset(raw_slot.get_indptr())?)?
    };

    let mut writer = CellGeneSparseWriter::new(
        bin_path,
        true,
        quality.cells_to_keep.len(),
        quality.genes_to_keep.len(),
        cell_qc.target_size,
    )?;

    let start_write = Instant::now();

    let (nnz, lib_size) = write_h5_cell_major_cells(
        file_path.as_ref(),
        raw_slot.get_data(),
        raw_slot.get_indices(),
        &indptr_raw,
        &quality.cells_to_keep,
        &dense_index_map(&quality.genes_to_keep),
        0,
        cell_qc.target_size,
        &mut writer,
        verbose,
    )?;

    writer.finalise()?;

    if verbose {
        println!("  Writing complete in {:.2?}", start_write.elapsed());
    }

    Ok(CellQuality {
        cell_indices: Vec::new(),
        gene_indices: Vec::new(),
        lib_size,
        nnz,
    })
}

/////////////////////////
// Special case: dense //
/////////////////////////

/// Error helper for unsupported `DENSE_COL`.
///
/// ### Returns
///
/// Returns an [ScDataErrors::UnsupportH5ADFormat]
pub fn dense_col_not_supported<T>() -> Result<T, ScDataErrors> {
    Err(ScDataErrors::UnsupportH5ADFormat(
        "dense_col ingestion is not yet implemented".to_string(),
    ))
}

/// First-pass QC scan for `DENSE_ROW` (cells-as-rows) storage.
///
/// Reads contiguous row-chunks in parallel; per chunk it tracks per-cell
/// `(lib_size, unique_genes)` and per-gene `nnz`. Applies the `MinCellQuality`
/// thresholds to produce a `CellOnFileQuality`.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `shape` - Array shape
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot].
/// * `cell_quality` - Information on the which cells and genes to include after
///   a first pass of the file.
/// * `verbose` - Controls verbosity of the function
///
/// ### Returns
///
/// The [CellOnFileQuality]
pub fn parse_h5_dense_row_quality<P: AsRef<Path>>(
    file_path: P,
    shape: (usize, usize),
    raw_slot: &RawDataSlot,
    cell_quality: &MinCellQuality,
    verbose: bool,
) -> Result<CellOnFileQuality, ScDataErrors> {
    let file_path = file_path.as_ref();

    if verbose {
        println!(
            "  Reading dense (row-major) matrix structure (shape: {} x {} )...",
            shape.0.separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
    }

    {
        let file = File::open(file_path)?;
        check_h5ad_dense_is_raw(&file, raw_slot, None)?;
    }

    const CELL_CHUNK_SIZE: usize = 10_000;
    let chunk_starts: Vec<usize> = (0..shape.0).step_by(CELL_CHUNK_SIZE).collect();
    let num_chunks = chunk_starts.len();

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

    let first_pass_time = Instant::now();
    let completed_chunks = Arc::new(AtomicUsize::new(0));
    let report_interval = (num_chunks / 10).max(1);

    // Pass 1: per-gene nnz across all cells
    let gene_counts: Vec<Vec<usize>> = chunk_starts
        .par_iter()
        .map(|&chunk_start| -> Result<Vec<usize>, ScDataErrors> {
            let mut local_counts = vec![0usize; shape.1];

            let file = File::open(file_path)?;
            let ds = file.dataset(raw_slot.get_dense_path())?;

            let chunk_end = (chunk_start + CELL_CHUNK_SIZE).min(shape.0);
            let block: Array2<f32> = ds.read_slice_2d(s![chunk_start..chunk_end, 0..shape.1])?;

            for row in block.outer_iter() {
                for (gene_idx, &val) in row.iter().enumerate() {
                    if val != 0.0 {
                        local_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 * 10.0).round() as usize) * 10;
                    println!(
                        "  Processed {}% of chunks ({}/{})",
                        progress, completed, num_chunks
                    );
                }
            }

            Ok(local_counts)
        })
        .collect::<Result<Vec<_>, _>>()?;

    let mut no_cells_exp_gene = vec![0usize; shape.1];
    for local_counts in gene_counts {
        for (i, c) in local_counts.into_iter().enumerate() {
            no_cells_exp_gene[i] += c;
        }
    }

    if verbose {
        let max_expr = no_cells_exp_gene.iter().max().unwrap_or(&0);
        let min_expr = no_cells_exp_gene.iter().min().unwrap_or(&0);
        let avg_expr = no_cells_exp_gene
            .iter()
            .sum::<usize>()
            .checked_div(shape.1)
            .unwrap_or(0);
        println!(
            "  Gene expression stats: min = {} | max = {} | avg = {} cells per gene",
            min_expr.separate_with_underscores(),
            max_expr.separate_with_underscores(),
            avg_expr.separate_with_underscores()
        );
    }

    let genes_to_keep: Vec<usize> = (0..shape.1)
        .filter(|&i| no_cells_exp_gene[i] >= cell_quality.min_cells)
        .collect();

    if verbose {
        println!("First pass done: {:.2?}", first_pass_time.elapsed());
        println!(
            "  Genes passing filter: {} / {}",
            genes_to_keep.len().separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
        println!("Second pass - cell statistics over kept genes:");
    }

    let mut genes_to_keep_lookup = vec![false; shape.1];
    for &g in &genes_to_keep {
        genes_to_keep_lookup[g] = true;
    }

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

    // Pass 2: per-cell stats restricted to kept genes
    let cell_stats: Vec<(Vec<usize>, Vec<f32>)> = chunk_starts
        .par_iter()
        .map(
            |&chunk_start| -> Result<(Vec<usize>, Vec<f32>), ScDataErrors> {
                let chunk_end = (chunk_start + CELL_CHUNK_SIZE).min(shape.0);
                let n = chunk_end - chunk_start;
                let mut local_unique = vec![0usize; n];
                let mut local_lib_size = vec![0.0f32; n];

                let file = File::open(file_path)?;
                let ds = file.dataset(raw_slot.get_dense_path())?;

                let block: Array2<f32> =
                    ds.read_slice_2d(s![chunk_start..chunk_end, 0..shape.1])?;

                for (local_row, row) in block.outer_iter().enumerate() {
                    for (gene_idx, &val) in row.iter().enumerate() {
                        if val != 0.0 && genes_to_keep_lookup[gene_idx] {
                            local_unique[local_row] += 1;
                            local_lib_size[local_row] += val;
                        }
                    }
                }

                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 * 10.0).round() as usize) * 10;
                        println!(
                            "  Processed {}% of chunks ({}/{})",
                            progress, completed, num_chunks
                        );
                    }
                }

                Ok((local_unique, local_lib_size))
            },
        )
        .collect::<Result<Vec<_>, _>>()?;

    let mut cell_unique_genes = vec![0usize; shape.0];
    let mut cell_lib_size = vec![0.0f32; shape.0];

    for (chunk_idx, (local_unique, local_lib)) in cell_stats.into_iter().enumerate() {
        let chunk_start = chunk_starts[chunk_idx];
        for (i, (u, l)) in local_unique.into_iter().zip(local_lib).enumerate() {
            cell_unique_genes[chunk_start + i] = u;
            cell_lib_size[chunk_start + i] = l;
        }
    }

    if verbose {
        let max_genes = cell_unique_genes.iter().max().unwrap_or(&0);
        let min_genes = cell_unique_genes.iter().min().unwrap_or(&0);
        let max_lib = cell_lib_size.iter().fold(0.0f32, |a, &b| a.max(b));
        let min_lib = cell_lib_size.iter().fold(f32::INFINITY, |a, &b| a.min(b));
        println!(
            "  Cell stats: genes per cell: min = {} | max = {}",
            min_genes.separate_with_underscores(),
            max_genes.separate_with_underscores()
        );
        println!(
            "  Cell stats: library size: min = {:.1} | max = {:.1}",
            min_lib, max_lib
        );
        println!("Second pass done: {:.2?}", second_pass_time.elapsed());
    }

    let cells_to_keep: Vec<usize> = (0..shape.0)
        .filter(|&i| {
            cell_unique_genes[i] >= cell_quality.min_unique_genes
                && cell_lib_size[i] >= cell_quality.min_lib_size as f32
        })
        .collect();

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

/// Helper function that reads in full data from a dense, row-major (cells x
/// genes) h5ad file.
///
/// Mirrors [read_h5ad_x_data_csr] but reads contiguous row-blocks from the
/// dense dataset and filters non-zeros on the fly.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `quality` - Information on the which cells and genes to include after a
///   first pass of the file.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot].
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// The `CompressedSparseData2` with the counts stored as u32.
pub fn read_h5ad_x_data_dense_row<P: AsRef<Path>>(
    file_path: P,
    quality: &CellOnFileQuality,
    raw_slot: &RawDataSlot,
    verbose: bool,
) -> Result<CompressedSparseData2<u32>, ScDataErrors> {
    let file = File::open(file_path)?;
    let ds = file.dataset(raw_slot.get_dense_path())?;
    let no_genes_total = ds.shape().get(1).copied().unwrap_or(0);

    let mut new_data: Vec<u32> = Vec::new();
    let mut new_indices: Vec<usize> = Vec::new();
    let mut new_indptr: Vec<usize> = Vec::with_capacity(quality.cells_to_keep.len() + 1);
    new_indptr.push(0);
    let gene_map = dense_index_map(&quality.genes_to_keep);

    let total_cells = quality.cells_to_keep.len();

    let start_write = Instant::now();

    if verbose {
        println!(
            "  Processing {} cells in chunks (dense row-major)...",
            total_cells.separate_with_underscores()
        );
    }

    const CHUNK_SIZE: usize = 1000;

    for (chunk_idx, cell_chunk) in quality.cells_to_keep.chunks(CHUNK_SIZE).enumerate() {
        if verbose && chunk_idx % 100 == 0 {
            let processed = chunk_idx * CHUNK_SIZE;
            println!(
                "   Processed {} / {} cells",
                processed.min(total_cells).separate_with_underscores(),
                total_cells.separate_with_underscores()
            );
        }

        // `chunks` never yields an empty slice.
        let row_start = *cell_chunk
            .iter()
            .min()
            .expect("batch non-empty by construction");
        let row_end = *cell_chunk
            .iter()
            .max()
            .expect("batch non-empty by construction")
            + 1;
        let block = ds.read_slice_2d::<f32, _>((row_start..row_end, 0..no_genes_total))?;

        for &old_cell_idx in cell_chunk {
            let row = block.row(old_cell_idx - row_start);

            // Dense iteration already yields ascending gene_idx; no sort needed.
            for (gene_idx, &val) in row.iter().enumerate() {
                if val == 0.0 {
                    continue;
                }
                let new_gene_idx = gene_map.get(gene_idx).copied().unwrap_or(INDEX_DROPPED);
                if new_gene_idx != INDEX_DROPPED {
                    new_data.push(val as u32);
                    new_indices.push(new_gene_idx as usize);
                }
            }

            new_indptr.push(new_data.len());
        }
    }

    let end_write = start_write.elapsed();

    if verbose {
        println!(
            "   Processed {} / {} cells (complete) in {:.2?}.",
            total_cells.separate_with_underscores(),
            total_cells.separate_with_underscores(),
            end_write
        );
    }

    let shape = (quality.genes_to_keep.len(), quality.cells_to_keep.len());

    Ok(CompressedSparseData2 {
        data: new_data,
        indices: new_indices.index_cast(),
        indptr: new_indptr.index_cast(),
        cs_type: CompressedSparseFormat::Csc,
        data_2: None::<Vec<u32>>,
        shape,
    })
}

/// Write the dense row-major h5 in a streaming fashion on disk
///
/// Iterates kept cells in row-chunks, scans non-zeros, drops entries pointing
/// at filtered genes, remaps gene indices via `quality.gene_old_to_new`, and
/// emits a `CsrCellChunk` per kept cell.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file
/// * `bin_path` - Path to the to-be-written binary file for the cell data
/// * `quality` - Information on the which cells and genes to include after a
///   first pass of the file.
/// * `cell_qc` - Structure containing the information on which minimum criteria
///   cells and genes need to pass.
/// * `raw_slot` - Where to find the raw counts in the h5ad object, see
///   [RawDataSlot].
/// * `verbose` - Controls verbosity of the function
///
/// ### Returns
///
/// The [CellQuality] or errors
pub fn write_h5_dense_row_streaming<P: AsRef<Path>>(
    file_path: P,
    bin_path: P,
    quality: &CellOnFileQuality,
    cell_qc: MinCellQuality,
    raw_slot: &RawDataSlot,
    verbose: bool,
) -> Result<CellQuality, ScDataErrors> {
    let file = File::open(&file_path)?;
    let ds = file.dataset(raw_slot.get_dense_path())?;
    let no_genes_total = ds.shape().get(1).copied().unwrap_or(0);

    let mut writer = CellGeneSparseWriter::new(
        bin_path,
        true,
        quality.cells_to_keep.len(),
        quality.genes_to_keep.len(),
        cell_qc.target_size,
    )?;

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

    let total_cells = quality.cells_to_keep.len();
    let num_batches = total_cells.div_ceil(H5_CELL_SLICE_SIZE);

    if verbose {
        println!(
            "  Processing {} cells in batches of {}...",
            total_cells.separate_with_underscores(),
            H5_CELL_SLICE_SIZE.separate_with_underscores()
        );
    }

    let start_write = Instant::now();

    let gene_map = dense_index_map(&quality.genes_to_keep);
    let mut offset = 0;

    for (batch_idx, cell_batch) in quality.cells_to_keep.chunks(H5_CELL_SLICE_SIZE).enumerate() {
        if verbose
            && (batch_idx.is_multiple_of((num_batches / 10).max(1)) || batch_idx == num_batches - 1)
        {
            let progress = ((batch_idx as f64 / num_batches as f64 * 10.0).round() as usize) * 10;
            let processed = (batch_idx + 1) * H5_CELL_SLICE_SIZE;
            println!(
                "  Processed {}% ({} / {} cells)",
                progress,
                processed.min(total_cells).separate_with_underscores(),
                total_cells.separate_with_underscores()
            );
        }

        // Batched read: span = [min(batch), max(batch)+1]. Kept cells should be
        // sorted (they came from an ordered scan), but be defensive.
        let row_start = *cell_batch
            .iter()
            .min()
            .expect("batch non-empty by construction");
        let row_end = *cell_batch
            .iter()
            .max()
            .expect("batch non-empty by construction")
            + 1;
        let block: Array2<f32> = ds.read_slice_2d(s![row_start..row_end, 0..no_genes_total])?;

        let built = cell_batch
            .par_iter()
            .enumerate()
            .map_init(Vec::<(u32, u32)>::new, |buf, (k, &old_cell_idx)| {
                buf.clear();
                let row = block.row(old_cell_idx - row_start);
                for (gene_idx, &val) in row.iter().enumerate() {
                    let new_gene = gene_map.get(gene_idx).copied().unwrap_or(INDEX_DROPPED);
                    if val != 0.0 && new_gene != INDEX_DROPPED {
                        buf.push((new_gene, val as u32));
                    }
                }
                compress_cell_row(buf, offset + k, cell_qc.target_size, writer.format())
            })
            .collect::<Result<Vec<_>, ScDataErrors>>()?;

        let mut payloads = Vec::with_capacity(built.len());
        for (nnz_i, payload) in built {
            nnz.push(nnz_i);
            lib_size.push(payload.library_size);
            payloads.push(payload);
        }
        writer.write_compressed_cell_chunks(&payloads)?;
        offset += cell_batch.len();
    }

    writer.finalise()?;

    if verbose {
        println!("  Writing complete in {:.2?}", start_write.elapsed());
    }

    Ok(CellQuality {
        cell_indices: Vec::new(),
        gene_indices: Vec::new(),
        lib_size,
        nnz,
    })
}

////////////////////////////////////
// Special case: only norm counts //
////////////////////////////////////

//////////////////
// QC functions //
//////////////////

/// Get the cell quality data from a CSC file (with only normalised counts)
///
/// This is a specialised function that will deal with cases in which authors
/// only provide normalised counts.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file
/// * `shape` - Tuple with `(no_cells, no_genes)`.
/// * `lib_sizes` - Library size per cell/spot from the observations.
/// * `cell_quality` - Structure defining the minimum quality values that are
///   expected here.
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// `CellOnFileQuality` structure that contains all of the information about
/// which cells and genes to include.
pub fn parse_h5_normalised_quality_csc<P: AsRef<Path>>(
    file_path: P,
    shape: (usize, usize),
    lib_sizes: &[f32],
    cell_quality: &MinCellQuality,
    verbose: bool,
) -> hdf5::Result<CellOnFileQuality> {
    let file_path = file_path.as_ref();

    if verbose {
        println!(
            "  Reading CSC matrix structure (shape: {} x {})...",
            shape.0.separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
    }

    let file = hdf5::File::open(file_path)?;
    let indptr = read_indptr(&file.dataset("X/indptr")?)?;

    // NNZ per gene is directly readable from indptr
    let no_cells_exp_gene: Vec<usize> = (0..shape.1).map(|i| indptr[i + 1] - indptr[i]).collect();

    let genes_to_keep: Vec<usize> = (0..shape.1)
        .filter(|&i| no_cells_exp_gene[i] >= cell_quality.min_cells)
        .collect();

    if verbose {
        println!(
            "  Genes passing filter: {} / {}",
            genes_to_keep.len().separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
        println!("  Calculating unique genes per cell...");
    }

    // Still need unique genes per cell - requires scanning indices
    const GENE_CHUNK_SIZE: usize = 10000;
    let mut cell_unique_genes = vec![0usize; shape.0];

    let indices_rd = H5ChunkedReader::new(file.dataset("X/indices")?);
    for gene_chunk in genes_to_keep.chunks(GENE_CHUNK_SIZE) {
        let start_pos = gene_chunk.iter().map(|&g| indptr[g]).min().unwrap_or(0);
        let end_pos = gene_chunk.iter().map(|&g| indptr[g + 1]).max().unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

        let chunk_indices: Vec<u32> = indices_rd
            .read(start_pos..end_pos)
            .map_err(|e| e.to_string())?;

        for &gene_idx in gene_chunk {
            let gene_start = indptr[gene_idx];
            let gene_end = indptr[gene_idx + 1];

            for idx in gene_start..gene_end {
                let local_idx = idx - start_pos;
                let cell_idx = chunk_indices[local_idx] as usize;
                if cell_idx < shape.0 {
                    cell_unique_genes[cell_idx] += 1;
                }
            }
        }
    }

    // QC using obs lib sizes, not summed X values
    let cells_to_keep: Vec<usize> = (0..shape.0)
        .filter(|&i| {
            cell_unique_genes[i] >= cell_quality.min_unique_genes
                && lib_sizes[i] >= cell_quality.min_lib_size as f32
        })
        .collect();

    if verbose {
        println!(
            "  Cells passing filter: {} / {}",
            cells_to_keep.len().separate_with_underscores(),
            shape.0.separate_with_underscores()
        );
    }

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

    Ok(file_quality_data)
}

/// Get the cell quality data from a CSR file (with only normalised counts)
///
/// This is a specialised function that will deal with cases in which authors
/// only provide normalised counts.
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file
/// * `shape` - Tuple with `(no_cells, no_genes)`.
/// * `lib_sizes` - Library size per cell/spot from the observations.
/// * `cell_quality` - Structure defining the minimum quality values that are
///   expected here.
/// * `verbose` - Controls verbosity of the function.
///
/// ### Returns
///
/// `CellOnFileQuality` structure that contains all of the information about
/// which cells and genes to include.
pub fn parse_h5_normalised_quality_csr<P: AsRef<Path>>(
    file_path: P,
    shape: (usize, usize),
    lib_sizes: &[f32],
    cell_quality: &MinCellQuality,
    verbose: bool,
) -> hdf5::Result<CellOnFileQuality> {
    let file_path = file_path.as_ref();

    if verbose {
        println!(
            "  Reading CSR matrix structure (shape: {} x {})...",
            shape.0.separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
    }

    let file = hdf5::File::open(file_path)?;
    let indptr = read_indptr(&file.dataset("X/indptr")?)?;
    let indices_rd = H5ChunkedReader::new(file.dataset("X/indices")?);

    // Unique genes per cell is directly readable from indptr
    let cell_unique_genes: Vec<usize> = (0..shape.0).map(|i| indptr[i + 1] - indptr[i]).collect();

    // Still need cells-per-gene for gene filtering - requires scanning
    const CELL_CHUNK_SIZE: usize = 10000;
    let mut no_cells_exp_gene = vec![0usize; shape.1];

    for chunk_start in (0..shape.0).step_by(CELL_CHUNK_SIZE) {
        let chunk_end = (chunk_start + CELL_CHUNK_SIZE).min(shape.0) - 1;
        let data_start = indptr[chunk_start];
        let data_end = indptr[chunk_end + 1];

        if data_start >= data_end {
            continue;
        }

        let chunk_indices: Vec<u32> = indices_rd
            .read(data_start..data_end)
            .map_err(|e| e.to_string())?;

        for cell_idx in chunk_start..=chunk_end {
            let cell_start = indptr[cell_idx] - data_start;
            let cell_end = indptr[cell_idx + 1] - data_start;
            for local_idx in cell_start..cell_end {
                let gene_idx = chunk_indices[local_idx] as usize;
                if gene_idx < shape.1 {
                    no_cells_exp_gene[gene_idx] += 1;
                }
            }
        }
    }

    let genes_to_keep: Vec<usize> = (0..shape.1)
        .filter(|&i| no_cells_exp_gene[i] >= cell_quality.min_cells)
        .collect();

    if verbose {
        println!(
            "  Genes passing filter: {} / {}",
            genes_to_keep.len().separate_with_underscores(),
            shape.1.separate_with_underscores()
        );
    }

    // QC using obs lib sizes, not summed X values
    let cells_to_keep: Vec<usize> = (0..shape.0)
        .filter(|&i| {
            cell_unique_genes[i] >= cell_quality.min_unique_genes
                && lib_sizes[i] >= cell_quality.min_lib_size as f32
        })
        .collect();

    if verbose {
        println!(
            "  Cells passing filter: {} / {}",
            cells_to_keep.len().separate_with_underscores(),
            shape.0.separate_with_underscores()
        );
    }

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

    Ok(file_quality_data)
}

////////////////////////////
// Reconstruction helpers //
////////////////////////////

/// Reconstruct a raw integer count from a log1p-normalised value
///
/// Reverses: norm = ln1p(x / lib_size * target_size)
///
/// ### Params
///
/// * `norm_val` - The normalised count
/// * `lib_size` - The library size
/// * `target_size` - The target size the authors chose
///
/// ### Returns
///
/// The reconstructed normalised counts as u32
#[inline(always)]
fn reconstruct_raw_count(norm_val: f32, lib_size: f32, target_size: f32) -> u32 {
    if norm_val <= 0.0 || lib_size <= 0.0 {
        return 0;
    }
    let reconstructed = norm_val.exp_m1() * lib_size / target_size;
    reconstructed.round().clamp(0.0, u32::MAX as f32) as u32
}

/// Reconstruct raw counts from normalised CSR data and write to binary
///
/// Specialised function that deals with cases where authors only provide
/// normalised counts
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `bin_path` - Path to the binary file.
/// * `quality` - The quality reference.
/// * `lib_sizes` - Vector of library sizes for reconstructing the raw counts.
/// * `target_size` - The target size the authors chose.
/// * `cell_qc` - MinCellQc parameters you wish to apply
/// * `verbose` - Controls verbosity of the function
///
/// ### Returns
///
/// `CellQuality` structure
fn reconstruct_and_write_csr<P: AsRef<Path>>(
    file_path: P,
    bin_path: P,
    quality: &CellOnFileQuality,
    lib_sizes: &[f32],
    target_size: f32,
    cell_qc: &MinCellQuality,
    verbose: bool,
) -> Result<CellQuality, ScDataErrors> {
    let file = hdf5::File::open(file_path)?;
    let data_rd = H5ChunkedReader::new(file.dataset("X/data")?);
    let indices_rd = H5ChunkedReader::new(file.dataset("X/indices")?);
    let indptr_ds = file.dataset("X/indptr")?;
    let indptr_raw = read_indptr(&indptr_ds)?;

    let mut writer = CellGeneSparseWriter::new(
        bin_path,
        true,
        quality.cells_to_keep.len(),
        quality.genes_to_keep.len(),
        // the norm layer is rebuilt at our target, not the authors' one
        cell_qc.target_size,
    )?;

    let mut lib_size_out = Vec::with_capacity(quality.cells_to_keep.len());
    let mut nnz_out = Vec::with_capacity(quality.cells_to_keep.len());

    let total_cells = quality.cells_to_keep.len();
    // (gene_index, raw_count) - gene index as usize, count as u32
    let gene_map = dense_index_map(&quality.genes_to_keep);
    let cell_map = dense_index_map(&quality.cells_to_keep);
    for (batch_idx, cell_batch) in quality.cells_to_keep.chunks(H5_CELL_SLICE_SIZE).enumerate() {
        if verbose && batch_idx % 10 == 0 {
            let processed = (batch_idx * H5_CELL_SLICE_SIZE).min(total_cells);
            println!(
                "   Processed {} / {} cells",
                processed.separate_with_underscores(),
                total_cells.separate_with_underscores()
            );
        }

        let start_pos = cell_batch.iter().map(|&c| indptr_raw[c]).min().unwrap_or(0);
        let end_pos = cell_batch
            .iter()
            .map(|&c| indptr_raw[c + 1])
            .max()
            .unwrap_or(0);

        // an all-empty batch has nothing to read; every cell range is empty
        let (chunk_data, chunk_indices): (Vec<f32>, Vec<u32>) = if start_pos >= end_pos {
            (Vec::new(), Vec::new())
        } else {
            (
                data_rd.read(start_pos..end_pos)?,
                indices_rd.read(start_pos..end_pos)?,
            )
        };

        let (nnz_batch, lib_batch) = write_cell_chunks_built(
            cell_batch.len(),
            |k| {
                let old_cell_idx = cell_batch[k];
                let cell_start = indptr_raw[old_cell_idx];
                let cell_end = indptr_raw[old_cell_idx + 1];
                let lib_size = lib_sizes[old_cell_idx];

                // (gene_index, raw_count)
                let mut cell_data: Vec<(u32, u32)> = Vec::new();
                for idx in cell_start..cell_end {
                    let local_idx = idx - start_pos;
                    let old_gene_idx = chunk_indices[local_idx] as usize;

                    let new_gene_idx = gene_map.get(old_gene_idx).copied().unwrap_or(INDEX_DROPPED);
                    if new_gene_idx != INDEX_DROPPED {
                        let raw_count =
                            reconstruct_raw_count(chunk_data[local_idx], lib_size, target_size);
                        if raw_count > 0 {
                            cell_data.push((new_gene_idx, raw_count));
                        }
                    }
                }

                if cell_data.windows(2).any(|w| w[0].0 > w[1].0) {
                    cell_data.sort_unstable_by_key(|&(g, _)| g);
                }
                let gene_indices: Vec<u32> = cell_data.iter().map(|&(g, _)| g).collect();
                let gene_counts: Vec<u32> = cell_data.iter().map(|&(_, c)| c).collect();

                CsrCellChunk::from_data(
                    &gene_counts,
                    &gene_indices,
                    cell_map[old_cell_idx] as usize,
                    cell_qc.target_size,
                    true,
                )
            },
            &mut writer,
        )?;
        nnz_out.extend(nnz_batch);
        lib_size_out.extend(lib_batch);
    }

    writer.finalise()?;

    Ok(CellQuality {
        cell_indices: Vec::new(),
        gene_indices: Vec::new(),
        lib_size: lib_size_out,
        nnz: nnz_out,
    })
}

/// Reconstruct raw counts from normalised CSC data and write to binary
///
/// Specialised function that deals with cases where authors only provide
/// normalised counts
///
/// ### Params
///
/// * `file_path` - Path to the h5ad file.
/// * `bin_path` - Path to the binary file.
/// * `quality` - The quality reference.
/// * `lib_sizes` - Vector of library sizes for reconstructing the raw counts.
/// * `target_size` - The target size the authors chose.
/// * `cell_qc` - MinCellQc parameters you wish to apply
/// * `verbose` - Controls verbosity of the function
///
/// ### Returns
///
/// `CellQuality` structure
fn reconstruct_and_write_csc<P: AsRef<Path>>(
    file_path: P,
    bin_path: P,
    quality: &CellOnFileQuality,
    lib_sizes: &[f32],
    target_size: f32,
    cell_qc: &MinCellQuality,
    verbose: bool,
) -> Result<CellQuality, ScDataErrors> {
    let file = hdf5::File::open(file_path)?;
    let data_rd = H5ChunkedReader::new(file.dataset("X/data")?);
    let indices_rd = H5ChunkedReader::new(file.dataset("X/indices")?);
    let indptr_ds = file.dataset("X/indptr")?;
    let indptr_raw = read_indptr(&indptr_ds)?;

    // Accumulate per-cell data in memory - necessary for CSR output
    // (gene_index, raw_count) - gene index as u32 to support >65k features
    let mut cell_data: Vec<Vec<(u32, u32)>> = vec![Vec::new(); quality.cells_to_keep.len()];
    let cell_map = dense_index_map(&quality.cells_to_keep);

    const GENE_CHUNK_SIZE: usize = 5000;
    let total_genes = quality.genes_to_keep.len();

    for (chunk_idx, gene_chunk) in quality.genes_to_keep.chunks(GENE_CHUNK_SIZE).enumerate() {
        if verbose && chunk_idx % 10 == 0 {
            let processed = (chunk_idx * GENE_CHUNK_SIZE).min(total_genes);
            println!(
                "   Processed {} / {} genes",
                processed.separate_with_underscores(),
                total_genes.separate_with_underscores()
            );
        }

        let start_pos = gene_chunk.iter().map(|&g| indptr_raw[g]).min().unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1])
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

        let chunk_data: Vec<f32> = data_rd.read(start_pos..end_pos)?;
        let chunk_indices: Vec<u32> = indices_rd.read(start_pos..end_pos)?;

        for &gene_idx in gene_chunk {
            let new_gene_idx = quality.gene_old_to_new[&gene_idx] as u32;
            let gene_start = indptr_raw[gene_idx];
            let gene_end = indptr_raw[gene_idx + 1];

            for idx in gene_start..gene_end {
                let local_idx = idx - start_pos;
                let old_cell_idx = chunk_indices[local_idx] as usize;

                let new_cell_idx = cell_map.get(old_cell_idx).copied().unwrap_or(INDEX_DROPPED);
                if new_cell_idx != INDEX_DROPPED {
                    let norm_val = chunk_data[local_idx];
                    let lib_size = lib_sizes[old_cell_idx];
                    let raw_count = reconstruct_raw_count(norm_val, lib_size, target_size);
                    if raw_count > 0 {
                        cell_data[new_cell_idx as usize].push((new_gene_idx, raw_count));
                    }
                }
            }
        }
    }

    if verbose {
        println!(
            "   Processed {} / {} genes (complete)",
            total_genes.separate_with_underscores(),
            total_genes.separate_with_underscores()
        );
        println!(
            "  Writing {} cells to disk...",
            quality.cells_to_keep.len().separate_with_underscores()
        );
    }

    let mut writer = CellGeneSparseWriter::new(
        bin_path,
        true,
        quality.cells_to_keep.len(),
        quality.genes_to_keep.len(),
        // the norm layer is rebuilt at our target, not the authors' one
        cell_qc.target_size,
    )?;
    let (nnz_out, lib_size_out) =
        write_cell_rows(&mut cell_data, 0, cell_qc.target_size, &mut writer)?;

    writer.finalise()?;

    Ok(CellQuality {
        cell_indices: Vec::new(),
        gene_indices: Vec::new(),
        lib_size: lib_size_out,
        nnz: nnz_out,
    })
}

/// Writes h5ad normalised counts to disk by reconstructing raw counts
///
/// For datasets where only normalised counts are available in X, this function
/// reads the library sizes from an obs column, reverses the normalisation to
/// reconstruct integer counts, and writes them in the standard binary format.
///
/// Assumes normalisation was: norm = ln1p(x / lib_size * target_size)
/// Reconstruction: x = round(expm1(norm) * lib_size / target_size)
///
/// ### Params
///
/// * `h5_path` - Path to the h5ad file
/// * `bin_path` - Path to the binary file to write to
/// * `cs_type` - Storage format of X: "csc" or "csr"
/// * `no_cells` - Total number of observations
/// * `no_genes` - Total number of variables
/// * `obs_lib_size_col` - Name of the obs column containing total counts per cell
/// * `target_size` - Target size used in the original normalisation (e.g. 1e4)
/// * `cell_quality` - Minimum QC thresholds
/// * `verbose` - Controls verbosity
///
/// ### Returns
///
/// A tuple of `(no_cells_kept, no_genes_kept, CellQuality)`
#[allow(clippy::too_many_arguments)]
pub fn write_h5_normalised_counts<P: AsRef<Path>>(
    h5_path: P,
    bin_path: P,
    cs_type: &str,
    no_cells: usize,
    no_genes: usize,
    obs_lib_size_col: &str,
    target_size: f32,
    cell_quality: MinCellQuality,
    verbose: bool,
) -> Result<(usize, usize, CellQuality), ScDataErrors> {
    if verbose {
        println!(
            "Step 1/4: Reading library sizes from obs/{}...",
            obs_lib_size_col
        );
    }

    let file_format = parse_compressed_sparse_format(cs_type)
        .ok_or_else(|| ScDataErrors::UnknownSparseFormat(cs_type.to_string()))?;

    let file = hdf5::File::open(h5_path.as_ref())?;
    let lib_size_path = format!("obs/{}", obs_lib_size_col);

    check_h5_filters(&file, &[&lib_size_path, "X/data", "X/indices", "X/indptr"])?;

    let lib_sizes_raw: Vec<f32> = file
        .dataset(&lib_size_path)
        .map_err(|_| ScDataErrors::ObsColumnMissing(obs_lib_size_col.to_string()))?
        .read_1d()?
        .to_vec();

    if lib_sizes_raw.len() != no_cells {
        return Err(ScDataErrors::LibSizeLengthMismatch {
            expected: no_cells,
            found: lib_sizes_raw.len(),
        });
    }

    if verbose {
        let max_lib = lib_sizes_raw.iter().cloned().fold(0.0f32, f32::max);
        let min_lib = lib_sizes_raw.iter().cloned().fold(f32::INFINITY, f32::min);
        println!(
            "  Library sizes: min = {:.0} | max = {:.0}",
            min_lib, max_lib
        );
    }

    if verbose {
        println!("Step 2/4: Calculating QC metrics...");
    }

    let file_quality = match file_format {
        CompressedSparseFormat::Csc => parse_h5_normalised_quality_csc(
            &h5_path,
            (no_cells, no_genes),
            &lib_sizes_raw,
            &cell_quality,
            verbose,
        )?,
        CompressedSparseFormat::Csr => parse_h5_normalised_quality_csr(
            &h5_path,
            (no_cells, no_genes),
            &lib_sizes_raw,
            &cell_quality,
            verbose,
        )?,
    };

    if verbose {
        println!("Step 3/4: QC results:");
        println!(
            "  Genes passing QC: {} / {}",
            file_quality.genes_to_keep.len().separate_with_underscores(),
            no_genes.separate_with_underscores()
        );
        println!(
            "  Cells passing QC: {} / {}",
            file_quality.cells_to_keep.len().separate_with_underscores(),
            no_cells.separate_with_underscores()
        );
        println!("Step 4/4: Reconstructing raw counts and writing to binary format...");
    }

    let mut cell_qc = match file_format {
        CompressedSparseFormat::Csc => reconstruct_and_write_csc(
            &h5_path,
            &bin_path,
            &file_quality,
            &lib_sizes_raw,
            target_size,
            &cell_quality,
            verbose,
        )?,
        CompressedSparseFormat::Csr => reconstruct_and_write_csr(
            &h5_path,
            &bin_path,
            &file_quality,
            &lib_sizes_raw,
            target_size,
            &cell_quality,
            verbose,
        )?,
    };

    cell_qc.set_cell_indices(&file_quality.cells_to_keep);
    cell_qc.set_gene_indices(&file_quality.genes_to_keep);

    Ok((
        file_quality.cells_to_keep.len(),
        file_quality.genes_to_keep.len(),
        cell_qc,
    ))
}

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

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

    #[test]
    fn test_read_indptr_past_u32() {
        let path = std::env::temp_dir().join("bixverse_h5ad_io_indptr.h5");
        let big = [0i64, 5_000_000_000, 5_000_000_007];
        {
            let file = File::create(&path).unwrap();
            file.new_dataset_builder()
                .with_data(&big)
                .create("indptr")
                .unwrap();
        }
        let file = File::open(&path).unwrap();
        let ds = file.dataset("indptr").unwrap();
        let got = read_indptr(&ds);
        let _ = std::fs::remove_file(&path);
        assert_eq!(got.unwrap(), vec![0, 5_000_000_000, 5_000_000_007]);
    }

    #[test]
    fn test_raw_check_fewer_values_than_n_check() {
        let path = std::env::temp_dir().join("bixverse_h5ad_io_small_raw.h5");
        {
            let file = File::create(&path).unwrap();
            let x = file.create_group("X").unwrap();
            x.new_dataset_builder()
                .with_data(&[1.0f32, 2.0, 3.0])
                .create("data")
                .unwrap();
        }
        let file = File::open(&path).unwrap();
        let res = check_h5ad_is_raw(&file, &RawDataSlot::DataX, None);
        let _ = std::fs::remove_file(&path);
        assert!(res.is_ok());
    }
}