"""Methods and classes for writing data to disk.

- Methods:
    - create_zarr_dataset: Creates and returns a Zarr hierarchy/dataset.
    - create_zarr_obj_array: Creates and returns a Zarr object array.
    - create_zarr_count_assay: Creates and returns a Zarr array with name 'counts'.
    - subset_assay_zarr: Selects a subset of the data in an assay in the specified Zarr hierarchy.
    - dask_to_zarr: Creates a Zarr hierarchy from a Dask array.
    - to_h5ad: Convert a Zarr file to H5ad format
    - to_mtx: Convert a Zarr file to MTX format

- Classes:
    - ZarrMerge: Merge multiple Zarr files into a single Zarr file.
    - SubsetZarr: Extracts a subset of cells from the given Zarr file and saves into a new Zarr file.
    - CrToZarr: A class for converting data in the Cellranger format to a Zarr hierarchy.
    - H5adToZarr: A class for converting data in the Cellranger Matrix Market format to a Zarr hierarchy.
    - NaboH5ToZarr: A class for converting data in a h5 file generated by Nabo, to a Zarr hierarchy.
    - LoomToZarr: A class for converting data in a Loom file to a Zarr hierarchy.
"""

import zarr
from typing import Any, Tuple, List, Union, Dict, Optional
import numpy as np
from .readers import CrReader, H5adReader, NaboH5Reader, LoomReader, CSVReader
import os
import pandas as pd
from .utils import controlled_compute, logger, tqdmbar, show_dask_progress, load_zarr
from scipy.sparse import csr_matrix

__all__ = [
    "create_zarr_dataset",
    "create_zarr_obj_array",
    "create_zarr_count_assay",
    "subset_assay_zarr",
    "dask_to_zarr",
    "ZarrMerge",
    "SubsetZarr",
    "CrToZarr",
    "H5adToZarr",
    "NaboH5ToZarr",
    "LoomToZarr",
    "SparseToZarr",
    "to_h5ad",
    "to_mtx",
    "CSVtoZarr",
]


def create_zarr_dataset(
    g: zarr.Group,
    name: str,
    chunks: tuple,
    dtype: Any,
    shape: Tuple,
    overwrite: bool = True,
) -> zarr.hierarchy:
    """Creates and returns a Zarr array.

    Args:
        g (zarr.hierarchy):
        name (str):
        chunks (tuple):
        dtype (Any):
        shape (Tuple):
        overwrite (bool):

    Returns:
        A Zarr Array.
    """
    from numcodecs import Blosc

    compressor = Blosc(cname="lz4", clevel=5, shuffle=Blosc.BITSHUFFLE)
    return g.create_dataset(
        name,
        chunks=chunks,
        dtype=dtype,
        shape=shape,
        compressor=compressor,
        overwrite=overwrite,
    )


def dtype_fix(dtype, data: np.ndarray):
    if dtype is None or dtype == object:
        return "U" + str(max([len(str(x)) for x in data]))
    if np.issubdtype(data.dtype, np.dtype("S")):
        try:
            adata = data.astype("U")
        except UnicodeDecodeError:
            adata = np.array([x.decode("UTF-8") for x in data]).astype("U")
        return adata.dtype
    return dtype


def create_zarr_obj_array(
    g: zarr.Group,
    name: str,
    data,
    dtype: Union[str, Any] = None,
    overwrite: bool = True,
    chunk_size: int = 100000,
    shape: Optional[int] = None,
) -> zarr.hierarchy:
    """Creates and returns a Zarr object array.

    A Zarr object array can contain any type of object.
    https://zarr.readthedocs.io/en/stable/tutorial.html#object-arrays

    Args:
        g (zarr.hierarchy):
        name (str):
        data ():
        dtype (Union[str, Any]):
        overwrite (bool):
        chunk_size (int):
        shape:

    Returns:
        A Zarr object Array.
    """

    from numcodecs import Blosc

    compressor = Blosc(cname="lz4", clevel=5, shuffle=Blosc.BITSHUFFLE)

    if chunk_size is None or chunk_size is False:
        chunks = False
    else:
        chunks = (chunk_size,)

    if data is not None:
        data = np.array(data)
        dtype = dtype_fix(dtype, data)

        return g.create_dataset(
            name,
            data=data,
            chunks=chunks,
            shape=len(data),
            dtype=dtype,
            overwrite=overwrite,
            compressor=compressor,
        )
    else:
        return g.create_dataset(
            name,
            chunks=chunks,
            shape=shape,
            dtype=dtype,
            overwrite=overwrite,
            compressor=compressor,
        )


def create_zarr_count_assay(
    z: zarr.Group,
    assay_name: str,
    workspace: Union[str, None],
    chunk_size: Tuple[int, int],
    n_cells: int,
    feat_ids: Union[np.ndarray, List[str]],
    feat_names: Union[np.ndarray, List[str]],
    dtype: str = "uint32",
) -> zarr.hierarchy:
    """Creates and returns a Zarr array with name 'counts'.

    Args:
        z (zarr.Group):
        assay_name (str):
        workspace (Union[str, None]):
        chunk_size (Tuple[int, int]):
        n_cells (int):
        feat_ids (Union[np.ndarray, List[str]]):
        feat_names (Union[np.ndarray, List[str]]):
        dtype (str = 'uint32'):

    Returns:
        A Zarr array.
    """
    if workspace is None:
        g = z.create_group(assay_name, overwrite=True)
    else:
        g = z.create_group(f"{workspace}/{assay_name}", overwrite=True)
    g.attrs["is_assay"] = True
    g.attrs["misc"] = {}
    create_zarr_obj_array(g, "featureData/ids", feat_ids)
    create_zarr_obj_array(g, "featureData/names", feat_names)
    create_zarr_obj_array(
        g, "featureData/I", [True for _ in range(len(feat_ids))], "bool"
    )
    if workspace is not None:
        g = z.create_group(f"matrices/{assay_name}", overwrite=True)
    return create_zarr_dataset(
        g, "counts", chunk_size, dtype, (n_cells, len(feat_ids)), overwrite=True
    )


def create_cell_data(z, workspace, ids, names):
    if workspace is None:
        g = z.create_group("cellData")
    else:
        g = z.create_group(f"{workspace}/cellData")
    create_zarr_obj_array(g, "ids", ids, ids.dtype)
    create_zarr_obj_array(g, "names", names, names.dtype)
    create_zarr_obj_array(g, "I", [True for _ in range(len(ids))], "bool")
    return g


def sparse_writer(
    store: zarr.hierarchy, data_stream, n_cells: int, batch_size: int
) -> int:
    s, e, = (
        0,
        0,
    )
    n_chunks = n_cells // batch_size + 1
    for a in tqdmbar(data_stream, total=n_chunks):
        e += a.shape[0]
        store.set_coordinate_selection((a.row + s, a.col), a.data)
        s = e
    return e


class CrToZarr:
    """A class for converting data in the Cellranger format to a Zarr
    hierarchy.

    Args:
        cr: A CrReader object, containing the Cellranger data.
        zarr_loc: The file name for the Zarr hierarchy or a store
        chunk_size: The requested size of chunks to load into memory and process.
        dtype: the dtype of the data.

    Attributes:
        cr: A CrReader object, containing the Cellranger data.
        chunkSizes: The requested size of chunks to load into memory and process.
        z: The Zarr hierarchy (array or group).
    """

    def __init__(
        self,
        cr: CrReader,
        zarr_loc: str,
        chunk_size=(1000, 1000),
        dtype: str = "uint32",
        workspace: str = None,
    ):
        self.cr = cr
        self.chunkSizes = chunk_size
        self.workspace = workspace
        self.z = load_zarr(zarr_loc=zarr_loc, mode="w")
        create_cell_data(
            z=self.z,
            workspace=self.workspace,
            ids=np.array(self.cr.cell_names()),
            names=np.array(self.cr.cell_names()),
        )
        for assay_name in self.cr.assayFeats.columns:
            create_zarr_count_assay(
                z=self.z,
                assay_name=assay_name,
                workspace=workspace,
                chunk_size=chunk_size,
                n_cells=self.cr.nCells,
                feat_ids=self.cr.feature_ids(assay_name),
                feat_names=self.cr.feature_names(assay_name),
                dtype=dtype,
            )

    @staticmethod
    def _prep_assay_input_ranges(af: pd.DataFrame) -> Dict[str, List[List[int]]]:
        assay_order = (
            af.T.nFeatures.groupby(af.columns).sum().sort_values(ascending=False).index
        )
        ranges = {}
        for assay in assay_order:
            temp = []
            if len(af[assay].shape) == 2:
                for i in af[assay].values[1:3].T:
                    temp.append([i[0], i[1]])
            else:
                idx = af[assay]
                temp = [[idx.start, idx.end]]
            ranges[assay] = temp
        return ranges

    @staticmethod
    def _prep_feat_index_offset(
        ranges: Dict[str, List[List[int]]]
    ) -> Dict[str, List[int]]:
        feat_offset = {}
        for i in ranges:
            feat_offset[i] = []
            lv = 0
            for j in ranges[i]:
                feat_offset[i].append(-j[0] + lv)
                lv += j[1] - j[0]
        return feat_offset

    def dump(self, batch_size: int = 1000, lines_in_mem: int = 100000) -> None:
        """Writes the count values into the Zarr matrix.

        Args:
            batch_size: Number of cells to save at a time. (Default value: 1000)
            lines_in_mem: Number of lines to read at a time from MTX file (only used for CrDirReader)
                          (Default value: 100000)

        Raises:
            AssertionError: Catches eventual bugs in the class, if number of cells does not match after transformation.

        Returns:
            None
        """
        input_ranges = self._prep_assay_input_ranges(self.cr.assayFeats)
        if self.workspace is None:
            stores = {x: self.z[f"{x}/counts"] for x in input_ranges}
        else:
            stores = {x: self.z[f"matrices/{x}/counts"] for x in input_ranges}
        feat_offset = self._prep_feat_index_offset(input_ranges)
        s = 0
        n_chunks = self.cr.nCells // batch_size + 1
        for a in tqdmbar(self.cr.consume(batch_size, lines_in_mem), total=n_chunks):
            for assay in input_ranges:
                idx = np.zeros(a.col.shape[0]).astype(bool)
                feat_coords = a.col.copy()
                for r, of in zip(input_ranges[assay], feat_offset[assay]):
                    temp = (a.col >= r[0]) & (a.col < r[1])
                    if of != 0:
                        feat_coords[temp] = (
                            feat_coords[temp] + of
                        )  # of is already a negative value
                    idx = idx | temp
                if idx.sum() > 0:
                    stores[assay].set_coordinate_selection(
                        (s + a.row[idx], feat_coords[idx]), a.data[idx]
                    )
                else:
                    logger.warning(
                        f"No feature captured from chunk {s} to {s+a.shape[0]} for assay: {assay}"
                    )
            s += a.shape[0]
        if s != self.cr.nCells:
            raise AssertionError(
                "ERROR: This is a bug in CrToZarr. All cells might not have been successfully "
                "written into the zarr file. Please report this issue"
            )


class H5adToZarr:
    """A class for converting data in anndata's H5ad format to Zarr hierarchy.

    Args:
        h5ad: A H5adReader object, containing the Cellranger data.
        zarr_loc: The file name for the Zarr hierarchy or a store
        assay_name: the name of the assay (e. g. 'RNA')
        workspace: An optional workspace id.
        chunk_size: The requested size of chunks to load into memory and process.

    Attributes:
        h5ad: A h5ad object (h5 file with added AnnData structure).
        chunkSizes: The requested size of chunks to load into memory and process.
        assayName: The Zarr hierarchy (array or group).
        z: The Zarr hierarchy (array or group).
    """

    def __init__(
        self,
        h5ad: H5adReader,
        zarr_loc: str,
        assay_name: str = None,
        workspace: Union[str, None] = None,
        chunk_size=(1000, 1000),
    ):
        # TODO: support for multiple assay. One of the `var` datasets can be used to group features in separate assays
        self.h5ad = h5ad
        self.chunkSizes = chunk_size
        self.workspace = workspace
        if assay_name is None:
            logger.info(
                f"No value provided for assay names. Will use default value: 'RNA'"
            )
            self.assayName = "RNA"
        else:
            self.assayName = assay_name
        self.z = load_zarr(zarr_loc=zarr_loc, mode="w")
        self._ini_cell_data()
        create_zarr_count_assay(
            z=self.z,
            assay_name=self.assayName,
            workspace=workspace,
            chunk_size=chunk_size,
            n_cells=self.h5ad.nCells,
            feat_ids=self.h5ad.feat_ids(),
            feat_names=self.h5ad.feat_names(),
            dtype=self.h5ad.matrixDtype,
        )
        self._ini_feature_data()

    def _ini_cell_data(self):
        ids = self.h5ad.cell_ids()
        g = create_cell_data(
            z=self.z,
            workspace=self.workspace,
            ids=ids,
            names=ids,
        )
        for i, j in self.h5ad.get_cell_columns():
            create_zarr_obj_array(g, i, j, j.dtype)

    def _ini_feature_data(self):
        if self.workspace is None:
            g = self.z[f"{self.assayName}/featureData"]
        else:
            g = self.z[f"{self.workspace}/{self.assayName}/featureData"]
        for i, j in self.h5ad.get_feat_columns():
            if i not in g:
                create_zarr_obj_array(g, i, j, j.dtype)

    def dump(self, batch_size: int = 1000) -> None:
        # TODO: add informed description to docstring
        """
        Raises:
            AssertionError: Catches eventual bugs in the class, if number of cells does not match after transformation.

        Returns:
            None
        """
        if self.workspace is None:
            store = self.z[f"{self.assayName}/counts"]
        else:
            store = self.z[f"matrices/{self.assayName}/counts"]

        total_cells_written = sparse_writer(
            store=store,
            data_stream=self.h5ad.consume(batch_size),
            n_cells=self.h5ad.nCells,
            batch_size=batch_size,
        )
        if total_cells_written != self.h5ad.nCells:
            raise AssertionError(
                "ERROR: This is a bug in H5adToZarr. All cells might not have been successfully "
                "written into the zarr file. Please report this issue"
            )


class NaboH5ToZarr:
    """A class for converting data in a h5 file generated by Nabo, to a Zarr
    hierarchy.

    Args:
        h5: A Nabo h5 object containing the data.
        zarr_fn: The file name for the Zarr hierarchy.
        assay_name: the name of the assay (e. g. 'RNA')
        chunk_size: The requested size of chunks to load into memory and process.
        dtype: the dtype of the data.

    Attributes:
        h5: A Nabo h5 object.
        fn: The file name for the Zarr hierarchy.
        chunkSizes: The requested size of chunks to load into memory and process.
        assayName: The Zarr hierarchy (array or group).
        z: The Zarr hierarchy (array or group).
    """

    def __init__(
        self,
        h5: NaboH5Reader,
        zarr_fn: str,
        assay_name: str = None,
        workspace: Union[str, None] = None,
        chunk_size=(1000, 1000),
        dtype: str = "uint32",
    ):
        self.h5 = h5
        self.fn = zarr_fn
        self.chunkSizes = chunk_size
        self.workspace = workspace
        if assay_name is None:
            logger.info(
                f"No value provided for assay names. Will use default value: 'RNA'"
            )
            self.assayName = "RNA"
        else:
            self.assayName = assay_name
        self.z = zarr.open(self.fn, mode="w")
        self._ini_cell_data()
        create_zarr_count_assay(
            z=self.z,
            assay_name=self.assayName,
            workspace=workspace,
            chunk_size=chunk_size,
            n_cells=self.h5.nCells,
            feat_ids=self.h5.feat_ids(),
            feat_names=self.h5.feat_names(),
            dtype=dtype,
        )

    def _ini_cell_data(self):
        ids = self.h5.cell_ids()
        _ = create_cell_data(
            z=self.z,
            workspace=self.workspace,
            ids=np.array(ids),
            names=np.array(ids),
        )

    def dump(self, batch_size: int = 500) -> None:
        # TODO: add informed description to docstring
        """
        Raises:
            AssertionError: Catches eventual bugs in the class, if number of cells does not match after transformation.

        Returns:
            None
        """
        if self.workspace is None:
            store = self.z[f"{self.assayName}/counts"]
        else:
            store = self.z[f"matrices/{self.assayName}/counts"]
        s, e, = (
            0,
            0,
        )
        n_chunks = self.h5.nCells // batch_size + 1
        for a in tqdmbar(self.h5.consume(batch_size), total=n_chunks):
            e += a.shape[0]
            store[s:e] = a
            s = e
        if e != self.h5.nCells:
            raise AssertionError(
                "ERROR: This is a bug in NaboH5ToZarr. All cells might not have been successfully "
                "written into the zarr file. Please report this issue"
            )


class LoomToZarr:
    """A class for converting data in a Loom file to a Zarr hierarchy. Converts
    a Loom file read using scarf.LoomReader into Scarf's Zarr format.

    Args:
        loom: LoomReader object used to open Loom format file
        zarr_fn: Output Zarr filename with path
        assay_name: Name for the output assay. If not provided then automatically set to RNA
        chunk_size: Chunk size for the count matrix saved in Zarr file.

    Attributes:
        loom: A scarf.LoomReader object used to open Loom format file.
        fn: The file name for the Zarr hierarchy.
        chunkSizes: The requested size of chunks to load into memory and process.
        assayName: The Zarr hierarchy (array or group).
        z: The Zarr hierarchy (array or group).
    """

    def __init__(
        self,
        loom: LoomReader,
        zarr_fn: str,
        assay_name: str = None,
        workspace: Union[str, None] = None,
        chunk_size=(1000, 1000),
    ):
        # TODO: support for multiple assay. Data from within individual layers can be treated as separate assays
        self.loom = loom
        self.fn = zarr_fn
        self.chunkSizes = chunk_size
        self.workspace = workspace
        if assay_name is None:
            logger.info(
                f"No value provided for assay names. Will use default value: 'RNA'"
            )
            self.assayName = "RNA"
        else:
            self.assayName = assay_name
        self.z = zarr.open(self.fn, mode="w")
        self._ini_cell_data()
        create_zarr_count_assay(
            z=self.z,
            assay_name=self.assayName,
            workspace=workspace,
            chunk_size=chunk_size,
            n_cells=self.loom.nCells,
            feat_ids=self.loom.feature_ids(),
            feat_names=self.loom.feature_names(),
            dtype=self.loom.matrixDtype,
        )
        self._ini_feature_data()

    def _ini_cell_data(self):
        ids = np.array(self.loom.cell_ids())
        g = create_cell_data(
            z=self.z,
            workspace=self.workspace,
            ids=ids,
            names=ids,
        )
        for i, j in self.loom.get_cell_attrs():
            try:
                create_zarr_obj_array(g, i, j, j.dtype)
            except UnicodeDecodeError:
                logger.warning(f"Could not import {i} cell(column) attribute")

    def _ini_feature_data(self):
        if self.workspace is None:
            g = self.z[f"{self.assayName}/featureData"]
        else:
            g = self.z[f"{self.workspace}/{self.assayName}/featureData"]
        for i, j in self.loom.get_feature_attrs():
            create_zarr_obj_array(g, i, j, j.dtype)

    def dump(self, batch_size: int = 1000) -> None:
        # TODO: add informed description to docstring
        """
        Raises:
            AssertionError: Catches eventual bugs in the class, if number of cells does not match after transformation.

        Returns:
            None
        """
        if self.workspace is None:
            store = self.z[f"{self.assayName}/counts"]
        else:
            store = self.z[f"matrices/{self.assayName}/counts"]
        total_cells_written = sparse_writer(
            store=store,
            data_stream=self.loom.consume(batch_size),
            n_cells=self.loom.nCells,
            batch_size=batch_size,
        )
        if total_cells_written != self.loom.nCells:
            raise AssertionError(
                "ERROR: This is a bug in LoomToZarr. All cells might not have been successfully "
                "written into the zarr file. Please report this issue"
            )


class SparseToZarr:
    """A class for converting data in a sparse matrix to a Zarr hierarchy.

    Args:
        csr_mat: A CSR format sparse matrix
        zarr_fn: Output Zarr filename with path
        cell_ids: Cell IDs for the cells in the dataset.
        feature_ids: Feature IDs for the features in the dataset.
        assay_name: Name for the output assay. If not provided then automatically set to RNA.
        chunk_size: The requested size of chunks to load into memory and process.

    Raises:
        ValueError: Raised if number of input cell or feature IDs does not match the matrix.
        AssertionError: Catches eventual bugs in the class, if number of cells does not match after transformation.

    Attributes:
        mat: Input CSR matrix
        fn: The file name for the Zarr hierarchy.
        chunkSizes: The requested size of chunks to load into memory and process.
        assayName: The Zarr hierarchy (array or group).
        z: The Zarr hierarchy (array or group).
    """

    def __init__(
        self,
        csr_mat: csr_matrix,
        zarr_fn: str,
        cell_ids: Union[np.ndarray, List[str]],
        feature_ids: Union[np.ndarray, List[str]],
        assay_name: str = None,
        workspace: Union[str, None] = None,
        feature_names: Union[np.ndarray, List[str], None] = None,
        chunk_size=(1000, 1000),
        matrix_dtype: Optional[np.dtype] = None,
    ):
        self.mat = csr_mat
        self.fn = zarr_fn
        self.chunkSizes = chunk_size
        self.workspace = workspace
        cell_ids = np.array(cell_ids)
        if matrix_dtype is None:
            self.matrixDtype = self.mat.dtype
        else:
            self.matrixDtype = matrix_dtype
        if assay_name is None:
            logger.info(
                f"No value provided for assay names. Will use default value: 'RNA'"
            )
            self.assayName = "RNA"
        else:
            self.assayName = assay_name
        self.nCells, self.nFeatures = self.mat.shape
        if len(cell_ids) != self.nCells:
            raise ValueError(
                "ERROR: Number of cell ids are not same as number of cells in the matrix"
            )
        if len(feature_ids) != self.nFeatures:
            raise ValueError(
                "ERROR: Number of feature ids are not same as number of features in the matrix"
            )

        self.z = zarr.open(self.fn, mode="w")
        _ = create_cell_data(
            z=self.z,
            workspace=self.workspace,
            ids=cell_ids,
            names=cell_ids,
        )
        if feature_names is None:
            feature_names = feature_ids
        create_zarr_count_assay(
            z=self.z,
            assay_name=self.assayName,
            workspace=workspace,
            chunk_size=chunk_size,
            n_cells=self.nCells,
            feat_ids=feature_ids,
            feat_names=feature_names,
            dtype=self.matrixDtype,
        )

    def dump(self, batch_size: Optional[int] = None) -> None:
        """Write out the data matrix into the Zarr hierarchy.

        Args:
            batch_size: Number of cells to be written in one go. By default, this value will automatically be chosen
                        based on the chunk size in the cell dimension.

         Raises:
            ValueError: Raised if there is any unexpected errors when writing to the Zarr hierarchy.
            AssertionError: Catches eventual bugs in the class, if number of cells does not match after transformation.

        Returns:
            None
        """

        if self.workspace is None:
            store = self.z[f"{self.assayName}/counts"]
        else:
            store = self.z[f"matrices/{self.assayName}/counts"]

        if batch_size is None:
            batch_size = store.chunks[0]
        s, e, = (
            0,
            0,
        )
        n_chunks = self.nCells // batch_size + 1
        for e in tqdmbar(
            range(batch_size, self.nCells + batch_size, batch_size),
            total=n_chunks,
            desc="Writing data matrix",
        ):
            if s == self.nCells:
                raise ValueError(
                    "Unexpected error encountered in writing to Zarr. The last iteration has failed. "
                    "Please report this issue."
                )
            if e > self.nCells:
                e = self.nCells
            a = self.mat[s:e].tocoo()
            store.set_coordinate_selection(
                (a.row + s, a.col), a.data.astype(self.matrixDtype)
            )
            s = e
        if e != self.nCells:
            raise AssertionError(
                "ERROR: This is a bug in SparseToZarr. All cells might not have been successfully "
                "written into the zarr file. Please report this issue"
            )


class CSVtoZarr:
    """A class for converting data from CSV format to a Zarr hierarchy.

    Args:
        cr: A CSVReader object
        zarr_fn: The file name for the Zarr hierarchy.
        assay_name: A label for the assay. Ex. "RNA" or "ATAC"
        chunk_size: The requested size of chunks to load into memory and process.
        dtype: the dtype of the data.

    Attributes:
        csvr: A CSVReader object
        fn: The file name for the Zarr hierarchy.
        chunkSizes: The requested size of chunks to store in Zarr file
        z: The Zarr hierarchy (array or group).
    """

    def __init__(
        self,
        cr: CSVReader,
        zarr_fn: str,
        assay_name: str,
        chunk_size=(1000, 1000),
        workspace: Union[str, None] = None,
        dtype: Optional[np.dtype] = None,
    ):
        self.csvr = cr
        self.fn = zarr_fn
        self.assayName = assay_name
        self.chunkSizes = chunk_size
        self.workspace = workspace
        self.z = zarr.open(self.fn, mode="w")
        if dtype is not None:
            self.dtype = dtype
        else:
            self.dtype = next(self.csvr.consume())[0].dtype
        cell_ids = self.csvr.cell_ids()
        _ = create_cell_data(
            z=self.z,
            workspace=workspace,
            ids=cell_ids,
            names=cell_ids,
        )
        create_zarr_count_assay(
            z=self.z,
            assay_name=self.assayName,
            workspace=workspace,
            chunk_size=chunk_size,
            n_cells=self.csvr.nCells,
            feat_ids=self.csvr.feature_ids(),
            feat_names=self.csvr.feature_ids(),
            dtype=self.dtype,
        )

    def dump(self) -> None:
        """Writes the count values into the Zarr matrix.

        Args:

        Raises:
            AssertionError: Catches eventual bugs in the class, if number of cells does not match after transformation.

        Returns:
            None
        """

        if self.workspace is None:
            store = self.z[f"{self.assayName}/counts"]
        else:
            store = self.z[f"matrices/{self.assayName}/counts"]
        cell_data = [
            create_zarr_obj_array(
                self.z["cellData"], name=x, data=None, dtype=y, shape=self.csvr.nCells
            )
            for x, y in zip(self.csvr.cellDataCols, self.csvr.cellDataDtypes)
        ]
        batch_size = self.csvr.pandas_kwargs["chunksize"]
        s, e, = (
            0,
            0,
        )
        if self.csvr.nCells % batch_size == 0:
            n_chunks = int(self.csvr.nCells / batch_size)
        else:
            n_chunks = (self.csvr.nCells // batch_size) + 1
        for a, c in tqdmbar(self.csvr.consume(), total=n_chunks):
            e += a.shape[0]
            if self.dtype is not None:
                store[s:e] = a.astype(self.dtype)
            else:
                store[s:e] = a
            if c is not None:
                for n, i in enumerate(c.T):
                    cell_data[n][s:e] = i
            s = e
        if e != self.csvr.nCells:
            raise AssertionError(
                "ERROR: This is a bug in LoomToZarr. All cells might not have been successfully "
                "written into the zarr file. Please report this issue"
            )


def subset_assay_zarr(
    zarr_fn: str,
    in_grp: str,
    out_grp: str,
    cells_idx: np.ndarray,
    feat_idx: np.ndarray,
    chunk_size: tuple,
):
    """Selects a subset of the data in an assay in the specified Zarr
    hierarchy.

    For the arguments `cells_idx` and `feat_idx`, refer to the documentation for numpy.split:
    https://numpy.org/doc/stable/reference/generated/numpy.split.html

    Args:
        zarr_fn: The file name for the Zarr hierarchy.
        in_grp: Group in Zarr hierarchy to subset.
        out_grp: Group name in Zarr hierarchy to write subsetted assay to.
        cells_idx: A list of cell indices to (keep | drop ?).
        feat_idx: A list of feature indices to (keep | drop ?).
        chunk_size: The requested size of chunks to load into memory and process.

    Returns:
        None
    """
    z = zarr.open(zarr_fn, "r+")
    ig = z[in_grp]
    og = create_zarr_dataset(
        z, out_grp, chunk_size, "uint32", (len(cells_idx), len(feat_idx))
    )
    pos_start, pos_end = 0, 0
    for i in tqdmbar(np.array_split(cells_idx, len(cells_idx) // chunk_size[0] + 1)):
        pos_end += len(i)
        og[pos_start:pos_end, :] = ig.get_orthogonal_selection((i, feat_idx))
        pos_start = pos_end
    return None


def dask_to_zarr(df, z, loc, chunk_size, nthreads: int, msg: str = None):
    # TODO: perhaps change name of Dask array so it does not get confused with a dataframe
    """Creates a Zarr hierarchy from a Dask array.

    Args:
        df (): Dask array.
        z (): Zarr hierarchy.
        loc (): Location to write data/Zarr hierarchy to.
        chunk_size (): Size of chunks to load into memory and process.
        nthreads (int): Number of threads to use.
        msg (str): Message to use with progress bar (Default: f"Writing data to {loc}").
    """
    if msg is None:
        msg = f"Writing data to {loc}"
    og = create_zarr_dataset(z, loc, chunk_size, "float64", df.shape)
    pos_start, pos_end = 0, 0
    for i in tqdmbar(df.blocks, total=df.numblocks[0], desc=msg):
        pos_end += i.shape[0]
        og[pos_start:pos_end, :] = controlled_compute(i, nthreads)
        pos_start = pos_end
    return None


class SubsetZarr:
    """Split Zarr file using a subset of cells.

    Args:
        zarr_loc: Path for the output (subsetted) Zarr file
        assays: Source assays to be subsetted. These assays must be from the same dataset
        cell_key: Name of a boolean column in cell metadata. The cells with value True are included in the
                  subset. Only used when cell_idx is None.
        cell_idx: Indices of the cells to be included in the subsetted.
        reset_cell_filter: If True, then the cell filtering information is removed, i.e. even the filtered out cells
                           are set as True as in the 'I' column. To keep the filtering information set the value for
                           this parameter to False. (Default value: True)
        overwrite_existing_file: If True, then overwrites the existing data. (Default value: False)
        overwrite_cell_data: If True, then overwrites cell data (Default value: True)
    """

    def __init__(
        self,
        zarr_loc: str,
        assays: list,
        in_workspace: Union[str, None] = None,
        out_workspace: Union[str, None] = None,
        cell_key: str = None,
        cell_idx: np.ndarray = None,
        reset_cell_filter: bool = True,
        overwrite_existing_file: bool = False,
        overwrite_cell_data: bool = False,
    ) -> None:
        self.resetCells = reset_cell_filter
        self.overFn = overwrite_existing_file
        self.overCells = overwrite_cell_data
        self.inWorkspace = in_workspace
        self.outWorkspace = out_workspace
        self.z = self._check_files(zarr_loc)
        self.assays = self._check_assays(assays)
        self.cellIdx = self._check_idx(cell_key, cell_idx)

    def _check_files(self, zarr_loc: str):
        if (
            isinstance(zarr_loc, str)
            and os.path.isdir(zarr_loc)
            and self.overFn is False
        ):
            raise ValueError(
                f"Zarr file with name: {zarr_loc} already exists.\n"
                f"If you want to overwrite it then please set  overwrite_existing_file to True. "
                f"No subsetting was performed."
            )
        return load_zarr(zarr_loc=zarr_loc, mode="w")

    @staticmethod
    def _check_assays(assays):
        if type(assays) != list:
            raise TypeError(
                "Value for parameter `assays` should be a list. For example, `[ds.RNA]`"
            )
        n = []
        for assay in assays:
            try:
                n.append(assay.cells.N)
            except AttributeError:
                raise ValueError(
                    "Please make sure you are passing actual assay objects and not assay names. "
                    "For example, `[ds.RNA]`"
                )
        if len(set(n)) != 1:
            raise ValueError(
                f"ERROR: Provided assays do not have the same numer of cells. Please make "
                f"sure that the assays are from the same DataStore."
            )
        return assays

    def _check_idx(self, cell_key, cell_idx):
        if cell_key is None and cell_idx is None:
            raise ValueError("Both `cell_key` and `cell_idx` parameters cannot be None")
        if cell_idx is None:
            for assay in self.assays:
                try:
                    idx = assay.cells.fetch_all(cell_key)
                except KeyError:
                    raise ValueError(
                        f"ERROR: Provided cell_key {cell_key} was not found in the assay: {assay.name}"
                    )
                if idx.dtype != bool:
                    raise ValueError(
                        f"ERROR: {cell_key} is not of boolean type. Cannot perform subsetting"
                    )
                if cell_idx is None:
                    cell_idx = idx
                else:
                    if np.all(cell_idx == idx) is False:
                        raise ValueError(
                            f"ERROR: Provided cell_key {cell_key} is not consistent across the assays. "
                            f"Please make sure that the assays are from the same DataStore."
                        )
            cell_idx = np.where(cell_idx)[0]
        else:
            cell_idx = np.array(cell_idx)
            if np.issubdtype(cell_idx.dtype, np.integer) is False:
                raise ValueError(
                    f"ERROR: `cell_idx` must be of integer type. Provided array has a dtype: {cell_idx.dtype}"
                )
            if max(cell_idx) >= self.assays[0].cells.N:
                raise ValueError(
                    f"ERROR: `cell_idx` max value is larger than the number of cells in the data."
                )
        return cell_idx

    def _prep_cell_data(self):
        if self.outWorkspace is None:
            cell_slot = "cellData"
        else:
            cell_slot = f"{self.outWorkspace}/cellData"
        if cell_slot in self.z:
            g = self.z[cell_slot]
        else:
            g = self.z.create_group(cell_slot)

        if self.outWorkspace is None:
            cell_data = self.assays[0].z["/cellData"]
        else:
            cell_data = self.assays[0].z[f"/{self.outWorkspace}/cellData"]

        n_cells = len(self.cellIdx)
        for i in cell_data.keys():
            if i in g and self.overCells is False:
                continue
            if i in ["I"] and self.resetCells:
                create_zarr_obj_array(g, "I", [True for _ in range(n_cells)], "bool")
                continue
            v = cell_data[i][:][self.cellIdx]
            create_zarr_obj_array(g, i, v, dtype=v.dtype)

    def _prep_counts(self):
        n_cells = len(self.cellIdx)
        for assay in self.assays:
            create_zarr_count_assay(
                z=self.z,
                assay_name=assay.name,
                workspace=self.outWorkspace,
                chunk_size=assay.rawData.chunksize,
                n_cells=n_cells,
                feat_ids=assay.feats.fetch_all("ids"),
                feat_names=assay.feats.fetch_all("names"),
                dtype=assay.rawData.dtype,
            )

    def dump(self):
        self._prep_cell_data()
        self._prep_counts()
        for assay in self.assays:
            raw_data = assay.rawData[self.cellIdx]
            if self.outWorkspace is None:
                store = self.z[f"{assay.name}/counts"]
            else:
                store = self.z[f"matrices/{assay.name}/counts"]
            s, e, = (
                0,
                0,
            )
            for a in tqdmbar(
                raw_data.blocks,
                desc=f"Subsetting assay: {assay.name}",
                total=raw_data.numblocks[0],
            ):
                if a.shape[0] > 0:
                    e += a.shape[0]
                    store[s:e] = a.compute()
                    s = e


class ZarrMerge:
    """Merge multiple Zarr files into a single Zarr file.

    Args:
        zarr_path: Name of the new, merged Zarr file with path.
        assays: List of assay objects to be merged. For example, [ds1.RNA, ds2.RNA].
        names: Names of each of the assay objects in the `assays` parameter. They should be in the same order as in
               `assays` parameter.
        merge_assay_name: Name of assay in the merged Zarr file. For example, for scRNA-Seq it could be simply,
                          'RNA'.
        chunk_size: Tuple of cell and feature chunk size. (Default value: (1000, 1000)).
        dtype: Dtype of the raw values in the assay. Dtype is automatically inferred from the provided assays. If
               assays have different dtypes then a float type is used.
        overwrite: If True, then overwrites previously created assay in the Zarr file. (Default value: False).
        prepend_text: This text is pre-appended to each column name (Default value: 'orig').
        reset_cell_filter: If True, then the cell filtering information is removed, i.e. even the filtered out cells
                           are set as True as in the 'I' column. To keep the filtering information set the value for
                           this parameter to False. (Default value: True)

    Attributes:
        assays: List of assay objects to be merged. For example, [ds1.RNA, ds2.RNA].
        names: Names of each assay objects in the `assays` parameter.
        mergedCells:
        nCells: Number of cells in dataset.
        featCollection:
        mergedFeats:
        nFeats: Number of features in the dataset.
        featOrder:
        z: The merged Zarr file.
        assayGroup:
    """

    def __init__(
        self,
        zarr_path: str,
        assays: list,
        names: List[str],
        merge_assay_name: str,
        in_workspaces: Union[str, None] = None,
        out_workspace: Union[str, None] = None,
        chunk_size=(1000, 1000),
        dtype: str = None,
        overwrite: bool = False,
        prepend_text: str = "orig",
        reset_cell_filter: bool = True,
    ):
        self.assays = assays
        self.names = names
        self.inWorkspaces = in_workspaces
        self.outWorkspace = out_workspace
        self.mergedCells: pd.DataFrame = self._merge_cell_table(
            reset_cell_filter, prepend_text
        )
        self.nCells: int = self.mergedCells.shape[0]
        self.featCollection: List[Dict[str, str]] = self._get_feat_ids(assays)
        self.mergedFeats = self._merge_order_feats()
        self.nFeats = self.mergedFeats.shape[0]
        self.featOrder = self._ref_order_feat_idx()
        self.z = self._use_existing_zarr(zarr_path, merge_assay_name, overwrite)
        self._ini_cell_data(overwrite)
        if dtype is None:
            if len(set([str(x.rawData.dtype) for x in self.assays])) == 1:
                dtype = str(self.assays[0].rawData.dtype)
            else:
                dtype = "float"

        self.assayGroup = create_zarr_count_assay(
            z=self.z,
            assay_name=merge_assay_name,
            workspace=self.outWorkspace,
            chunk_size=chunk_size,
            n_cells=self.nCells,
            feat_ids=self.mergedFeats.index.values,
            feat_names=self.mergedFeats.names.values,
            dtype=dtype,
        )

    def _merge_cell_table(self, reset: bool, prepend_text: str) -> pd.DataFrame:
        """Merges the cell metadata table for each sample.

        Args:
            reset: whether to remove filtering information
            prepend_text: string to add as prefix for each cell column

        Returns:
        """
        # TODO: This method is not very memory efficient

        if len(self.assays) != len(set(self.names)):
            raise ValueError(
                "ERROR: A unique name should be provided for each of the assay"
            )
        if prepend_text == "":
            prepend_text = None
        ret_val = []
        for assay, name in zip(self.assays, self.names):
            a = assay.cells.to_pandas_dataframe(assay.cells.columns)
            a["ids"] = [f"{name}__{x}" for x in a["ids"]]
            for i in a.columns:
                if i not in ["ids", "I", "names"] and prepend_text is not None:
                    a[f"{prepend_text}_{i}"] = assay.cells.fetch_all(i)
                    a = a.drop(columns=[i])
            if reset:
                a["I"] = np.ones(len(a["ids"])).astype(bool)
            ret_val.append(a)
        ret_val = pd.concat(ret_val).reset_index(drop=True)
        if sum([x.cells.N for x in self.assays]) != ret_val.shape[0]:
            raise AssertionError(
                "Unexpected number of cells in the merged table. This is unexpected, "
                " please report this bug"
            )
        return ret_val

    @staticmethod
    def _get_feat_ids(assays) -> List[Dict[str, str]]:
        """Fetches ID->names mapping of features from each assay.

        Args:
            assays: List of Assay objects

        Returns:
            A list of dictionaries. Each dictionary is a id to name
            mapping for each feature in the corresponding assay
        """
        ret_val = []
        for i in assays:
            ret_val.append(
                i.feats.to_pandas_dataframe(["names", "ids"])
                .set_index("ids")["names"]
                .to_dict()
            )
        return ret_val

    def _merge_order_feats(self) -> pd.DataFrame:
        """Merge features from all the assays and determine their order.

        Returns:
        """
        union_set = {}
        for ids in self.featCollection:
            for i in ids:
                if i not in union_set:
                    union_set[i] = ids[i]
        ret_val = pd.DataFrame(
            {
                "idx": range(len(union_set)),
                "names": list(union_set.values()),
                "ids": list(union_set.keys()),
            }
        ).set_index("ids")
        r = ret_val.shape[0] / sum([x.feats.N for x in self.assays])
        if r == 1:
            raise ValueError(
                "No overlapping features found! Will not merge the files. Please check the features ids "
                " are comparable across the assays"
            )
        if r > 0.9:
            logger.warning("The number overlapping features is very low.")
        return ret_val

    def _ref_order_feat_idx(self) -> List[np.ndarray]:
        ret_val = []
        for ids in self.featCollection:
            ret_val.append(self.mergedFeats["idx"].reindex(list(ids.keys())).values)
        return ret_val

    def _use_existing_zarr(self, zarr_loc, merge_assay_name, overwrite):
        if self.outWorkspace is None:
            cell_slot = "cellData"
            assay_slot = merge_assay_name
        else:
            cell_slot = f"{self.outWorkspace}/cellData"
            assay_slot = f"{self.outWorkspace}/merge_assay_name"

        try:
            z = load_zarr(zarr_loc, mode="r")
            if cell_slot not in z:
                raise ValueError(
                    f"ERROR: Zarr file exists but seems corrupted. Either delete the "
                    "existing file or choose another path"
                )
            if assay_slot in z:
                if overwrite is False:
                    raise ValueError(
                        f"ERROR: Zarr file already contains {merge_assay_name} assay. Choose "
                        "a different zarr path or a different assay name. Otherwise set overwrite to True"
                    )
            try:
                if not all(
                    z[cell_slot]["ids"][:] == np.array(self.mergedCells["ids"].values)
                ):
                    raise ValueError(
                        f"ERROR: order of cells does not match the one in existing file"
                    )
            except KeyError:
                raise ValueError(
                    f"ERROR: 'cell data seems corrupted. Either delete the "
                    "existing file or choose another path"
                )
            return load_zarr(zarr_loc, mode="r+")
        except ValueError:
            # So no zarr file with same name exists. Check if a non zarr folder with the same name exists
            if isinstance(zarr_loc, str) and os.path.exists(zarr_loc):
                raise ValueError(
                    f"ERROR: Directory/file with name `{zarr_loc}`exists. "
                    f"Either delete it or use another name"
                )
            # creating a new zarr file
            return load_zarr(zarr_loc, mode="w")

    def _ini_cell_data(self, overwrite) -> None:
        """Save cell attributes to Zarr.

        Returns:
            None
        """
        if self.outWorkspace is None:
            cell_slot = "cellData"
        else:
            cell_slot = f"{self.outWorkspace}/cellData"

        if (cell_slot in self.z and overwrite is True) or cell_slot not in self.z:
            g = self.z.create_group(cell_slot, overwrite=True)
            for i in self.mergedCells:
                vals = self.mergedCells[i].values
                create_zarr_obj_array(g, i, vals, vals.dtype, overwrite=True)
        else:
            logger.info(f"cellData already exists so skipping _ini_cell_data")

    def dump(self, nthreads=2):
        """Copy the values from individual assays to the merged assay.

        Args:
            nthreads: Number of compute threads to use. (Default value: 2)

        Returns:
        """
        pos_start, pos_end = 0, 0
        for assay, feat_order in zip(self.assays, self.featOrder):
            for i in tqdmbar(
                assay.rawData.blocks,
                total=assay.rawData.numblocks[0],
                desc=f"Writing data to merged file",
            ):
                pos_end += i.shape[0]
                a = np.zeros((i.shape[0], self.nFeats))
                a[:, feat_order] = controlled_compute(i, nthreads)
                # FIXME: make this operation sparse
                self.assayGroup[pos_start:pos_end, :] = a
                pos_start = pos_end


def to_h5ad(
    assay,
    h5ad_filename: str,
    embeddings_cols: Optional[List[str]] = None,
    skip_recalc_nfeats: bool = True,
    n_threads: int = 4,
) -> None:
    """Save an assay as H5ad file.

    Args:
        assay: Assay to save in H5ad format
        h5ad_filename: Name for the H5ad file to be created.
        embeddings_cols: Columns in cell metadata to be treated as embeddings e. UMAP, tSNE
                         (Default value: ['UMAP', 'tSNE'])
        skip_recalc_nfeats: Skip recalculating nFeatures per cell. (Default value: True)
        n_threads: Number of processing threads to use (Default value: 4)

    Returns:
        None
    """
    import h5py
    from dask import array as da

    def save_attr(group, col, scarf_col, md):
        d = md.fetch_all(scarf_col)
        d_type = d.dtype
        if np.issubdtype(d_type, np.number) or np.issubdtype(d_type, bool):
            pass
        else:
            d_type = h5py.special_dtype(vlen=str)
        try:
            h5[group].create_dataset(col, data=d.astype(d_type))
        except TypeError:
            logger.warning(f"Dtype issue in {col}, {d.type} ({d_type})")

    h5 = h5py.File(h5ad_filename, "w")
    for i in ["X", "obs", "var", "obsm"]:
        h5.create_group(i)

    # Recalculating nFeature here just to avoid potential issues with stale data.
    if skip_recalc_nfeats is False:
        assay.cells.insert(
            f"{assay.name}_nFeatures",
            show_dask_progress(
                da.count_nonzero(assay.rawData, axis=1),
                msg="Preflight: recalculating nFeatures",
                nthreads=n_threads,
            ),
            overwrite=True,
        )

    n_feats_per_cell = assay.cells.fetch_all(f"{assay.name}_nFeatures").astype(int)
    tot_counts = int(n_feats_per_cell.sum())

    for i, s in zip(
        ["indptr", "indices", "data"], [assay.cells.N + 1, tot_counts, tot_counts]
    ):
        if i == "data":
            mat_dtype = assay.rawData.dtype
        else:
            mat_dtype = int
        h5["X"].create_dataset(
            i, (s,), chunks=True, compression="gzip", dtype=mat_dtype
        )

    h5["X/indptr"][:] = np.array([0] + list(n_feats_per_cell.cumsum())).astype(int)

    s, e = 0, 0
    for i in tqdmbar(
        assay.rawData.blocks,
        total=assay.rawData.numblocks[0],
        desc="Writing raw counts",
    ):
        i = csr_matrix(i.compute())
        e += i.data.shape[0]
        h5["X/data"][s:e] = i.data
        h5["X/indices"][s:e] = i.indices
        s = e
    attrs = {
        "encoding-type": "csr_matrix",
        "encoding-version": "0.1.0",
        "shape": np.array([assay.cells.N, assay.feats.N]),
    }
    for i, j in attrs.items():
        h5["X"].attrs[i] = j

    out_cols = []
    emb_cols = []
    if embeddings_cols is None:
        embeddings_cols = ["UMAP", "tSNE"]
    for i in assay.cells.columns:
        if i == "ids":
            save_attr("obs", "_index", "ids", assay.cells)
            out_cols.append("_index")
        else:
            is_emb = False
            if len(embeddings_cols) > 0:
                for j in embeddings_cols:
                    if i.startswith(f"{assay.name}_{j}"):
                        emb_cols.append(i)
                        is_emb = True
                        break
            if is_emb is False:
                save_attr("obs", i, i, assay.cells)
                out_cols.append(i)

    attrs = {
        "_index": "_index",
        "column-order": np.array(out_cols, dtype=object),
        "encoding-type": "dataframe",
        "encoding-version": "0.1.0",
    }
    for i, j in attrs.items():
        h5["obs"].attrs[i] = j

    out_cols = []
    for i in assay.feats.columns:
        if i == "ids":
            save_attr("var", "_index", "ids", assay.feats)
            out_cols.append("_index")
        elif i == "names":
            save_attr("var", "gene_short_name", "names", assay.feats)
            out_cols.append("gene_short_name")
        else:
            save_attr("var", i, i, assay.feats)
            out_cols.append(i)

    attrs = {
        "_index": "_index",
        "column-order": np.array(out_cols, dtype=object),
        "encoding-type": "dataframe",
        "encoding-version": "0.1.0",
    }
    for i, j in attrs.items():
        h5["var"].attrs[i] = j

    if len(emb_cols) > 0:
        emb_cols = np.array(emb_cols)
        c = pd.Series([x[:-1] for x in emb_cols])
        for i in c.unique():
            data = np.array([assay.cells.fetch_all(x) for x in emb_cols[c == i]]).T
            h5["obsm"].create_dataset(
                i.lower().replace(f"{assay.name.lower()}_", "X_"), data=data
            )

    h5.close()
    return None


def to_mtx(assay, mtx_directory: str, compress: bool = False):
    """Save an assay as a Matrix Market directory.

    Args:
        assay: Scarf assay. For example: `ds.RNA`
        mtx_directory: Out directory where MTX file will be saved along with barcodes and features file
        compress: If True, then the files are compressed and saved with .gz extension. (Default value: False).

    Returns:
        None
    """
    from scipy.sparse import coo_matrix
    import gzip

    if os.path.isdir(mtx_directory) is False:
        os.mkdir(mtx_directory)

    n_feats_per_cell = assay.cells.fetch_all(f"{assay.name}_nFeatures").astype(int)
    tot_counts = int(n_feats_per_cell.sum())
    if compress:
        barcodes_fn = "barcodes.tsv.gz"
        features_fn = "features.tsv.gz"
        h = gzip.open(os.path.join(mtx_directory, "matrix.mtx.gz"), "wt")
    else:
        barcodes_fn = "barcodes.tsv"
        features_fn = "genes.tsv"
        h = open(os.path.join(mtx_directory, "matrix.mtx"), "w")
    h.write("%%MatrixMarket matrix coordinate integer general\n% Generated by Scarf\n")
    h.write(f"{assay.feats.N} {assay.cells.N} {tot_counts}\n")
    s = 0
    for i in tqdmbar(assay.rawData.blocks, total=assay.rawData.numblocks[0]):
        i = coo_matrix((i.compute()))
        df = pd.DataFrame({"col": i.col + 1, "row": i.row + s + 1, "d": i.data})
        try:
            # Support for older versions of pandas
            df.to_csv(
                h, sep=" ", header=False, index=False, mode="a", lineterminator="\n"
            )
        except TypeError:
            df.to_csv(
                h, sep=" ", header=False, index=False, mode="a", line_terminator="\n"
            )
        s += i.shape[0]
    h.close()
    assay.cells.to_pandas_dataframe(["ids"]).to_csv(
        os.path.join(mtx_directory, barcodes_fn), sep="\t", header=False, index=False
    )

    assay.feats.to_pandas_dataframe(["ids", "names"]).to_csv(
        os.path.join(mtx_directory, features_fn), sep="\t", header=False, index=False
    )


def bed_to_sparse_array(
    bed_fn: str,
    bin_size: int,
    chrom_sizes: Dict[str, int],
    min_counts_per_cell: int = 500,
    read_chunk_size=1e6,
    sep: str = "\t",
    chrom_col: int = 0,
    start_col: int = 1,
    end_col: int = 2,
    barcode_col: int = 3,
    count_col: int = 4,
    comments_startswith: str = "#",
    disable_tqdm: bool = False,
    chrom_modifier=None,
):
    """

    Args:
        bed_fn:
        bin_size:
        chrom_sizes:
        min_counts_per_cell:
        read_chunk_size:
        sep:
        chrom_col:
        start_col:
        end_col:
        barcode_col:
        count_col:
        comments_startswith:
        disable_tqdm:
        chrom_modifier:

    Returns:

    """
    import gc

    def feat_mapper(x):
        return feat_idx.get(x, n_feats)

    def default_chrom_modifier(x):
        return x + "_"

    feat_idx = {}
    for i in tqdmbar(chrom_sizes, disable=disable_tqdm, desc="Calculating bin indices"):
        for j in range((chrom_sizes[i] // bin_size) + 1):
            feat_idx[f"{i}_{j}"] = len(feat_idx)
    cell_idx = {}
    mat = []
    n_feats = len(feat_idx)
    if chrom_modifier is None:
        chrom_modifier = default_chrom_modifier

    stream = pd.read_csv(
        bed_fn,
        sep=sep,
        header=None,
        comment=comments_startswith,
        usecols=[chrom_col, start_col, end_col, barcode_col, count_col],
        chunksize=int(read_chunk_size),
    )
    for df in tqdmbar(
        stream, disable=disable_tqdm, desc="Building in memory sparse matrix"
    ):
        df[chrom_col] = df[chrom_col].map(chrom_modifier) + (
            (df[start_col] + (df[end_col] - df[start_col]) // 2).values // bin_size
        ).astype(str)
        for i in df[barcode_col].unique():
            if i not in cell_idx:
                cell_idx[i] = len(cell_idx)
        mat.append(
            np.vstack(
                [
                    np.fromiter(map(cell_idx.get, df[barcode_col].values), dtype=int),
                    np.fromiter(map(feat_mapper, df[chrom_col].values), dtype=int),
                    df[count_col].values,
                ]
            ).T
        )
    mat = np.vstack(mat)
    gc.collect()
    mat = csr_matrix(
        (mat[:, 2], (mat[:, 0], mat[:, 1])), shape=(len(cell_idx), n_feats + 1)
    )
    gc.collect()
    idx = np.array(mat.sum(axis=1))[:, 0] > min_counts_per_cell
    return mat[idx, :-1], pd.Series(cell_idx.keys())[idx], pd.Series(feat_idx.keys())
