bixverse-rs 0.4.2

//! 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::prelude::*;
use crate::single_cell::sc_data::H5_CELL_SLICE_SIZE;
use crate::single_cell::sc_data::data_io::{CellGeneSparseWriter, CellOnFileQuality};

////////////
// 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<(), BixverseErrors>
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(BixverseErrors::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<(), BixverseErrors> {
    let n_check = n_check.unwrap_or(N_CHECK);

    let data: Vec<f32> = file
        .dataset(raw_slot.get_data())?
        .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<(), BixverseErrors> {
    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), BixverseErrors> {
    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(|| BixverseErrors::UnsupportH5ADFormat(h5ad_format.to_string()))?;

    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 mut lib_size = Vec::with_capacity(n_cells);
    let mut nnz = Vec::with_capacity(n_cells);

    for i in 0..n_cells {
        let start_i = file_data.indptr[i] as usize;
        let end_i = file_data.indptr[i + 1] as usize;

        let cell_data = &file_data.data[start_i..end_i];
        let cell_indices = &file_data.indices[start_i..end_i];

        let cell_chunk =
            CsrCellChunk::from_data(cell_data, cell_indices, i, cell_quality.target_size, true);

        let (nnz_i, lib_size_i) = cell_chunk.get_qc_info();
        nnz.push(nnz_i);
        lib_size.push(lib_size_i);

        writer.write_cell_chunk(cell_chunk)?;

        if verbose && (i + 1) % 100000 == 0 {
            println!(
                "  Written {} / {} cells to disk.",
                (i + 1).separate_with_underscores(),
                n_cells.separate_with_underscores()
            );
        }
    }

    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,
    ))
}

/// 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), BixverseErrors> {
    if verbose {
        println!("Step 1/3: Analysing data structure and calculating QC metrics...");
    }

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

    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 => {
            if verbose {
                println!("  Pass 1/2: Scanning library sizes...");
            }
            let cell_lib_sizes = scan_h5_csc_library_sizes(&h5_path, &file_quality, &raw_slot)?;
            if verbose {
                println!("  Pass 2/2: Writing cells with normalisation...");
            }
            write_h5_csc_to_csr_streaming(
                &h5_path,
                &bin_path,
                &file_quality,
                &raw_slot,
                &cell_lib_sizes,
                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, BixverseErrors> {
    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: Vec<u32> = file.dataset(raw_slot.get_indptr())?.read_1d()?.to_vec();

    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]) as usize);
    }

    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 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 Ok(file) = File::open(file_path) else {
                return (local_unique, local_lib_size);
            };

            let (Ok(data_ds), Ok(indices_ds)) = (
                file.dataset(raw_slot.get_data()),
                file.dataset(raw_slot.get_indices()),
            ) else {
                return (local_unique, local_lib_size);
            };

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

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

            let (Ok(chunk_data), Ok(chunk_indices)) = (
                data_ds.read_slice_1d(data_start..data_end),
                indices_ds.read_slice_1d(data_start..data_end),
            ) else {
                return (local_unique, local_lib_size);
            };

            let chunk_data: Vec<f32> = chunk_data.to_vec();
            let chunk_indices: Vec<u32> = chunk_indices.to_vec();

            for &gene_idx in gene_chunk.iter() {
                let gene_data_start = indptr[gene_idx] as usize - data_start;
                let gene_data_end = indptr[gene_idx + 1] as usize - 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>, BixverseErrors> {
    let file = File::open(file_path)?;

    check_h5ad_is_raw(&file, raw_slot, None)?;

    let data_ds = file.dataset(raw_slot.get_data())?;
    let indices_ds = file.dataset(raw_slot.get_indices())?;
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;
    let indptr_raw: Vec<u32> = indptr_ds.read_1d()?.to_vec();

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

    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] as usize)
            .min()
            .unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1] as usize)
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

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

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

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

                if let Some(&new_cell_idx) = quality.cell_old_to_new.get(&old_cell_idx) {
                    cell_lib_sizes[new_cell_idx] += 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>, BixverseErrors> {
    let file = File::open(file_path)?;

    check_h5ad_is_raw(&file, raw_slot, None)?;

    let data_ds = file.dataset(raw_slot.get_data())?;
    let indices_ds = file.dataset(raw_slot.get_indices())?;
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;

    // Read indptr first (small array)
    let indptr_raw: Vec<u32> = indptr_ds.read_1d()?.to_vec();

    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 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] as usize)
            .min()
            .unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1] as usize)
            .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_ds.read_slice_1d(start_pos..end_pos)?.to_vec();
        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(start_pos..end_pos)?.to_vec();

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

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

                if let Some(&new_cell_idx) = quality.cell_old_to_new.get(&old_cell_idx) {
                    new_data.push(chunk_data[local_idx] as u32);
                    new_indices.push(new_cell_idx);
                }
            }
            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]
/// * `cell_lib_sizes` - Vector with the pre-calculated library sizes, see
///   scan_h5_csc_library_sizes().
/// * `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,
    cell_lib_sizes: &[u32],
    target_size: f32,
    verbose: bool,
) -> Result<CellQuality, BixverseErrors> {
    let file = File::open(file_path)?;
    let data_ds = file.dataset(raw_slot.get_data())?;
    let indices_ds = file.dataset(raw_slot.get_indices())?;
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;
    let indptr_raw: Vec<u32> = indptr_ds.read_1d()?.to_vec();

    // 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()];

    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] as usize)
            .min()
            .unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1] as usize)
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

        let chunk_data: Vec<f32> = data_ds.read_slice_1d(start_pos..end_pos)?.to_vec();
        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(start_pos..end_pos)?.to_vec();

        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] as usize;
            let gene_end = indptr_raw[gene_idx + 1] as usize;

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

                if let Some(&new_cell_idx) = quality.cell_old_to_new.get(&old_cell_idx) {
                    let raw_count = chunk_data[local_idx] as u32;
                    cell_data[new_cell_idx].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...",
            cell_lib_sizes.len().separate_with_underscores()
        );
    }

    // now write cells with correct normalisation
    let mut writer = CellGeneSparseWriter::new(
        bin_path,
        true,
        cell_lib_sizes.len(),
        quality.genes_to_keep.len(),
        target_size,
    )?;
    let mut lib_size = Vec::with_capacity(cell_lib_sizes.len());
    let mut nnz = Vec::with_capacity(cell_lib_sizes.len());

    let total_cells = cell_data.len();

    for (cell_idx, mut data) in cell_data.into_iter().enumerate() {
        if verbose && (cell_idx + 1) % 100000 == 0 {
            println!(
                "   Written {} / {} cells to disk.",
                (cell_idx + 1).separate_with_underscores(),
                total_cells.separate_with_underscores()
            );
        }

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

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

        let cell_chunk =
            CsrCellChunk::from_data(&gene_counts, &gene_indices, cell_idx, target_size, true);

        let (nnz_i, lib_size_i) = cell_chunk.get_qc_info();
        nnz.push(nnz_i);
        lib_size.push(lib_size_i);

        writer.write_cell_chunk(cell_chunk)?;
    }

    writer.finalise()?;

    if verbose {
        println!(
            "   Written {} / {} cells (complete).",
            total_cells.separate_with_underscores(),
            total_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, BixverseErrors> {
    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: Vec<u32> = file.dataset(raw_slot.get_indptr())?.read_1d()?.to_vec();

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

            let Ok(file) = File::open(file_path) else {
                return local_counts;
            };

            let Ok(indices_ds) = file.dataset(raw_slot.get_indices()) else {
                return local_counts;
            };

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

            if data_start >= data_end {
                return local_counts;
            }

            let Ok(chunk_indices) = indices_ds.read_slice_1d(data_start..data_end) else {
                return local_counts;
            };

            let chunk_indices: Vec<u32> = chunk_indices.to_vec();

            for cell_idx in chunk_start_cell..=chunk_end_cell {
                let cell_data_start = indptr[cell_idx] as usize - data_start;
                let cell_data_end = indptr[cell_idx + 1] as usize - 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 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 Ok(file) = File::open(file_path) else {
                return (local_unique, local_lib_size);
            };

            let (Ok(data_ds), Ok(indices_ds)) = (
                file.dataset(raw_slot.get_data()),
                file.dataset(raw_slot.get_indices()),
            ) else {
                return (local_unique, local_lib_size);
            };

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

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

            let (Ok(chunk_data), Ok(chunk_indices)) = (
                data_ds.read_slice_1d(data_start..data_end),
                indices_ds.read_slice_1d(data_start..data_end),
            ) else {
                return (local_unique, local_lib_size);
            };

            let chunk_data: Vec<f32> = chunk_data.to_vec();
            let chunk_indices: Vec<u32> = chunk_indices.to_vec();

            for cell_idx in chunk_start_cell..=chunk_end_cell {
                let cell_data_start = indptr[cell_idx] as usize - data_start;
                let cell_data_end = indptr[cell_idx + 1] as usize - 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>, BixverseErrors> {
    let file = File::open(file_path)?;

    let data_ds = file.dataset(raw_slot.get_data())?;
    let indices_ds = file.dataset(raw_slot.get_indices())?;
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;

    let indptr_raw: Vec<u32> = indptr_ds.read_1d()?.to_vec();

    // Build CSC format directly: indptr = cells, indices = genes
    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 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] as usize)
            .min()
            .unwrap_or(0);
        let end_pos = cell_chunk
            .iter()
            .map(|&c| indptr_raw[c + 1] as usize)
            .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_ds.read_slice_1d(start_pos..end_pos)?.to_vec();
        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(start_pos..end_pos)?.to_vec();

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

            // Collect this cell's data with gene filtering
            let mut cell_data: Vec<(usize, 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;

                if let Some(&new_gene_idx) = quality.gene_old_to_new.get(&old_gene_idx) {
                    cell_data.push((new_gene_idx, 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);
            }

            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, // Return CSC format
        data_2: None::<Vec<u32>>,
        shape,
    })
}

/// Stream h5 CSR data directly to disk with batched reading
///
/// This is memory-efficient because we process cells in batches and write
/// immediately. CSR format is ideal since we can calculate library size and
/// normalisation per-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
///
/// 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, BixverseErrors> {
    let file = File::open(&file_path)?;
    let data_ds = file.dataset(raw_slot.get_data())?;
    let indices_ds = file.dataset(raw_slot.get_indices())?;
    let indptr_ds = file.dataset(raw_slot.get_indptr())?;
    let indptr_raw: Vec<u32> = indptr_ds.read_1d()?.to_vec();

    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();

    // Reusable buffers to avoid allocations
    let mut cell_data: Vec<(usize, u32)> = Vec::with_capacity(10000);
    let mut gene_indices: Vec<u32> = Vec::with_capacity(10000);
    let mut gene_counts: Vec<u32> = Vec::with_capacity(10000);

    for (batch_idx, cell_batch) in quality.cells_to_keep.chunks(H5_CELL_SLICE_SIZE).enumerate() {
        if verbose && (batch_idx % ((num_batches / 10).max(1)) == 0 || 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()
            );
        }

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

        if start_pos >= end_pos {
            for &old_cell_idx in cell_batch {
                lib_size.push(0);
                nnz.push(0);
                let new_cell_idx = quality.cell_old_to_new[&old_cell_idx];
                let empty_chunk = CsrCellChunk::from_data(
                    &[] as &[u32],
                    &[] as &[u32],
                    new_cell_idx,
                    cell_qc.target_size,
                    true,
                );
                writer.write_cell_chunk(empty_chunk)?;
            }
            continue;
        }

        let chunk_data: Vec<f32> = data_ds.read_slice_1d(start_pos..end_pos)?.to_vec();
        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(start_pos..end_pos)?.to_vec();

        for &old_cell_idx in cell_batch {
            let cell_start = indptr_raw[old_cell_idx] as usize;
            let cell_end = indptr_raw[old_cell_idx + 1] as usize;

            cell_data.clear();
            gene_indices.clear();
            gene_counts.clear();

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

                if let Some(&new_gene_idx) = quality.gene_old_to_new.get(&old_gene_idx) {
                    let raw_val = chunk_data[local_idx] as u32;
                    cell_data.push((new_gene_idx, raw_val));
                }
            }

            if !cell_data.is_empty() {
                // Check if already sorted (common in CSR)
                let needs_sort = cell_data.windows(2).any(|w| w[0].0 > w[1].0);
                if needs_sort {
                    cell_data.sort_unstable_by_key(|&(gene_idx, _)| gene_idx);
                }

                gene_indices.extend(cell_data.iter().map(|(g, _)| *g as u32));
                gene_counts.extend(cell_data.iter().map(|(_, c)| *c));
            }

            let new_cell_idx = quality.cell_old_to_new[&old_cell_idx];
            let cell_chunk = CsrCellChunk::from_data(
                &gene_counts,
                &gene_indices,
                new_cell_idx,
                cell_qc.target_size,
                true,
            );

            let (nnz_i, lib_size_i) = cell_chunk.get_qc_info();
            nnz.push(nnz_i);
            lib_size.push(lib_size_i);

            writer.write_cell_chunk(cell_chunk)?;
        }
    }

    writer.finalise()?;

    let end_write = start_write.elapsed();

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

    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 [BixverseErrors::UnsupportH5ADFormat]
pub fn dense_col_not_supported<T>() -> Result<T, BixverseErrors> {
    Err(BixverseErrors::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, BixverseErrors> {
    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>, BixverseErrors> {
            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>), BixverseErrors> {
                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>, BixverseErrors> {
    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 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;
                }
                if let Some(&new_gene_idx) = quality.gene_old_to_new.get(&gene_idx) {
                    new_data.push(val as u32);
                    new_indices.push(new_gene_idx);
                }
            }

            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, BixverseErrors> {
    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 mut gene_indices_buf: Vec<u32> = Vec::with_capacity(1024);
    let mut gene_counts_buf: Vec<u32> = Vec::with_capacity(1024);

    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])?;

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

            gene_indices_buf.clear();
            gene_counts_buf.clear();

            for (gene_idx, &val) in row.iter().enumerate() {
                if val == 0.0 {
                    continue;
                }
                if let Some(&new_gene_idx) = quality.gene_old_to_new.get(&gene_idx) {
                    gene_indices_buf.push(new_gene_idx as u32);
                    gene_counts_buf.push(val as u32);
                }
            }

            let new_cell_idx = quality.cell_old_to_new[&old_cell_idx];
            let cell_chunk = CsrCellChunk::from_data(
                &gene_counts_buf,
                &gene_indices_buf,
                new_cell_idx,
                cell_qc.target_size,
                true,
            );

            let (nnz_i, lib_size_i) = cell_chunk.get_qc_info();
            nnz.push(nnz_i);
            lib_size.push(lib_size_i);

            writer.write_cell_chunk(cell_chunk)?;
        }
    }

    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: Vec<u32> = file.dataset("X/indptr")?.read_1d()?.to_vec();

    // 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]) as usize)
        .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];

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

        if start_pos >= end_pos {
            continue;
        }

        let indices_ds = file.dataset("X/indices")?;
        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(start_pos..end_pos)?.to_vec();

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

            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: Vec<u32> = file.dataset("X/indptr")?.read_1d()?.to_vec();
    let indices_ds = 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]) as usize)
        .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] as usize;
        let data_end = indptr[chunk_end + 1] as usize;

        if data_start >= data_end {
            continue;
        }

        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(data_start..data_end)?.to_vec();

        for cell_idx in chunk_start..=chunk_end {
            let cell_start = indptr[cell_idx] as usize - data_start;
            let cell_end = indptr[cell_idx + 1] as usize - 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, BixverseErrors> {
    let file = hdf5::File::open(file_path)?;
    let data_ds = file.dataset("X/data")?;
    let indices_ds = file.dataset("X/indices")?;
    let indptr_ds = file.dataset("X/indptr")?;
    let indptr_raw: Vec<u32> = indptr_ds.read_1d()?.to_vec();

    let mut writer = CellGeneSparseWriter::new(
        bin_path,
        true,
        quality.cells_to_keep.len(),
        quality.genes_to_keep.len(),
        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 mut cell_data: Vec<(usize, u32)> = Vec::with_capacity(10000);
    let mut gene_indices: Vec<u32> = Vec::with_capacity(10000);
    let mut gene_counts: Vec<u32> = Vec::with_capacity(10000);

    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] as usize)
            .min()
            .unwrap_or(0);
        let end_pos = cell_batch
            .iter()
            .map(|&c| indptr_raw[c + 1] as usize)
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            for &old_cell_idx in cell_batch {
                let new_cell_idx = quality.cell_old_to_new[&old_cell_idx];
                let empty_chunk = CsrCellChunk::from_data(
                    &[] as &[u32],
                    &[] as &[u32],
                    new_cell_idx,
                    cell_qc.target_size,
                    true,
                );
                lib_size_out.push(0);
                nnz_out.push(0);
                writer.write_cell_chunk(empty_chunk)?;
            }
            continue;
        }

        let chunk_data: Vec<f32> = data_ds.read_slice_1d(start_pos..end_pos)?.to_vec();
        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(start_pos..end_pos)?.to_vec();

        for &old_cell_idx in cell_batch {
            let cell_start = indptr_raw[old_cell_idx] as usize;
            let cell_end = indptr_raw[old_cell_idx + 1] as usize;
            let lib_size = lib_sizes[old_cell_idx];

            cell_data.clear();
            gene_indices.clear();
            gene_counts.clear();

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

                if let Some(&new_gene_idx) = quality.gene_old_to_new.get(&old_gene_idx) {
                    let norm_val = chunk_data[local_idx];
                    let raw_count = reconstruct_raw_count(norm_val, lib_size, target_size);
                    if raw_count > 0 {
                        cell_data.push((new_gene_idx, raw_count));
                    }
                }
            }

            if !cell_data.is_empty() {
                let needs_sort = cell_data.windows(2).any(|w| w[0].0 > w[1].0);
                if needs_sort {
                    cell_data.sort_unstable_by_key(|&(g, _)| g);
                }
                gene_indices.extend(cell_data.iter().map(|(g, _)| *g as u32));
                gene_counts.extend(cell_data.iter().map(|(_, c)| *c));
            }

            let new_cell_idx = quality.cell_old_to_new[&old_cell_idx];
            let cell_chunk = CsrCellChunk::from_data(
                &gene_counts,
                &gene_indices,
                new_cell_idx,
                cell_qc.target_size,
                true,
            );

            let (nnz_i, lib_size_i) = cell_chunk.get_qc_info();
            nnz_out.push(nnz_i);
            lib_size_out.push(lib_size_i);

            writer.write_cell_chunk(cell_chunk)?;
        }
    }

    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, BixverseErrors> {
    let file = hdf5::File::open(file_path)?;
    let data_ds = file.dataset("X/data")?;
    let indices_ds = file.dataset("X/indices")?;
    let indptr_ds = file.dataset("X/indptr")?;
    let indptr_raw: Vec<u32> = indptr_ds.read_1d()?.to_vec();

    // 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()];

    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] as usize)
            .min()
            .unwrap_or(0);
        let end_pos = gene_chunk
            .iter()
            .map(|&g| indptr_raw[g + 1] as usize)
            .max()
            .unwrap_or(0);

        if start_pos >= end_pos {
            continue;
        }

        let chunk_data: Vec<f32> = data_ds.read_slice_1d(start_pos..end_pos)?.to_vec();
        let chunk_indices: Vec<u32> = indices_ds.read_slice_1d(start_pos..end_pos)?.to_vec();

        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] as usize;
            let gene_end = indptr_raw[gene_idx + 1] as usize;

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

                if let Some(&new_cell_idx) = quality.cell_old_to_new.get(&old_cell_idx) {
                    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].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(),
        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 = cell_data.len();

    for (cell_idx, mut data) in cell_data.into_iter().enumerate() {
        if verbose && (cell_idx + 1) % 100000 == 0 {
            println!(
                "   Written {} / {} cells.",
                (cell_idx + 1).separate_with_underscores(),
                total_cells.separate_with_underscores()
            );
        }

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

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

        let cell_chunk = CsrCellChunk::from_data(
            &gene_counts,
            &gene_indices,
            cell_idx,
            cell_qc.target_size,
            true,
        );

        let (nnz_i, lib_size_i) = cell_chunk.get_qc_info();
        nnz_out.push(nnz_i);
        lib_size_out.push(lib_size_i);

        writer.write_cell_chunk(cell_chunk)?;
    }

    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), BixverseErrors> {
    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(|| BixverseErrors::UnknownSparseFormat(cs_type.to_string()))?;

    let file = hdf5::File::open(h5_path.as_ref())?;
    let lib_size_path = format!("obs/{}", obs_lib_size_col);
    let lib_sizes_raw: Vec<f32> = file
        .dataset(&lib_size_path)
        .map_err(|_| BixverseErrors::ObsColumnMissing(obs_lib_size_col.to_string()))?
        .read_1d()?
        .to_vec();

    if lib_sizes_raw.len() != no_cells {
        return Err(BixverseErrors::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,
    ))
}