"""internal utilities for chunkwise data processing"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/chunk_.ipynb.

# %% auto 0
__all__ = ['fill_slice', 'all_chunk_slices', 'all_chunk_slices_with_overlap', 'chunkwise_slicing_mapping',
           'chunkwise_knn_mapping']

# %% ../nbs/API/chunk_.ipynb 3
import numpy as np
import itertools
from .rtree import HilbertRtree
from scipy.spatial import KDTree

# %% ../nbs/API/chunk_.ipynb 4
def fill_slice(
    shape, # numpy arr, cupy arr, zarr,
    slices, # tuple of slice object, len == data_arr.ndim
):
    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<shape[i]
        if slice_i.stop is None:
            stop = shape[i]
        else:
            stop = slice_i.stop
        assert stop > start
        assert stop <= shape[i]
        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/API/chunk_.ipynb 5
def all_chunk_slices(
    shape, # np.array, cp.array,zarr
    chunks,
):
    '''get the slices for every input chunks'''
    out_slices = []
    for shape_, chunk_ in zip(shape,chunks):
        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(int(bound_1dim[j]),int(bound_1dim[j+1])))
        out_slices.append(slice_1dim)
    out_slices = list(itertools.product(*out_slices))
    return out_slices

# %% ../nbs/API/chunk_.ipynb 9
def all_chunk_slices_with_overlap(
    shape, 
    chunks, 
    depths,
):
    '''get the slices for every input chunks with overlap'''
    out_slices = []
    for shape_, chunk_, depth_ in zip(shape,chunks,depths):
        if chunk_ <0: chunk_ = shape_
        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(int(starts_1dim[j]),int(ends_1dim[j])))
        out_slices.append(slice_1dim)
    out_slices = list(itertools.product(*out_slices))
    return out_slices

# %% ../nbs/API/chunk_.ipynb 11
def chunkwise_slicing_mapping(
    shape,
    chunks,
    depths,
):
    '''get the slices for every input chunks with overlap
    output chunks without overlap and their mapping slices'''
    in_slices = []; out_slices = []; map_slices = []
    for shape_, chunk_, depth_ in zip(shape,chunks,depths):
        if chunk_ <0: chunk_ = shape_
        in_starts_1dim = np.arange(-depth_,shape_-depth_,chunk_)
        in_starts_1dim[in_starts_1dim<0] = 0
        in_ends_1dim = np.arange(chunk_+depth_,shape_+chunk_+depth_,chunk_)
        in_ends_1dim[in_ends_1dim>shape_] = shape_
        out_starts_1dim = np.arange(0,shape_,chunk_)
        out_ends_1dim = np.arange(chunk_,shape_+chunk_,chunk_)
        out_ends_1dim[out_ends_1dim>shape_] = shape_

        assert len(in_starts_1dim) == len(out_starts_1dim)
        assert len(in_ends_1dim) == len(out_ends_1dim)

        in_slice_1dim = []; out_slice_1dim = []; map_slice_1dim = []
        for in_start, in_end, out_start, out_end in zip(in_starts_1dim, in_ends_1dim, out_starts_1dim, out_ends_1dim):
            in_slice_1dim.append(slice(int(in_start),int(in_end)))
            out_slice_1dim.append(slice(int(out_start),int(out_end)))
            offset = out_start-in_start
            map_slice_1dim.append(slice(int(offset),int(offset+out_end-out_start)))

        in_slices.append(in_slice_1dim)
        out_slices.append(out_slice_1dim)
        map_slices.append(map_slice_1dim)

    in_slices = list(itertools.product(*in_slices))
    out_slices = list(itertools.product(*out_slices))
    map_slices = list(itertools.product(*map_slices))
    return in_slices, out_slices, map_slices

# %% ../nbs/API/chunk_.ipynb 13
def chunkwise_knn_mapping(
    x,
    y,
    chunks:int,
    k = 128
):
    n = y.shape[0]
    bound = np.arange(0,n+chunks,chunks)
    if bound[-1] > n: bound[-1] = n
    rtree = HilbertRtree.build(x,y,page_size=chunks)

    out_slices = []
    in_indices = []
    map_indices = []
    for i in range(bound.shape[0]-1):
        start = int(bound[i])
        end = int(bound[i+1])
        out_slices.append(slice(start,end))

        y_, x_ = y[start:end], x[start:end]
        # normalize x and y here should be better and save the ratio!
        y_max, y_min, x_max, x_min = y_.max(), y_.min(), x_.max(), x_.min()
        y_norm = (y_ - y_min)/(y_max - y_min)
        x_norm = (x_ - x_min)/(x_max - x_min)
        pos_norm = np.stack((x_norm,y_norm),axis=-1)

        # select boundary points
        if chunks<2*k:
            bound_idx = np.arange(chunks) # select all points as boundary points
        else:
            bound_ratio = 0.1
            while True:
                h_y, h_x = bound_ratio, bound_ratio
                mask_boundary = (y_norm > 1-h_y) | (y_norm < h_y)|(x_norm > 1-h_x)|(x_norm < h_x)
                bound_idx = np.nonzero(mask_boundary)[0]
                bound_ratio = bound_ratio + 0.1
                if len(bound_idx)>=k: break

        # esimate distance of kth NN of the boundary points
        tree = KDTree(pos_norm)
        dd = tree.query(pos_norm[bound_idx],k=[k,],workers=-1)[0]
        max_dd = np.max(dd)

        # decide halo size and query halo points
        y_halo_size, x_halo_size = max_dd*(y_max-y_min), max_dd*(x_max-x_min)
        halo_bounds = [x_min-x_halo_size, y_min-y_halo_size, x_max+x_halo_size, y_max+y_halo_size]
        halo_idx = rtree.bbox_query(halo_bounds, x, y)
        #print(halo_idx.shape)
        y_halo, x_halo = y[halo_idx], x[halo_idx]

        # find input points: union of knn in halo points for points in the chunk 
        tree = KDTree(np.stack((x_halo, y_halo),axis=-1))
        in_idx = tree.query(np.stack((x_, y_),axis=-1),k=k,workers=-1)[1]
        in_idx = np.unique(in_idx)
        in_idx = halo_idx[in_idx] # idx in all points
        in_indices.append(in_idx)
        # assert np.all(np.isin(np.arange(start,end),in_idx))

        # indices that map points in the chunk from input points
        map_idx = (in_idx>=start) & (in_idx<end)
        map_idx = np.flatnonzero(map_idx)
        map_indices.append(map_idx)

    return in_indices, out_slices, map_indices
