"""Accurate and interative big data visualization"""

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

# %% auto 0
__all__ = ['ras_pyramid', 'ras_plot', 'pc_pyramid', 'pc_plot']

# %% ../../nbs/CLI/plot.ipynb 4
import logging
import zarr
import numpy as np
import math
from pathlib import Path
from functools import partial
from typing import Callable
import numpy as np
from numba import prange
import holoviews as hv
from holoviews import streams

import dask
from dask import array as da
from dask import delayed
from dask.distributed import Client, LocalCluster, progress
import time

import toml
from ..utils_ import ngpjit
from ..rtree import HilbertRtree
from .logging import mc_logger
from ..coord_ import Coord
from . import mk_clean_dir, dask_from_zarr, dask_to_zarr, parallel_write_zarr, parallel_read_zarr

# %% ../../nbs/CLI/plot.ipynb 6
def _ras_downsample_all_and_save(ras,zarrs,channel_idx):
    slices = [slice(None),slice(None)]
    if len(channel_idx) != 0:
        for idx in channel_idx:
            slices.append(slice(idx,idx+1))
    slices = tuple(slices)

    for level in range(len(zarrs)):
        ras_ = ras[::2**level,::2**level]
        parallel_write_zarr(ras_,zarrs[level],slices)

# %% ../../nbs/CLI/plot.ipynb 7
@mc_logger
def ras_pyramid(
    ras:str, # path to input data, 2D zarr array (one single raster) or 3D zarr array (a stack of rasters)
    out_dir:str, # output directory to store rendered data
    chunks:tuple[int,int]=(256,256), # output raster tile size
    processes=False, # use process for dask worker over thread
    n_workers=1, # number of dask worker
    threads_per_worker=2, # number of threads per dask worker
    **dask_cluster_arg, # other dask local cluster args
):
    '''render raster data to pyramid of difference zoom levels.'''
    # I forget why threads_per_worker set to 2, maybe because one for data read and write and another one for process
    logger = logging.getLogger(__name__)
    logger.info('clean out dir')
    out_dir = Path(out_dir); mk_clean_dir(out_dir)

    ras_zarr = zarr.open(ras,mode='r')
    logger.zarr_info(ras, ras_zarr)
    
    ny, nx = ras_zarr.shape[0:2]
    n_channel = ras_zarr.ndim-2
    out_chunks = chunks
    channel_chunks = ((1,)*n_channel)
    maxlevel = math.floor(math.log2(min(nx,ny))) # so at least 2 pixels
    
    logger.info(f'rendered raster pyramid with zoom level ranging from 0 (finest resolution) to {maxlevel} (coarsest resolution).')

    with LocalCluster(processes=processes,
                      n_workers=n_workers,
                      threads_per_worker=threads_per_worker,
                      **dask_cluster_arg) as cluster, Client(cluster) as client:
        logger.info('dask local cluster started.')
        logger.dask_cluster_info(cluster)
        ras_data = dask_from_zarr(ras,chunks=(ny,nx,*channel_chunks))

        downsampled_ras_zarrs = []
        for level in range(maxlevel+1):
            shape = (math.ceil(ny/(2**level)), math.ceil(nx/(2**level)))
            #downsampled_ras_store = zarr.NestedDirectoryStore(out_dir/f'{level}.zarr')
            downsampled_ras_zarr = zarr.open(
                out_dir/f'{level}.zarr',mode='w',
                shape=(*shape,*ras_zarr.shape[2:]),
                dtype=ras_data.dtype,
                chunks=(*out_chunks,*channel_chunks),)
            logger.zarr_info(out_dir/f'{level}.zarr',downsampled_ras_zarr)
            downsampled_ras_zarrs.append(downsampled_ras_zarr)

        ras_data_delayed = ras_data.to_delayed().reshape(ras_zarr.shape[2:])
        out_delayed = np.empty_like(ras_data_delayed,dtype=object)
        downsample_save_delayed = delayed(_ras_downsample_all_and_save,pure=True,nout=0)
        
        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] = downsample_save_delayed(ras_data_delayed[idx],downsampled_ras_zarrs,idx)
                out_delayed[idx] = da.from_delayed(out_delayed[idx],shape=(1,),dtype=int)
        out = da.block(out_delayed.tolist())
        logger.info('computing graph setted. doing all the computing.')
        # dask.visualize(out,filename="ras_pyramid.svg", optimize_graph=True, color='order')
        futures = client.persist(out)
        progress(futures,notebook=False)
        time.sleep(0.1)
        da.compute(futures)
        logger.info('computing finished.')
    logger.info('dask cluster closed.')

# %% ../../nbs/CLI/plot.ipynb 13
# there should be better way to achieve variable kdims, but I don't find that.
def _hv_ras_callback_0(x_range,y_range,width,height,scale,data_dir,post_proc,coord,level_increase):
    # start = time.time()
    if x_range is None:
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if y_range is None:
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res,y_res)))
    level += level_increase
    level = sorted((0, level, coord.maxlevel))[1]
    data_zarr = zarr.open(data_dir/f'{level}.zarr',mode='r')
    xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
    coord_bbox = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
    # decide_slice = time.time()
    data = post_proc(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1))
    # post_proc_data = time.time()
    # print(f"It takes {post_proc_data-decide_slice} to post_proc the data", file = sourceFile)
    ### test shows data read takes only 0.006 s, post_proc and data_range takes only 0.001s
    ### the majority of time is used by holoviews that I can not optimize.
    return hv.Image(data[::-1,:],bounds=coord_bbox)
def _hv_ras_callback_1(x_range,y_range,width,height,scale,data_dir,post_proc,coord,level_increase,i=0):
    if x_range is None:
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if y_range is None:
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res,y_res)))
    level = sorted((0, level, coord.maxlevel))[1]
    level += level_increase
    data_zarr = zarr.open(data_dir/f'{level}.zarr',mode='r')
    xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
    coord_bbox = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
    data = post_proc(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1),i)
    return hv.Image(data[::-1,:],bounds=coord_bbox)
def _hv_ras_callback_2(x_range,y_range,width,height,scale,data_dir,post_proc,coord,level_increase,i=0,j=0):
    if x_range is None:
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if y_range is None:
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res,y_res)))
    level = sorted((0, level, coord.maxlevel))[1]
    level += level_increase
    data_zarr = zarr.open(data_dir/f'{level}.zarr',mode='r')
    xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
    coord_bbox = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
    data = post_proc(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1),i,j)
    return hv.Image(data[::-1,:],bounds=coord_bbox)

# %% ../../nbs/CLI/plot.ipynb 14
def _default_ras_post_proc(data_zarr, xslice, yslice, *kdims):
    data_n_kdim = data_zarr.ndim - 2
    assert len(kdims) == data_n_kdim
    if len(kdims) == 0:
        # zarr do not support empty tuple as input
        return data_zarr[yslice,xslice]
    else:
        index_tuple = (yslice, xslice, *kdims)
        return data_zarr[index_tuple]

# %% ../../nbs/CLI/plot.ipynb 15
def _ras_phase_post_proc(data_zarr, xslice, yslice, *kdims):
    data_n_kdim = data_zarr.ndim - 2
    assert len(kdims) == 1
    i = kdims[0]
    assert data_n_kdim == 1
    assert np.iscomplexobj(data_zarr)
    return np.angle(data_zarr[yslice,xslice,i])

# %% ../../nbs/CLI/plot.ipynb 16
def _ras_inf_0_post_proc(data_zarr, xslice, yslice, *kdims):
    data_n_kdim = data_zarr.ndim - 2
    assert len(kdims) == 1
    i = kdims[0]
    if data_n_kdim == 1:
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[yslice,xslice,0]*data_zarr[yslice,xslice,i].conj())
        else:
            return data_zarr[yslice,xslice,0]-data_zarr[yslice,xslice,i]
    else:
        assert data_n_kdim == 2
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[yslice,xslice,0,i])
        else:
            return data_zarr[yslice,xslice,0,i]

def _ras_inf_seq_post_proc(data_zarr, xslice, yslice, *kdims):
    data_n_kdim = data_zarr.ndim - 2
    assert len(kdims) == 1
    i = kdims[0]
    if data_n_kdim == 1:
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[yslice,xslice,i]*data_zarr[yslice,xslice,i+1].conj())
        else:
            return data_zarr[yslice,xslice,i]-data_zarr[yslice,xslice,i+1]
    else:
        assert data_n_kdim == 2
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[yslice,xslice,i,i+1])
        else:
            return data_zarr[yslice,xslice,i,i+1]
def _ras_inf_all_post_proc(data_zarr, xslice, yslice, *kdims):
    data_n_kdim = data_zarr.ndim - 2
    assert len(kdims) == 2
    i,j = kdims
    if data_n_kdim == 1:
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[yslice,xslice,i]*data_zarr[yslice,xslice,j].conj())
        else:
            return data_zarr[yslice,xslice,i]-data_zarr[yslice,xslice,j]
    else:
        assert data_n_kdim == 2
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[yslice,xslice,i,j])
        else:
            return data_zarr[yslice,xslice,i,j]

# %% ../../nbs/CLI/plot.ipynb 17
def ras_plot(
    pyramid_dir:str, # directory to the rendered ras pyramid
    post_proc:Callable=None, # function for the post processing, can be None, 'intf_0', 'intf_seq', 'intf_all' or user-defined function
    n_kdim:int=None, # number of key dimensions, can only be 0 or 1 or 2, ndim of raster dataset -2 by default
    bounds:tuple=None, # bounding box (x0, y0, x_max, y_max)
    level_increase=0, # amount of zoom level increase for more clear point show and faster responds time
):
    '''plot rendered stack of ras tiles.'''
    pyramid_dir = Path(pyramid_dir)
    data_zarr = zarr.open(pyramid_dir/'0.zarr',mode='r')
    ny, nx = data_zarr.shape[:2]
    if post_proc is None: 
        post_proc = _default_ras_post_proc
    elif post_proc == 'intf_0':
        post_proc = _ras_inf_0_post_proc
        n_kdim = 1
    elif post_proc == 'intf_seq':
        post_proc = _ras_inf_seq_post_proc
        n_kdim = 1
    elif post_proc == 'intf_all':
        post_proc = _ras_inf_all_post_proc
        n_kdim = 2
    elif post_proc == 'phase':
        post_proc = _ras_phase_post_proc
        n_kdim = 1

    if n_kdim is None: n_kdim = data_zarr.ndim -2 
    assert n_kdim <= 2, 'n_kdim can only be 0 or 1 or2.'
    kdims = ['i','j'][:n_kdim]

    if len(kdims) == 0:
        hv_ras_callback = _hv_ras_callback_0
    elif len(kdims) == 1:
        hv_ras_callback = _hv_ras_callback_1
    elif len(kdims) == 2:
        hv_ras_callback = _hv_ras_callback_2

    if bounds is None:
        x0 = 0; dx = 1; y0 = 0; dy = 1
    else:
        x0, y0, xm, ym = bounds
        dx = (xm-x0)/(nx-1); dy = (ym-y0)/(ny-1)
    coord = Coord(x0,dx,nx,y0,dy,ny)

    rangexy = streams.RangeXY()
    plotsize = streams.PlotSize()
    images = hv.DynamicMap(partial(hv_ras_callback,data_dir=pyramid_dir,
                                   post_proc=post_proc,coord=coord,level_increase=level_increase),streams=[rangexy,plotsize],kdims=kdims)
    return images

# %% ../../nbs/CLI/plot.ipynb 51
@ngpjit
def _next_level_idx_from_raster_of_integer(pc_idx, nan_value):
    '''return the raster indices to the next level of raster'''
    assert pc_idx.ndim == 2
    ny, nx = pc_idx.shape
    next_ny, next_nx = math.ceil(ny/2), math.ceil(nx/2)
    xi = np.empty((next_ny,next_nx), dtype=np.int32)
    yi = np.empty((next_ny,next_nx), dtype=np.int32)

    for i in range(next_ny):
        for j in prange(next_nx):
            # Select a 2x2 box from the original array
            box = pc_idx[i*2:min(i*2+2, ny), j*2:min(j*2+2, nx)]
            idx_ = np.argwhere(box != nan_value)
            if len(idx_) == 0:
                yi[i,j]= i*2
                xi[i,j] = j*2
            else:
                yi[i,j] = idx_[0,0] + i*2
                xi[i,j] = idx_[0,1] + j*2
    return yi, xi

# %% ../../nbs/CLI/plot.ipynb 53
def _pc_downsample_all_and_save(pc,coord,gix,yis,xis,pc_zarr,ras_zarrs,channel_idx):
    pc_slices = [slice(None),]
    ras_slices = [slice(None),slice(None)]
    if len(channel_idx) != 0:
        for idx in channel_idx:
            pc_slices.append(slice(idx,idx+1))
            ras_slices.append(slice(idx,idx+1))
    pc_slices = tuple(pc_slices)
    ras_slices = tuple(ras_slices)
    parallel_write_zarr(pc,pc_zarr,pc_slices)

    ras = coord.rasterize(pc,gix)
    parallel_write_zarr(ras,ras_zarrs[0],ras_slices)

    for level in range(1,len(ras_zarrs)):
        ras = ras[yis[level-1],xis[level-1]]
        parallel_write_zarr(ras,ras_zarrs[level],ras_slices)

# %% ../../nbs/CLI/plot.ipynb 54
@mc_logger
def pc_pyramid(
    pc:str, # path to point cloud data, 1D array (one single pc image) or 2D zarr array (a stack of pc images)
    out_dir:str, # output directory to store rendered data
    x:str=None, # path to x coordinate, e.g., longitude or web mercator x
    y:str=None, # path to y coordinate, e.g., latitude or web mercator y
    yx:str=None, # path to x and y coordinates. this coordinates should have shape [n_points,2]. e.g., gix
    ras_resolution:float=20, # minimum resolution of rendered raster data,
    ras_chunks:tuple[int,int]=(256,256), # output raster tile size
    pc_chunks:int=65536, # output pc tile size
    processes=False, # use process for dask worker over thread
    n_workers=1, # number of dask worker
    threads_per_worker=2, # number of threads per dask worker
    **dask_cluster_arg, # other dask local cluster args
):
    '''render point cloud data to pyramid of difference zoom levels.'''
    logger = logging.getLogger(__name__)
    logger.info('clean out dir')
    out_dir = Path(out_dir); mk_clean_dir(out_dir)

    pc_zarr = zarr.open(pc,mode='r')
    logger.zarr_info(pc, pc_zarr)
    
    n_pc = pc_zarr.shape[0]
    channel_chunks = (1,)*(pc_zarr.ndim-1)
    logger.info(f'rendering point cloud data coordinates:')
    if x is None and y is None:
        yx_zarr = zarr.open(yx,mode='r')
        assert yx_zarr.shape[1] == 2
        yx = parallel_read_zarr(yx_zarr,(slice(None),slice(0,2)))
    else:
        y_zarr = zarr.open(y,mode='r')
        yx = np.empty((y_zarr.shape[0],2),dtype=y_zarr.dtype)
        yx[:,0] = parallel_read_zarr(zarr.open(y,mode='r'),(slice(None),))
        yx[:,1] = parallel_read_zarr(zarr.open(x,mode='r'),(slice(None),))
    x, y = yx[:,1], yx[:,0]

    x0, xm, y0, ym = x.min(), x.max(), y.min(), y.max()
    nx, ny = math.ceil((xm-x0)/ras_resolution), math.ceil((ym-y0)/ras_resolution)
    coord = Coord(x0, ras_resolution, nx, y0, ras_resolution, ny)
    bounds = {'bounds':[x0, y0, coord.xm, coord.ym]}
    logger.info(f"rasterizing point cloud data to grid with bounds: {bounds['bounds']}.")
    with open(out_dir/'bounds.toml',mode='w') as f:
        toml.dump(bounds, f, encoder=toml.TomlNumpyEncoder())

    gix = coord.coords2gixs(yx)
    maxlevel = coord.maxlevel

    out_x_zarr = zarr.open(out_dir/f'x.zarr',mode='w',shape=x.shape,dtype=x.dtype,chunks=(pc_chunks,))
    out_y_zarr = zarr.open(out_dir/f'y.zarr',mode='w',shape=y.shape,dtype=y.dtype,chunks=(pc_chunks,))
    logger.zarr_info(out_dir/f'x.zarr',out_x_zarr)
    logger.zarr_info(out_dir/f'y.zarr',out_y_zarr)
    parallel_write_zarr(x, out_x_zarr,(slice(None),))
    parallel_write_zarr(y, out_y_zarr,(slice(None),))
    del x, y, yx
    logger.info('pc data coordinates rendering ends.')

    yis = []; xis = []
    for level in range(maxlevel+1):
        if level == 0:
            current_idx = coord.rasterize_iidx(gix)
        else:
            yi, xi = _next_level_idx_from_raster_of_integer(last_idx,-1)
            yis.append(yi); xis.append(xi)
            current_idx = last_idx[yi,xi]
        out_idx_zarr = zarr.open(out_dir/f'idx_{level}.zarr',shape=current_idx.shape,dtype=current_idx.dtype,chunks=ras_chunks)
        logger.zarr_info(out_dir/f'idx_{level}.zarr',out_idx_zarr)
        parallel_write_zarr(current_idx,out_idx_zarr,(slice(None),slice(None)))
        last_idx = current_idx
    logger.info('rasterized idx rendering ends')

    with LocalCluster(processes=processes,
                      n_workers=n_workers,
                      threads_per_worker=threads_per_worker,
                      **dask_cluster_arg) as cluster, Client(cluster) as client:
        logger.info('dask local cluster started to render pc data.')
        logger.dask_cluster_info(cluster)

        out_pc_zarr = zarr.open(out_dir/f'pc.zarr',mode='w',shape=pc_zarr.shape, dtype=pc_zarr.dtype, chunks=(pc_chunks,*channel_chunks))
        logger.zarr_info(out_dir/f'pc.zarr', out_pc_zarr)
    
        downsampled_ras_zarrs = []
        for level in range(maxlevel+1):
            shape = (math.ceil(ny/(2**level)), math.ceil(nx/(2**level)))
            #downsampled_ras_store = zarr.NestedDirectoryStore(out_dir/f'{level}.zarr')
            downsampled_ras_zarr = zarr.open(
                out_dir/f'{level}.zarr',mode='w',
                shape=(*shape,*pc_zarr.shape[1:]),
                dtype=pc_zarr.dtype,
                chunks=(*ras_chunks,*channel_chunks),)
            logger.zarr_info(out_dir/f'{level}.zarr',downsampled_ras_zarr)
            downsampled_ras_zarrs.append(downsampled_ras_zarr)

        pc_darr = dask_from_zarr(pc,chunks=(n_pc,*channel_chunks))
        pc_delayed = pc_darr.to_delayed().reshape(pc_zarr.shape[1:])
        out_delayed = np.empty_like(pc_delayed,dtype=object)
        downsample_save_delayed = delayed(_pc_downsample_all_and_save,pure=True,nout=0)
        
        with np.nditer(out_delayed,flags=['multi_index','refs_ok'], op_flags=['readwrite']) as arr_it:
            for arr_block in arr_it:
                channel_idx = arr_it.multi_index
                out_delayed[channel_idx] = downsample_save_delayed(pc_delayed[channel_idx],coord,gix, yis, xis, out_pc_zarr, downsampled_ras_zarrs,channel_idx)
                out_delayed[channel_idx] = da.from_delayed(out_delayed[channel_idx],shape=(1,),dtype=int)
        out = da.block(out_delayed.tolist())

        logger.info('computing graph setted. doing all the computing.')
        futures = client.persist(out)
        progress(futures,notebook=False)
        time.sleep(0.1)
        da.compute(futures)
        logger.info('computing finished.')
    logger.info('dask cluster closed.')

# %% ../../nbs/CLI/plot.ipynb 59
def _is_nan_range(x_range):
    if x_range is None:
        return True
    if np.isnan(x_range[0]):
        return True
    if abs(x_range[1]-x_range[0]) == 0:
        return True
    return False

# %% ../../nbs/CLI/plot.ipynb 60
def _hv_pc_Image_callback_0(x_range,y_range,width,height,scale,data_dir,post_proc_ras,coord,level_increase):
    if _is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if _is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level += level_increase
    level = sorted((-1, level, coord.maxlevel))[1]
    # level = -1
    images = []
    if level > -1:
        data_zarr = zarr.open(data_dir/f'{level}.zarr',mode='r')
        idx_zarr = zarr.open(data_dir/f'idx_{level}.zarr',mode='r')
        xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
        x0, y0, xm, ym = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
        data = post_proc_ras(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1))
        idx = idx_zarr[yi0:yim+1,xi0:xim+1]
        return hv.Image((np.linspace(x0,xm,data.shape[1]), np.linspace(y0,ym,data.shape[0]),data,idx),vdims=['z','idx'])
    else:
        return hv.Image([],vdims=['z','idx'])
def _hv_pc_Image_callback_1(x_range,y_range,width,height,scale,data_dir,post_proc_ras,coord,level_increase,i):
    if _is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if _is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level += level_increase
    level = sorted((-1, level, coord.maxlevel))[1]
    images = []
    if level > -1:
        data_zarr = zarr.open(data_dir/f'{level}.zarr',mode='r')
        idx_zarr = zarr.open(data_dir/f'idx_{level}.zarr',mode='r')
        xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
        x0, y0, xm, ym = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
        data = post_proc_ras(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1),i)
        idx = idx_zarr[yi0:yim+1,xi0:xim+1]
        return hv.Image((np.linspace(x0,xm,data.shape[1]), np.linspace(y0,ym,data.shape[0]),data,idx),vdims=['z','idx'])
    else:
        return hv.Image([],vdims=['z','idx'])
def _hv_pc_Image_callback_2(x_range,y_range,width,height,scale,data_dir,post_proc_ras,coord,level_increase,i,j):
    if _is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if _is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level += level_increase
    level = sorted((-1, level, coord.maxlevel))[1]
    images = []
    if level > -1:
        data_zarr = zarr.open(data_dir/f'{level}.zarr',mode='r')
        idx_zarr = zarr.open(data_dir/f'idx_{level}.zarr',mode='r')
        xi0, yi0, xim, yim = coord.hv_bbox2gix_bbox((x0,y0,xm,ym),level)
        x0, y0, xm, ym = coord.gix_bbox2hv_bbox((xi0, yi0, xim, yim),level)
        data = post_proc_ras(data_zarr,slice(xi0,xim+1),slice(yi0,yim+1),i,j)
        idx = idx_zarr[yi0:yim+1,xi0:xim+1]
        return hv.Image((np.linspace(x0,xm,data.shape[1]), np.linspace(y0,ym,data.shape[0]),data,idx),vdims=['z','idx'])
    else:
        return hv.Image([],vdims=['z','idx'])

# %% ../../nbs/CLI/plot.ipynb 61
def _hv_pc_Points_callback_0(x_range,y_range,width,height,scale,data_dir,post_proc_pc,coord,rtree,level_increase):
    if _is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if _is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level += level_increase
    level = sorted((-1, level, coord.maxlevel))[1]
    images = []
    if level > -1:
        return hv.Points([],vdims=['z','idx'])
    else:
        data_zarr, x_zarr, y_zarr = (zarr.open(data_dir/file,mode='r') for file in ('pc.zarr', 'x.zarr', 'y.zarr'))
        idx = rtree.bbox_query((x0, y0, xm, ym), x_zarr, y_zarr)
        x, y = (zarr_[idx] for zarr_ in (x_zarr, y_zarr))
        data = post_proc_pc(data_zarr,idx)
        return hv.Points((x,y,data,idx),vdims=['z','idx'])
def _hv_pc_Points_callback_1(x_range,y_range,width,height,scale,data_dir,post_proc_pc,coord,rtree,level_increase,i=0):
    if _is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if _is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level += level_increase
    level = sorted((-1, level, coord.maxlevel))[1]
    # level = -1
    images = []
    if level > -1:
        return hv.Points([],vdims=['z','idx'])
    else:
        data_zarr, x_zarr, y_zarr = (zarr.open(data_dir/file,mode='r') for file in ('pc.zarr', 'x.zarr', 'y.zarr'))
        idx = rtree.bbox_query((x0, y0, xm, ym), x_zarr, y_zarr)
        x, y = (zarr_[idx] for zarr_ in (x_zarr, y_zarr))
        data = post_proc_pc(data_zarr,idx,i)
        return hv.Points((x,y,data,idx),vdims=['z','idx'])
def _hv_pc_Points_callback_2(x_range,y_range,width,height,scale,data_dir,post_proc_pc,coord,rtree,level_increase,i=0,j=0):
    if _is_nan_range(x_range):
        x0 = coord.x0; xm = coord.xm
    else:
        x0, xm = x_range
    if _is_nan_range(y_range):
        y0 = coord.y0; ym = coord.ym
    else:
        y0, ym = y_range
    if height is None: height = hv.plotting.bokeh.ElementPlot.height
    if width is None: width = hv.plotting.bokeh.ElementPlot.width

    x_res = (xm-x0)/width; y_res = (ym-y0)/height
    level = math.floor(math.log2(min(x_res/coord.dx,y_res/coord.dy)))
    level += level_increase
    level = sorted((-1, level, coord.maxlevel))[1]
    # level = -1
    images = []
    if level > -1:
        return hv.Points([],vdims=['z','idx'])
    else:
        data_zarr, x_zarr, y_zarr = (zarr.open(data_dir/file,mode='r') for file in ('pc.zarr', 'x.zarr', 'y.zarr'))
        idx = rtree.bbox_query((x0, y0, xm, ym), x_zarr, y_zarr)
        x, y = (zarr_[idx] for zarr_ in (x_zarr, y_zarr))
        data = post_proc_pc(data_zarr,idx,i,j)
        return hv.Points((x,y,data,idx),vdims=['z','idx'])

# %% ../../nbs/CLI/plot.ipynb 62
def _default_pc_post_proc(data_zarr, idx_array, *kdims):
    data_n_kdim = data_zarr.ndim - 1
    assert len(kdims) == data_n_kdim
    if len(kdims) == 0:
        return data_zarr[idx_array]
    else:
        index_tuple = (idx_array, *kdims)
        return data_zarr[index_tuple]

# %% ../../nbs/CLI/plot.ipynb 63
def _pc_phase_post_proc(data_zarr, idx_array, *kdims):
    data_n_kdim = data_zarr.ndim - 1
    assert len(kdims) == 1
    assert data_n_kdim == 1
    i = kdims[0]
    assert np.iscomplexobj(data_zarr)
    return np.angle(data_zarr[idx_array,i])

# %% ../../nbs/CLI/plot.ipynb 64
def _pc_inf_0_post_proc(data_zarr, idx_array, *kdims):
    data_n_kdim = data_zarr.ndim - 1
    assert len(kdims) == 1
    i = kdims[0]
    if data_n_kdim == 1:
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[idx_array,0]*data_zarr[idx_array,i].conj())
        else:
            return data_zarr[idx_array,0]-data_zarr[idx_array,i]
    else:
        assert data_n_kdim == 2
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[idx_array,0,i])
        else:
            return data_zarr[idx_array,0,i]

def _pc_inf_seq_post_proc(data_zarr, idx_array, *kdims):
    data_n_kdim = data_zarr.ndim - 1
    assert len(kdims) == 1
    i = kdims[0]
    if data_n_kdim == 1:
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[idx_array,i]*data_zarr[idx_array,i+1].conj())
        else:
            return data_zarr[idx_array,i]-data_zarr[idx_array,i+1]
    else:
        assert data_n_kdim == 2
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[idx_array,i,i+1])
        else:
            return data_zarr[idx_array,i,i+1]

def _pc_inf_all_post_proc(data_zarr, idx_array, *kdims):
    data_n_kdim = data_zarr.ndim - 1
    assert len(kdims) == 2
    i,j = kdims
    if data_n_kdim == 1:
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[idx_array,i]*data_zarr[idx_array,j].conj())
        else:
            return data_zarr[idx_array,i]-data_zarr[idx_array,j]
    else:
        assert data_n_kdim == 2
        if np.iscomplexobj(data_zarr):
            return np.angle(data_zarr[idx_array,i,j])
        else:
            return data_zarr[idx_array,i,j]

# %% ../../nbs/CLI/plot.ipynb 65
def pc_plot(
    pyramid_dir:str, # directory to the rendered point cloud pyramid
    post_proc_ras:Callable=None, # function for the post processing
    post_proc_pc:Callable=None, # function for the post processing
    n_kdim:int=None, # number of key dimensions, can only be 0 or 1 or 2, ndim of point cloud dataset -1 by default
    rtree=None, # rtree, if not provide, will be automatically generated but may slow the program
    level_increase=0, # amount of zoom level increase for more clear point show and faster responds time
):
    '''plot rendered point cloud pyramid.'''    
    if post_proc_ras is None: post_proc_ras = _default_ras_post_proc
    if post_proc_pc is None: post_proc_pc = _default_pc_post_proc

    pyramid_dir = Path(pyramid_dir)
    data_zarr = zarr.open(pyramid_dir/'0.zarr',mode='r')
    ny, nx = data_zarr.shape[:2]

    if post_proc_ras is None:
        post_proc_ras = _default_ras_post_proc
        post_proc_pc = _default_pc_post_proc
    elif post_proc_ras == 'intf_0':
        post_proc_ras = _ras_inf_0_post_proc
        post_proc_pc = _pc_inf_0_post_proc
        n_kdim = 1
    elif post_proc_ras == 'intf_seq':
        post_proc_ras = _ras_inf_seq_post_proc
        post_proc_pc = _pc_inf_seq_post_proc
        n_kdim = 1
    elif post_proc_ras == 'intf_all':
        post_proc_ras = _ras_inf_all_post_proc
        post_proc_pc = _pc_inf_all_post_proc
        n_kdim = 2
    elif post_proc_pc == 'phase':
        post_proc_ras = _ras_phase_post_proc
        post_proc_pc = _pc_phase_post_proc
        n_kdim = 1

    if n_kdim is None: n_kdim = data_zarr.ndim -2
    assert n_kdim <= 2, 'n_kdim can only be 0 or 1 or2.'
    kdims = ['i','j'][:n_kdim]
    
    with open(pyramid_dir/'bounds.toml','r') as f:
        x0, y0, xm, ym = toml.load(f)['bounds']

    dx = (xm-x0)/(nx-1); dy = (ym-y0)/(ny-1)
    coord = Coord(x0,dx,nx,y0,dy,ny)
    
    if rtree is None:
        x = zarr.open(pyramid_dir/'x.zarr',mode='r')[:]
        y = zarr.open(pyramid_dir/'y.zarr',mode='r')[:]
        rtree = HilbertRtree.build(x,y,page_size=512)

    if len(kdims) == 0:
        hv_pc_Image_callback = _hv_pc_Image_callback_0
        hv_pc_Points_callback = _hv_pc_Points_callback_0
    elif len(kdims) == 1:
        hv_pc_Image_callback = _hv_pc_Image_callback_1
        hv_pc_Points_callback = _hv_pc_Points_callback_1
    elif len(kdims) == 2:
        hv_pc_Image_callback = _hv_pc_Image_callback_2
        hv_pc_Points_callback = _hv_pc_Points_callback_2

    rangexy = streams.RangeXY()
    plotsize = streams.PlotSize()
    images = hv.DynamicMap(partial(hv_pc_Image_callback,data_dir=pyramid_dir,
                                   post_proc_ras=post_proc_ras,coord=coord,level_increase=level_increase),
                           streams=[rangexy,plotsize],kdims=kdims)
    points = hv.DynamicMap(partial(hv_pc_Points_callback,data_dir=pyramid_dir,
                                   post_proc_pc=post_proc_pc,coord=coord,rtree=rtree,level_increase=level_increase),
                           streams=[rangexy,plotsize],kdims=kdims)
    return images*points
