"""Point Cloud data manipulation"""

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

# %% auto 0
__all__ = ['pc_hix', 'pc_gix', 'pc_sort', 'pc2ras', 'ras2pc', 'pc_union', 'pc_intersect', 'pc_diff']

# %% ../nbs/API/pc.ipynb 7
import numpy as np
from numba import prange
import math
from .utils_ import is_cuda_available, get_array_module
if is_cuda_available():
    import cupy as cp
from .utils_ import ngjit, ngpjit
from .coord_ import Coord

# %% ../nbs/API/pc.ipynb 9
def _ras_dims(
    gix1:np.ndarray, # int array, grid index of the first point cloud
    gix2:np.ndarray=None, # int array, grid index of the second point cloud
)->tuple: # the shape of the original raster image
    '''Get the shape of the original raster image from two index, the shape could be smaller than the truth but it doesn't matter.'''
    xp = get_array_module(gix1)
    if gix2 is None:
        dims_az = gix1[-1,0]+1
        dims_r = int(xp.max(gix1[:,1]))+1
    else:
        dims_az = max(int(gix1[-1,0]),int(gix2[-1,0]))+1
        dims_r = max(int(xp.max(gix1[:,1])),int(xp.max(gix2[:,1])))+1
    return (dims_az,dims_r)

# %% ../nbs/API/pc.ipynb 10
def _ravel_gix(gix,dims):
    xp = get_array_module(gix)
    return xp.ravel_multi_index((gix[:,0],gix[:,1]),dims=dims)

# %% ../nbs/API/pc.ipynb 11
def _unravel_gix(gix_1d,dims):
    xp = get_array_module(gix_1d)
    return xp.stack(xp.unravel_index(gix_1d,dims),axis=-1).astype(xp.int32)

# %% ../nbs/API/pc.ipynb 13
def _check_idx_sorted(idx,shape=None):
    xp = get_array_module(idx)
    if idx.ndim == 2:
        idx_1d = _ravel_gix(idx,dims=shape)
    else:
        idx_1d = idx
    assert (xp.diff(idx_1d)>0).any(), "idx is not sorted or unique!"

# %% ../nbs/API/pc.ipynb 15
# Some functions adapted from spatialpandas at https://github.com/holoviz/spatialpandas under BSD-2-Clause license,
# Which is Initially based on https://github.com/galtay/hilbert_curve, but specialized
# for 2 dimensions with numba acceleration

# Further enhanced with numba prange.
@ngjit
def _int_2_binary(v, width): # Return a binary byte array representation of `v` zero padded to `width` bits.
    res = np.zeros(width, dtype=np.uint8)
    for i in range(width):
        res[width - i - 1] = v % 2
        v = v >> 1
    return res
@ngjit
def _binary_2_int(bin_vec):
    # Convert a binary byte array to an integer
    res = 0
    next_val = 1
    width = len(bin_vec)
    for i in range(width):
        res += next_val*bin_vec[width - i - 1]
        next_val <<= 1
    return res
@ngjit
def _hilbert_integer_to_transpose(p:int, # iterations to use in the hilbert curve
                                  h:int, # integer distance along hilbert curve
                                 )->np.array: # x (two element np.array): transpose of h (n components with values between 0 and 2**p-1)
    #Store a hilbert integer (`h`) as its transpose (`x`).
    h_bits = _int_2_binary(h, 2*p)
    x0, x1 = _binary_2_int(h_bits[0::2]), _binary_2_int(h_bits[1::2])
    x = np.array([x0,x1])
    return x
@ngjit
def _transpose_to_hilbert_integer(p:int, # iterations to use in the hilbert curve
                                  coord:list, # transpose of h (2 components with values between 0 and 2**p-1)
                                 )->int: # h (int): integer distance along hilbert curve
    # Restore a hilbert integer (`h`) from its transpose (`x`).
    bins = [_int_2_binary(v, p) for v in coord]
    concat = np.zeros(2*p, dtype=np.uint8)
    for i in range(p):
        for j in range(2):
            concat[2*i + j] = bins[j][i]

    h = _binary_2_int(concat)
    return h

@ngjit
def _coordinate_from_distance(p:int, # iterations to use in the hilbert curve
                              h:int, # integer distance along hilbert curve
                             )->list: # coord (list): Coordinate as length-n list
    # Return the coordinate for a hilbert distance.
    coord = _hilbert_integer_to_transpose(p, h)
    Z = 2 << (p-1)

    # Gray decode by H ^ (H/2)
    t = coord[1] >> 1
    for i in range(1, 0, -1):
        coord[i] ^= coord[i-1]
    coord[0] ^= t

    # Undo excess work
    Q = 2
    while Q != Z:
        P = Q - 1
        for i in range(1, -1, -1):
            if coord[i] & Q:
                # invert
                coord[0] ^= P
            else:
                # exchange
                t = (coord[0] ^ coord[i]) & P
                coord[0] ^= t
                coord[i] ^= t
        Q <<= 1

    return coord
@ngjit
def _distance_from_coordinate(p:int, # iterations to use in the hilbert curve
                              coord:np.ndarray, # coordinate as 1d array
                             )->int: # distance
    # Return the hilbert distance for a given coordinate.
    M = 1 << (p - 1)
    # Inverse undo excess work
    Q = M
    while Q > 1:
        P = Q - 1
        for i in range(2):
            if coord[i] & Q:
                coord[0] ^= P
            else:
                t = (coord[0] ^ coord[i]) & P
                coord[0] ^= t
                coord[i] ^= t
        Q >>= 1
    # Gray encode
    for i in range(1, 2):
        coord[i] ^= coord[i - 1]
    t = 0
    Q = M
    while Q > 1:
        if coord[1] & Q:
            t ^= Q - 1
        Q >>= 1
    for i in range(2):
        coord[i] ^= t
    h = _transpose_to_hilbert_integer(p, coord)
    return h
@ngpjit
def _coordinates_from_distances(p:int, # iterations to use in the hilbert curve
                                h:np.ndarray, # 1d array of integer distances along hilbert curve
                               )->np.ndarray: # 2d array of coordinate, each row a coordinate corresponding to associated distance value in input.
    # Return the coordinates for an array of hilbert distances.
    result = np.zeros((len(h), 2), dtype=np.int32)
    for i in prange(len(h)):
        result[i] = _coordinate_from_distance(p, h[i])
    return result
@ngpjit
def _distances_from_coordinates(p:int, # iterations to use in the hilbert curve
                                coords:np.ndarray, # 2d array of coordinates, one coordinate per row
                               )->np.ndarray: # 1d array of distances
    # Return the hilbert distances for a given set of coordinates.
    coords = np.atleast_2d(coords).copy()
    result = np.zeros(coords.shape[0], dtype=np.int64)
    for i in prange(coords.shape[0]):
        coord = coords[i]
        result[i] = _distance_from_coordinate(p, coord)
    return result

# %% ../nbs/API/pc.ipynb 16
def pc_hix(
    gix, # grid index
    shape:tuple, # (nlines, width)
)->np.ndarray:
    '''Compute the hillbert index for point cloud data.
    '''
    ny, nx = shape
    p = math.ceil(math.log2(max(nx,ny)))
    hix = _distances_from_coordinates(p, gix)
    return hix

# %% ../nbs/API/pc.ipynb 21
def pc_gix(
    hix, # hillbert index
    shape:tuple, # (nlines, width)
)->np.ndarray:
    '''convert hillbert index to grid index.
    '''
    ny, nx = shape
    p = math.ceil(math.log2(max(nx,ny)))
    gix = _coordinates_from_distances(p, hix)
    return gix

# %% ../nbs/API/pc.ipynb 23
def pc_sort(idx:np.ndarray, # unsorted `gix` (2D) or `hix`(1D)
            shape:tuple=None, # (nline, width), faster if provided for grid index input
           )->np.ndarray: # indices that sort input
    '''Get the indices that sort the input.'''
    xp = get_array_module(idx)
    if idx.ndim == 2:
        if shape is not None:
            dims_az, dims_r = shape
        else:
            dims_az = int(xp.max(idx[:,0]))+1
            dims_r = int(xp.max(idx[:,1]))+1
        idx_1d = xp.ravel_multi_index((idx[:,0],idx[:,1]),dims=(dims_az,dims_r))
    else:
        idx_1d = idx
    key = xp.argsort(idx_1d,kind='stable')
    return key

# %% ../nbs/API/pc.ipynb 27
def pc2ras(idx:np.ndarray, # gix or hix array
           pc_data:np.ndarray, # data, 1D or more
           shape:tuple, # image shape
          ):
    '''convert point cloud data to original raster, filled with nan'''
    xp = get_array_module(pc_data)
    raster = xp.empty((*shape,*pc_data.shape[1:]),dtype=pc_data.dtype)
    raster[:] = xp.nan
    if idx.ndim == 2:
        gix = idx
    else:
        gix = pc_gix(idx, shape)
    raster[gix[:,0],gix[:,1]] = pc_data
    return raster

# %% ../nbs/API/pc.ipynb 29
def ras2pc(ras,gix):
    return ras[gix[:,0],gix[:,1]]

# %% ../nbs/API/pc.ipynb 30
@ngjit
def _pc_union_numba(idx1, idx2):
    # hand write merge sort
    size1 = idx1.shape[0]
    size2 = idx2.shape[0]
    idx = np.empty(size1+size2,dtype=idx1.dtype)
    inv_iidx1 = np.empty_like(idx1) # 0:size1 for inv_iidx1, size1: for inv_iidx2
    inv_iidx2 = np.empty_like(idx2)
    ninv_iidx2 = np.empty_like(idx2)
    i1 = 0; i2 = 0; i = 0; inv1_i = 0; inv2_i = 0; ninv2_i = 0;
    while((i1<size1) and (i2<size2)):
        if idx1[i1] < idx2[i2]:
            idx[i] = idx1[i1]
            inv_iidx1[inv1_i] = i
            i1 +=1; inv1_i += 1
        elif idx1[i1] > idx2[i2]:
            idx[i] = idx2[i2]
            inv_iidx2[inv2_i] = i
            ninv_iidx2[ninv2_i] = i2
            i2 +=1; inv2_i += 1; ninv2_i += 1
        else:
            idx[i] = idx1[i1]
            inv_iidx1[inv1_i] = i
            i1 += 1; inv1_i += 1; i2 += 1
        i +=1
    while(i1<size1):
        idx[i] = idx1[i1]
        inv_iidx1[inv1_i] = i
        i1 += 1; inv1_i += 1; i += 1
    while(i2<size2):
            idx[i] = idx2[i2]
            inv_iidx2[inv2_i] = i
            ninv_iidx2[ninv2_i] = i2
            i2 +=1; inv2_i += 1; ninv2_i += 1; i += 1

    return idx[:i], inv_iidx1[:inv1_i], inv_iidx2[:inv2_i], ninv_iidx2[:ninv2_i]

# %% ../nbs/API/pc.ipynb 31
def pc_union(idx1:np.ndarray, # int array, grid index or hillbert index of the first point cloud
             idx2:np.ndarray, # int array, grid index or hillbert index of the second point cloud
             shape:tuple=None, # image shape, faster if provided for grid index input
             # the union index `idx`; 
             # index of the point in output union index that originally in the first point cloud `inv_iidx`;
             # index of the point in output union index that only exist in the second point cloud `inv_iidx2`;
             # index of the point in the second input index that are not in the first input point cloud
            )->tuple:
    '''Get the union of two point cloud dataset. For points at their intersection, prefer idx1 rather than idx2'''
    assert idx1.ndim == idx2.ndim
    xp = get_array_module(idx1)
    n1 = idx1.shape[0]; n2 = idx2.shape[0]
    if idx1.ndim == 2:
        dims = shape if shape is not None else _ras_dims(idx1,idx2)
        idx1_1d = _ravel_gix(idx1,dims=dims) # automatically the returned 1d index is in int64
        idx2_1d = _ravel_gix(idx2,dims=dims)
    else:
        idx1_1d = idx1; idx2_1d = idx2
    _check_idx_sorted(idx1_1d); _check_idx_sorted(idx2_1d)

    if xp is np:
        # on cpu use handwrite merge sort
        idx_1d, inv_iidx1, inv_iidx2, ninv_iidx2 = _pc_union_numba(idx1_1d,idx2_1d)
        if idx1.ndim == 2:
            idx = _unravel_gix(idx_1d,dims)
        else:
            idx = idx_1d
        return idx, inv_iidx1, inv_iidx2, ninv_iidx2
    else:
        # this part is modified from np.unique
        idx_1d = xp.concatenate((idx1_1d,idx2_1d),axis=-1)
    
        iidx = xp.argsort(idx_1d,kind='stable') # test shows argsort is faster than lexsort, that is why use ravel and unravel index
        idx_1d = idx_1d[iidx]
    
        inv_iidx = xp.empty_like(iidx)
        inv_iidx[iidx] = xp.arange(iidx.shape[0]) # idea taken from https://stackoverflow.com/questions/2483696/undo-or-reverse-argsort-python
    
        mask = xp.empty(idx_1d.shape, dtype=bool)
        mask[:1] = True
        mask[1:] = idx_1d[1:] != idx_1d[:-1]
        
        idx_1d = idx_1d[mask]
        
        _mask = mask[inv_iidx] # the mask in the original cat order
        mask1 = _mask[:n1]
        mask2 = _mask[n1:]
        
        imask = xp.cumsum(mask) - 1
        inv_iidx = xp.empty(mask.shape, dtype=np.int64)
        inv_iidx[iidx] = imask # inverse the mapping
        inv_iidx = inv_iidx[_mask]
        
        inv_iidx1, inv_iidx2, ninv_iidx2 = inv_iidx[:n1], inv_iidx[n1:], *xp.where(mask2)

        if idx1.ndim == 2:
            idx = _unravel_gix(idx_1d,dims)
        else:
            idx = idx_1d
    
        return idx, inv_iidx[:n1], inv_iidx[n1:], *xp.where(mask2)

# %% ../nbs/API/pc.ipynb 42
def pc_intersect(idx1:np.ndarray, # int array, grid index or hillbert index of the first point cloud
                 idx2:np.ndarray, # int array, grid index or hillbert index of the second point cloud
                 # the intersect index `idx`,
                 # index of the point in first point cloud index that also exist in the second point cloud,
                 # index of the point in second point cloud index that also exist in the first point cloud
                 shape:tuple=None, # image shape, faster if provided for grid index input
                )->tuple:
    '''Get the intersection of two point cloud dataset.'''
    # Here I do not write the core function by myself since cupy have a different implementation of intersect1d
    assert idx1.ndim == idx2.ndim
    xp = get_array_module(idx1)
    if idx1.ndim == 2:
        dims = shape if shape is not None else _ras_dims(idx1,idx2)
        idx1_1d = _ravel_gix(idx1,dims=dims) # automatically the returned 1d index is in int64
        idx2_1d = _ravel_gix(idx2,dims=dims) # automatically the returned 1d index is in int64
    else:
        idx1_1d = idx1; idx2_1d = idx2

    _check_idx_sorted(idx1_1d); _check_idx_sorted(idx2_1d)
    idx, iidx1, iidx2 = xp.intersect1d(idx1_1d,idx2_1d,assume_unique=True,return_indices=True)
    if idx1.ndim == 2:
        idx = _unravel_gix(idx,dims)
    return idx, iidx1, iidx2

# %% ../nbs/API/pc.ipynb 45
def pc_diff(idx1:np.ndarray, # int array, grid index or hillbert index of the first point cloud
            idx2:np.ndarray, # int array, grid index or hillbert index of the second point cloud
            shape:tuple=None, # image shape, faster if provided for grid index input
            # the diff index `idx`,
            # index of the point in first point cloud index that do not exist in the second point cloud,
           )->tuple:
    '''Get the point cloud in `idx1` that are not in `idx2`.'''
    assert idx1.ndim == idx2.ndim
    xp = get_array_module(idx1)
    if idx1.ndim == 2:
        dims = shape if shape is not None else _ras_dims(idx1,idx2)
        idx1_1d = _ravel_gix(idx1,dims) # automatically the returned 1d index is in int64
        idx2_1d = _ravel_gix(idx2,dims) # automatically the returned 1d index is in int64
    else:
        idx1_1d = idx1; idx2_1d = idx2

    _check_idx_sorted(idx1_1d); _check_idx_sorted(idx2_1d)
    mask = xp.in1d(idx1_1d, idx2_1d, assume_unique=True, invert=True)
    idx = idx1_1d[mask]
    
    if idx1.ndim == 2:
        idx = _unravel_gix(idx,dims)
        
    return idx, xp.where(mask)[0]

# %% ../nbs/API/pc.ipynb 49
@ngpjit
def _pc_split_by_chunk(
    gix,
    chunks,
    shape,
):
    '''utility for rearrange pc into raster chunks order'''
    nlines, width = shape
    az_chunk, r_chunk = chunks
    az_bounds = np.arange(0,nlines+az_chunk,az_chunk)
    if az_bounds[-1] > nlines: az_bounds[-1] = nlines
    r_bounds = np.arange(0,width+r_chunk,r_chunk)
    if r_bounds[-1] > width: r_bounds[-1] = width

    n_az_chunk = math.ceil(nlines/az_chunk)
    n_r_chunk = math.ceil(width/r_chunk)
    gix_az_bounds = np.searchsorted(gix[:,0],az_bounds)
    
    chunk_idx = np.empty(gix.shape[0],dtype=np.int64) # index that sort gix ordered pc into ras chunk order
    chunk_bounds = np.empty(n_az_chunk*n_r_chunk+1,dtype=np.int64) # boundarys that split ras chunk ordered pc into individual chunk
    invert_idx = np.empty(gix.shape[0],dtype=np.int64) # index that sort ras-chunk ordered pc to gix order back
    for i in prange(n_az_chunk):
        pc_start,pc_end = gix_az_bounds[i:i+2]
        gix_r_ = gix[pc_start:pc_end,1]
        r_chunk_idx = np.searchsorted(r_bounds,gix_r_,side='right')-1
        _chunk_idx = np.argsort(r_chunk_idx,kind='mergesort')

        chunk_bounds[i*n_r_chunk:(i+1)*n_r_chunk] = np.searchsorted(r_chunk_idx[_chunk_idx],np.arange(n_r_chunk))+pc_start

        chunk_idx[pc_start:pc_end] = _chunk_idx+pc_start
        
        _invert_idx = np.empty_like(r_chunk_idx,dtype=np.int64)
        _invert_idx[_chunk_idx] = np.arange(pc_start,pc_end)
        invert_idx[pc_start:pc_end] = _invert_idx
    chunk_bounds[-1] = gix.shape[0]
    return chunk_idx, chunk_bounds, invert_idx

# %% ../nbs/API/pc.ipynb 53
@ngpjit
def _gix_ras_chunk(
    sorted_gix,
    chunk_bounds,
    chunks,
    shape,
    overlap=(0,0),
):
    '''split gix into raster chunks'''
    nlines, width = shape
    az_chunk, r_chunk = chunks
    n_az_chunk = math.ceil(nlines/az_chunk)
    n_r_chunk = math.ceil(width/r_chunk)

    ras_chunk_order_gix = np.empty_like(sorted_gix)    
    for i in prange(chunk_bounds.shape[0]-1):
        start = chunk_bounds[i]
        stop = chunk_bounds[i+1]
        az_chunk_idx = i//n_r_chunk
        r_chunk_idx = i%n_r_chunk
        az_offset = -az_chunk_idx*az_chunk
        if az_chunk_idx != 0: az_offset += overlap[0]
        r_offset = -r_chunk_idx*r_chunk
        if r_chunk_idx != 0: r_offset += overlap[1]
        ras_chunk_order_gix[start:stop,0] = sorted_gix[start:stop,0]+az_offset
        ras_chunk_order_gix[start:stop,1] = sorted_gix[start:stop,1]+r_offset
    return ras_chunk_order_gix
