# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/CLI/dask_.ipynb.

# %% auto 0
__all__ = ['parallel_read_zarr', 'dask_from_zarr', 'dask_from_zarr_overlap', 'ZarrDir', 'parallel_write_zarr', 'dask_to_zarr']

# %% ../../nbs/CLI/dask_.ipynb 3
import math
import itertools
import concurrent
import logging
from pathlib import Path

import numpy as np
from dask import array as da
from dask import delayed
import zarr

from .logging import mc_logger_

# %% ../../nbs/CLI/dask_.ipynb 4
def _fill_slice(data_zarr,slices):
    out_slices = []
    for i in range(len(slices)):
        slice_i = slices[i]
        if slice_i.start is None:
            start = 0
        else:
            start = slice_i.start
        assert start>=0
        assert start<data_zarr.shape[i]
        #assert start%data_zarr.chunks[i] == 0
        if slice_i.stop is None:
            stop = data_zarr.shape[i]
        else:
            stop = slice_i.stop
        assert stop > start
        assert stop <=data_zarr.shape[i]
            #assert stop%data_zarr.chunks[i] == 0
        assert (slice_i.step is None) or (slice_i.step == 1)
        step = 1
        out_slices.append(slice(start,stop,step))
    return tuple(out_slices)

# %% ../../nbs/CLI/dask_.ipynb 5
def _one_chunk_slices_and_out_shape(data_zarr,slices):
    zarr_1chunk_slice = []
    out_1chunk_slice = []
    out_shape = []
    for i in range(len(slices)):
        chunk_size = data_zarr.chunks[i]
        shape = data_zarr.shape[i]
        zarr_slices = slices[i]
        start = zarr_slices.start
        stop = zarr_slices.stop
        out_shape.append(stop-start)

        start_ = math.floor(start/chunk_size)*chunk_size    
        zarr_1dim_bound = np.arange(start_,stop+chunk_size,chunk_size)
        zarr_1dim_bound[0] = start
        if zarr_1dim_bound[-1] > stop: zarr_1dim_bound[-1] = stop
        out_1dim_bound = zarr_1dim_bound-start

        zarr_1chunk_1dim_slice = []
        out_1chunk_1dim_slice = []
        for j in range(zarr_1dim_bound.shape[0]-1):
            zarr_1chunk_1dim_slice.append(slice(zarr_1dim_bound[j],zarr_1dim_bound[j+1]))
            out_1chunk_1dim_slice.append(slice(out_1dim_bound[j],out_1dim_bound[j+1]))

        zarr_1chunk_slice.append(zarr_1chunk_1dim_slice)
        out_1chunk_slice.append(out_1chunk_1dim_slice)

    zarr_1chunk_slice = list(itertools.product(*zarr_1chunk_slice))
    out_1chunk_slice = list(itertools.product(*out_1chunk_slice))

    return zarr_1chunk_slice, out_1chunk_slice, tuple(out_shape)

# %% ../../nbs/CLI/dask_.ipynb 6
def _read_one_chunk(data_zarr,out,zarr_slices,out_slices):
    out[out_slices] = data_zarr[zarr_slices]

# %% ../../nbs/CLI/dask_.ipynb 7
def parallel_read_zarr(data_zarr,slices,thread_pool_size=None,fill_slice=True):
    if fill_slice:
        slices = _fill_slice(data_zarr,slices)
    zarr_1chunk_slices, out_1chunk_slices, out_shape = \
    _one_chunk_slices_and_out_shape(data_zarr,slices)
    # global out
    out = np.empty(out_shape,dtype=data_zarr.dtype)

    with concurrent.futures.ThreadPoolExecutor(thread_pool_size) as executor:
        for zarr_1chunk_slice,out_1chunk_slice in zip(zarr_1chunk_slices,out_1chunk_slices):
            future = executor.submit(_read_one_chunk,data_zarr,out,zarr_1chunk_slice,out_1chunk_slice)
    return out

# %% ../../nbs/CLI/dask_.ipynb 14
def _all_chunk_slices(data_zarr, chunks):
    out_slices = []
    for i in range(data_zarr.ndim):
        shape = data_zarr.shape[i]
        chunk = chunks[i]
        if chunk <0: chunk=shape
        bound_1dim = np.arange(0,shape+chunk,chunk)
        if bound_1dim[-1] > shape: bound_1dim[-1] = shape

        slice_1dim = []
        for j in range(bound_1dim.shape[0]-1):
            slice_1dim.append(slice(bound_1dim[j],bound_1dim[j+1]))
        out_slices.append(slice_1dim)
    out_slices = list(itertools.product(*out_slices))
    return out_slices

# %% ../../nbs/CLI/dask_.ipynb 15
def dask_from_zarr(
    data_path:str,
    parallel_dims:int|tuple=None,
    chunks:tuple=None,
):
    # chunks provide more control of chunking compared with parallel_dims
    data_zarr = zarr.open(data_path,'r')
    if chunks is None: chunks = data_zarr.chunks
    chunks = list(chunks)
    if parallel_dims is not None:
        if isinstance(parallel_dims,int):
            parallel_dims = (parallel_dims,)

        for pdim in parallel_dims:
            chunks[pdim] = data_zarr.shape[pdim]
    for i in range(len(chunks)):
        if chunks[i] <0: chunks[i] = data_zarr.shape[i]
    chunks = tuple(chunks)

    slices = _all_chunk_slices(data_zarr,chunks)

    _p_read_zarr_delayed = delayed(parallel_read_zarr,pure=True,nout=1)
    out_delayed = np.empty(len(slices),dtype=object)
    for i in range(len(slices)):
        slices_i = _fill_slice(data_zarr,slices[i])
        out_delayed[i] = _p_read_zarr_delayed(data_zarr,slices_i,fill_slice=False)
        shape = []
        for slice_ in slices_i:
            shape.append(slice_.stop-slice_.start)
        shape = tuple(shape)
        out_delayed[i] = da.from_delayed(out_delayed[i],shape=shape,dtype=data_zarr.dtype,meta=np.array((),dtype=data_zarr.dtype))

    nchunks = [math.ceil(shape/chunk) for chunk, shape in zip(chunks,data_zarr.shape)]
    out_delayed = out_delayed.reshape(tuple(nchunks))
    data_darr = da.block(out_delayed.tolist())
    return data_darr

# %% ../../nbs/CLI/dask_.ipynb 20
def _all_chunk_slices_with_overlap(data_zarr, chunks, depths):
    out_slices = []
    for i in range(data_zarr.ndim):
        shape = data_zarr.shape[i]
        chunk = chunks[i]
        if chunk <0: chunk=shape
        depth = depths[i]
        starts_1dim = np.arange(-depth,shape-depth,chunk)
        starts_1dim[starts_1dim<0] = 0
        ends_1dim = np.arange(chunk+depth,shape+chunk+depth,chunk)
        ends_1dim[ends_1dim>shape] = shape

        slice_1dim = []
        for j in range(starts_1dim.shape[0]):
            slice_1dim.append(slice(starts_1dim[j],ends_1dim[j]))
        out_slices.append(slice_1dim)
    out_slices = list(itertools.product(*out_slices))
    return out_slices

# %% ../../nbs/CLI/dask_.ipynb 22
def dask_from_zarr_overlap(
    data_path:str,
    chunks:tuple,
    depth:dict|tuple,
):
    # chunks provide more control of chunking compared with parallel_dims
    data_zarr = zarr.open(data_path,'r')

    chunks = list(chunks)
    for i in range(len(chunks)):
        if chunks[i] <0: chunks[i] = data_zarr.shape[i]
    chunks = tuple(chunks)

    slices = _all_chunk_slices_with_overlap(data_zarr,chunks,depth)

    _p_read_zarr_delayed = delayed(parallel_read_zarr,pure=True,nout=1)
    out_delayed = np.empty(len(slices),dtype=object)
    for i in range(len(slices)):
        slices_i = _fill_slice(data_zarr,slices[i])
        out_delayed[i] = _p_read_zarr_delayed(data_zarr,slices_i,fill_slice=False)
        shape = []
        for slice_ in slices_i:
            shape.append(slice_.stop-slice_.start)
        shape = tuple(shape)
        out_delayed[i] = da.from_delayed(out_delayed[i],shape=shape,dtype=data_zarr.dtype,meta=np.array((),dtype=data_zarr.dtype))

    nchunks = [math.ceil(shape/chunk) for chunk, shape in zip(chunks,data_zarr.shape)]
    out_delayed = out_delayed.reshape(tuple(nchunks))
    data_darr = da.block(out_delayed.tolist())
    return data_darr

# %% ../../nbs/CLI/dask_.ipynb 25
class ZarrDir():
    def __init__(self,zarr_path_list:list):
        self.zarr_path_list = zarr_path_list
        zarr0 = zarr.open(self.zarr_path_list[0])
        self.chunksize = zarr0.chunks
        self.ndim = zarr0.ndim
        self.dtype = zarr0.dtype
        
        dim0_shape = 0
        dim0_chunks = []
        for zarr_path in self.zarr_path_list:
            data_zarr = zarr.open(zarr_path,'r')
            assert data_zarr.chunks[0] == data_zarr.shape[0]
            assert data_zarr.chunks[1:] == (1,)*(data_zarr.ndim-1)
            dim0_shape += data_zarr.shape[0]
            dim0_chunks.append(data_zarr.chunks[0])
        self.shape = (dim0_shape,*zarr0.shape[1:])
        self.dim0_chunks = tuple(dim0_chunks)

    @classmethod
    def from_dir(cls,zarr_dir):
        zarr_dir = Path(zarr_dir)
        zarr_path_list = sorted(zarr_dir.glob('*.zarr'),key=lambda path: int(path.stem)) # if one chunk is missing, it is ok
        return cls(zarr_path_list)

# %% ../../nbs/CLI/dask_.ipynb 26
def _parallel_read_pc_dir(
    zarr_dir, # zarr_dir object
    idx, # index for dim 1,2,...
    thread_pool_size=None):
    # global out
    if isinstance(idx,int):
        idx = (idx,)
    out = np.empty(zarr_dir.shape[0],dtype=zarr_dir.dtype)
    out_slices = []
    start = 0
    for chunk in zarr_dir.dim0_chunks:
        end = start+chunk
        out_slices.append(slice(start,end))
        start = end

    def _read_one_zarr_one_chunk(zarr_path,out,idx,out_slice):
        data_zarr = zarr.open(zarr_path,'r')
        idx = (slice(None),*idx)
        out[out_slice] = data_zarr[idx]

    with concurrent.futures.ThreadPoolExecutor(thread_pool_size) as executor:
        for zarr_path,out_slice in zip(zarr_dir.zarr_path_list,out_slices):
            future = executor.submit(_read_one_zarr_one_chunk,zarr_path,out,idx,out_slice)
    return out

# %% ../../nbs/CLI/dask_.ipynb 29
def _dask_from_pc_zarr_dir(zarrs):
    if isinstance(zarrs,list):
        zarr_dir = ZarrDir(zarrs)
    else:
        zarr_dir = ZarrDir.from_dir(zarrs)
    out_shape = zarr_dir.shape[1:]
    out_chunk_shape = (1,)*(zarr_dir.ndim-1)
    out_delayed = np.empty((1,*out_shape),dtype=object)
    parallel_read_delayed = delayed(_parallel_read_pc_dir,pure=True,nout=1)
    with np.nditer(out_delayed,flags=['multi_index','refs_ok'], op_flags=['readwrite']) as arr_it:
        for arr_block in arr_it:
            idx = arr_it.multi_index
            out_delayed[idx] = parallel_read_delayed(zarr_dir,idx[1:])
            out_delayed[idx] = da.from_delayed(out_delayed[idx],shape=(zarr_dir.shape[0],),meta=np.array((),dtype=zarr_dir.dtype))
            out_delayed[idx] = out_delayed[idx].reshape((zarr_dir.shape[0],*out_chunk_shape))
    return da.block(out_delayed.tolist())

# %% ../../nbs/CLI/dask_.ipynb 31
def _write_one_chunk(data_zarr,data,zarr_slices,data_slices):
    data_zarr[zarr_slices] = data[data_slices]

# %% ../../nbs/CLI/dask_.ipynb 32
def parallel_write_zarr(data,data_zarr,slices,thread_pool_size=None,fill_slice=True):
    if fill_slice:
        slices = _fill_slice(data_zarr,slices)
    zarr_1chunk_slices, data_1chunk_slices = \
    _one_chunk_slices_and_out_shape(data_zarr,slices)[:2]

    with concurrent.futures.ThreadPoolExecutor(thread_pool_size) as executor:
        for zarr_1chunk_slice,data_1chunk_slice in zip(zarr_1chunk_slices,data_1chunk_slices):
            future = executor.submit(_write_one_chunk,data_zarr,data,zarr_1chunk_slice,data_1chunk_slice)

# %% ../../nbs/CLI/dask_.ipynb 37
def _all_chunk_slices_except_pdims(data_zarr, pdims):
    out_slices = []
    for i in range(data_zarr.ndim):
        if i in pdims:
            slice_1dim = [slice(None)]
        else:
            shape = data_zarr.shape[i]
            chunk = data_zarr.chunks[i]
            if chunk <0: chunk=shape

            bound_1dim = np.arange(0,shape+chunk,chunk)
            if bound_1dim[-1] > shape: bound_1dim[-1] = shape

            slice_1dim = []
            for j in range(bound_1dim.shape[0]-1):
                slice_1dim.append(slice(bound_1dim[j],bound_1dim[j+1]))
        out_slices.append(slice_1dim)
    out_slices = list(itertools.product(*out_slices))
    return out_slices

# %% ../../nbs/CLI/dask_.ipynb 38
@mc_logger_
def dask_to_zarr(data_darr,url,chunks,path=None,log_zarr=True):
    logger = logging.getLogger(__name__)
    if isinstance(url, zarr.Array):
        data_zarr = url
        if log_zarr:
            logger.zarr_info(path,data_zarr)
    else:
        data_zarr = zarr.open(url,'w',shape=data_darr.shape,dtype=data_darr.dtype,chunks=chunks)
        if log_zarr:
            logger.zarr_info(url,data_zarr)
    pdims = []
    for i in range(data_darr.ndim):
        if data_darr.chunksize[i] != data_zarr.chunks[i]:
            assert data_darr.chunksize[i] == data_darr.shape[i], "dask array is not properly rechunked."
            pdims.append(i)
    pdims = tuple(pdims)
    slices = _all_chunk_slices_except_pdims(data_zarr,pdims)
    n_slices = len(slices)

    _p_write_zarr_delayed = delayed(parallel_write_zarr,pure=True,nout=0)
    out_delayed = np.empty(n_slices,dtype=object)
    data_darr_delayed = data_darr.to_delayed()
    data_darr_delayed = data_darr_delayed.reshape(n_slices)
    for i in range(n_slices):
        slices_i = _fill_slice(data_zarr,slices[i])
        out_delayed[i] = _p_write_zarr_delayed(data_darr_delayed[i],data_zarr,slices_i,fill_slice=False)
        out_delayed[i] = da.from_delayed(out_delayed[i],shape=(1,),dtype=int)

    return da.block(out_delayed.tolist())
