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

# %% auto 0
__all__ = ['Coord']

# %% ../nbs/API/coord_.ipynb 2
import numba
import numpy as np
import math
from .utils_ import ngjit,ngpjit
from numba import prange

# %% ../nbs/API/coord_.ipynb 3
# test shows this automatical parallel can accelerate 3x, 
# probablely because the data calculting is accelerated but the data writing is not
@ngpjit
def _coords2gixs(coords, x0, dx, nx, out):
    '''input is 1d'''
    for i in prange(coords.shape[0]):
        _out = round((coords[i] - x0)/dx)
        if _out < 0: _out = 0 # numba do not support raise error within prange, we can not raise error now
        if _out > nx-1: _out=nx-1
        out[i] = _out
    # np.round((coords - x0)/dx,out=out) # the test shows the automatical parallel is not fast enough
    # out[out < 0] = 0
    # out[out > nx - 1] = nx - 1

# %% ../nbs/API/coord_.ipynb 4
class Coord(object):
    '''utils for digitize raster and point cloud data, in a pyramid.
    The coord is defined as the continuous coordinates, e.g., longitude & latitude.
    The index is defined as the digitized index, (0,0,nx-1,ny-1).
    '''
    def __init__(self,x0,dx,nx,y0,dy,ny):
        self.x0 = x0
        self.dx = dx
        assert nx <2**31-1, "nx too big, may cause integer overflow."
        self.nx = nx
        self.xm = x0+(nx-1)*dx
        self.y0 = y0
        self.dy = dy
        assert ny <2**31-1, "ny too big, may cause integer overflow."
        self.ny = ny
        self.ym = y0+(ny-1)*dy
        self.maxlevel = math.floor(math.log2(min(nx,ny)))
        self.p = math.ceil(math.log2(max(nx,ny))) # order of the hillbert curve
    
    def max_gix(self,level):
        return math.ceil(self.nx/(2**level))-1, math.ceil(self.ny/(2**level))-1
    
    def coord2gix(self,x,y,level):
        # find the nearest xi and yi if x, y are not exactly on the grid
        xi = round((x-self.x0)/self.dx/2**level)
        yi = round((y-self.y0)/self.dy/2**level)
        xi_max, yi_max = self.max_gix(level)
        return sorted((0,xi,xi_max))[1], sorted((0,yi,yi_max))[1]
    
    def gix2coord(self,xi,yi,level):
        xi_max, yi_max = self.max_gix(level)
        xi, yi = sorted((0,xi,xi_max))[1], sorted((0,yi,yi_max))[1]
        return xi*2**level*self.dx+self.x0, yi*2**level*self.dy+self.y0

    def hv_bbox2gix_bbox(self,coord_bbox,level):
        # holoview bounding box is always in continous coordinates
        # index always truncated to include more data
        x0, y0, xm, ym = coord_bbox
        xi0 = math.floor((x0-self.x0)/self.dx/2**level+0.5)
        yi0 = math.floor((y0-self.y0)/self.dy/2**level+0.5)
        xim = math.ceil((xm-self.x0)/self.dx/2**level-0.5)
        yim = math.ceil((ym-self.y0)/self.dy/2**level-0.5)

        xi_max, yi_max = self.max_gix(level)
        return sorted((0,xi0,xi_max))[1], sorted((0,yi0,yi_max))[1], sorted((0,xim,xi_max))[1], sorted((0,yim,yi_max))[1]

    def gix_bbox2hv_bbox(self,gix_bbox, level):
        # the bounding box in holoviews is not at the pixel center but at the pixel corner
        xi_max, yi_max = self.max_gix(level)
        xi0, yi0, xim, yim = gix_bbox
        xi0, yi0, xim, yim = sorted((0,xi0,xi_max))[1], sorted((0,yi0,yi_max))[1], sorted((0,xim,xi_max))[1], sorted((0,yim,yi_max))[1]
        return (xi0-0.5)*2**level*self.dx+self.x0, (yi0-0.5)*2**level*self.dy+self.y0, (xim+0.5)*2**level*self.dx+self.x0, (yim+0.5)*2**level*self.dy+self.y0
    
    def coords2gixs(
        self,
        coords, # shape [N, 2], first column: y coordinates, secondary column, x coordinates
    ):
        '''inputs are 2d arrays. Only zoom level = 0 supported.'''
        out = np.empty_like(coords,dtype=np.int32)
        _coords2gixs(coords[:,1], self.x0, self.dx, self.nx, out[:,1])
        _coords2gixs(coords[:,0], self.y0, self.dy, self.ny, out[:,0])
        return out
    def rasterize_iidx(
        self,
        gix
    ):
        # the output is the inverse index (2D, same shape with ras), i.e., rasterized_ras = pc[iidx]
        # but the nan value in rasterized ras is not filled
        # it should be updated as rasterized_ras[iidx==-1] = np.nan
        iidx = np.full((self.ny,self.nx),fill_value=-1,dtype=np.int64)
        iidx[gix[:,0],gix[:,1]] = np.arange(gix.shape[0])
        return iidx
    def rasterize(self,pc, # point cloud data, can be a stack, shape (n_pc,...)
                  gix, # 2D grid index of the pc, value within (0,0,im,jm), shape (2,n_pc)
                  fill_value=np.nan, # fill value
                 )->np.ndarray: # the raster data, the index of the original point cloud list 
        # the fill_value can not be np.nan if pc is integer
        ras = np.full((self.ny,self.nx,*pc.shape[1:]),fill_value=fill_value,dtype=pc.dtype)
        ras[gix[:,0],gix[:,1]] = pc
        return ras
