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

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

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

# %% ../nbs/API/plot.ipynb 3
import numpy as np
from typing import Callable
from functools import partial
import math

import holoviews as hv
from holoviews import streams
from .coord_ import Coord
from .rtree import HilbertRtree

# %% ../nbs/API/plot.ipynb 7
# 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,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_ = data[::2**level,::2**level]
    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_,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,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_ = data[::2**level,::2**level]
    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_,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,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_ = data[::2**level,::2**level]
    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_,slice(xi0,xim+1),slice(yi0,yim+1),i,j)
    return hv.Image(data_[::-1,:],bounds=coord_bbox)

# %% ../nbs/API/plot.ipynb 8
def _default_ras_post_proc(data, xslice, yslice, *kdims):
    data_n_kdim = data.ndim - 2
    assert len(kdims) == data_n_kdim
    if len(kdims) == 0:
        return data[yslice,xslice]
    else:
        index_tuple = (yslice, xslice, *kdims)
        return data[index_tuple]

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

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

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

# %% ../nbs/API/plot.ipynb 11
def ras_plot(
    data: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.'''
    ny, nx = data.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.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=data,
                                   post_proc=post_proc,coord=coord,level_increase=level_increase),streams=[rangexy,plotsize],kdims=kdims)
    return images

# %% ../nbs/API/plot.ipynb 37
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/API/plot.ipynb 38
def _hv_pc_Image_callback_0(
    x_range,y_range,width,height,scale,
    data,y,x,rtree,
    post_proc_pc,
    coord,level_increase,
    ras_resolution,
):
    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*2**level_increase; y_res = (ym-y0)/height*2**level_increase
    if min(x_res, y_res) > ras_resolution:
        # slice new data
        idx = rtree.bbox_query((x0, y0, xm, ym), x, y)
        if len(idx) == 0:
            return hv.Image([], vdims=['z', 'idx'])
        data_ = post_proc_pc(data,idx); y_ = y[idx]; x_ = x[idx]

        # rasterize on sliced small dataset
        nx = math.ceil((xm - x0) / x_res)
        ny = math.ceil((ym - y0) / y_res)
        coord_ = Coord(x0, x_res, nx, y0, y_res, ny)

        yxi_ = coord_.coords2gixs(np.stack((y_, x_),axis=-1))
        ras_ = coord_.rasterize(data_, yxi_, fill_value=np.nan)
        idx_ = coord_.rasterize(idx, yxi_, fill_value=-1)

        x_vals = x0 + x_res * np.arange(ras_.shape[1])
        y_vals = y0 + y_res * np.arange(ras_.shape[0])

        return hv.Image((x_vals, y_vals, ras_, idx_), vdims=['z', 'idx'])

    else:
        return hv.Image([],vdims=['z','idx'])

def _hv_pc_Points_callback_0(
    x_range,y_range,width,height,scale,
    data,y,x,rtree,
    post_proc_pc,
    coord,level_increase,
    ras_resolution,
):
    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*2**level_increase; y_res = (ym-y0)/height*2**level_increase
    if min(x_res, y_res) > ras_resolution:
        return hv.Points([],vdims=['z','idx'])
    else:
        # slice new data
        idx = rtree.bbox_query((x0, y0, xm, ym), x, y)
        data_ = post_proc_pc(data,idx); y_ = y[idx]; x_ = x[idx]
        return hv.Points((x_,y_,data_,idx),vdims=['z','idx'])

# %% ../nbs/API/plot.ipynb 39
def _hv_pc_Image_callback_1(
    x_range,y_range,width,height,scale,
    data,y,x,rtree,
    post_proc_pc,coord,level_increase,ras_resolution,
    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*2**level_increase; y_res = (ym-y0)/height*2**level_increase
    if min(x_res, y_res) > ras_resolution:
        # slice new data
        idx = rtree.bbox_query((x0, y0, xm, ym), x, y)
        if len(idx) == 0:
            return hv.Image([], vdims=['z', 'idx'])
        data_ = post_proc_pc(data,idx,i); y_ = y[idx]; x_ = x[idx]

        # rasterize on sliced small dataset
        nx = math.ceil((xm - x0) / x_res)
        ny = math.ceil((ym - y0) / y_res)
        coord_ = Coord(x0, x_res, nx, y0, y_res, ny)

        yxi_ = coord_.coords2gixs(np.stack((y_, x_),axis=-1))
        ras_ = coord_.rasterize(data_, yxi_, fill_value=np.nan)
        idx_ = coord_.rasterize(idx, yxi_, fill_value=-1)

        x_vals = x0 + x_res * np.arange(ras_.shape[1])
        y_vals = y0 + y_res * np.arange(ras_.shape[0])

        return hv.Image((x_vals, y_vals, ras_, idx_), vdims=['z', 'idx'])

    else:
        return hv.Image([],vdims=['z','idx'])

def _hv_pc_Points_callback_1(
    x_range,y_range,width,height,scale,
    data,y,x,rtree,
    post_proc_pc,coord,level_increase,ras_resolution,
    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*2**level_increase; y_res = (ym-y0)/height*2**level_increase
    if min(x_res, y_res) > ras_resolution:
        return hv.Points([],vdims=['z','idx'])
    else:
        # slice new data
        idx = rtree.bbox_query((x0, y0, xm, ym), x, y)
        data_ = post_proc_pc(data,idx,i); y_ = y[idx]; x_ = x[idx]
        return hv.Points((x_,y_,data_,idx),vdims=['z','idx'])

# %% ../nbs/API/plot.ipynb 40
def _hv_pc_Image_callback_2(
    x_range,y_range,width,height,scale,
    data,y,x,rtree,
    post_proc_pc,coord,level_increase,ras_resolution,
    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*2**level_increase; y_res = (ym-y0)/height*2**level_increase
    if min(x_res, y_res) > ras_resolution:
        # slice new data
        idx = rtree.bbox_query((x0, y0, xm, ym), x, y)
        if len(idx) == 0:
            return hv.Image([], vdims=['z', 'idx'])
        data_ = post_proc_pc(data,idx,i,j); y_ = y[idx]; x_ = x[idx]

        # rasterize on sliced small dataset
        nx = math.ceil((xm - x0) / x_res)
        ny = math.ceil((ym - y0) / y_res)
        coord_ = Coord(x0, x_res, nx, y0, y_res, ny)

        yxi_ = coord_.coords2gixs(np.stack((y_, x_),axis=-1))
        ras_ = coord_.rasterize(data_, yxi_, fill_value=np.nan)
        idx_ = coord_.rasterize(idx, yxi_, fill_value=-1)

        x_vals = x0 + x_res * np.arange(ras_.shape[1])
        y_vals = y0 + y_res * np.arange(ras_.shape[0])

        return hv.Image((x_vals, y_vals, ras_, idx_), vdims=['z', 'idx'])

    else:
        return hv.Image([],vdims=['z','idx'])

def _hv_pc_Points_callback_2(
    x_range,y_range,width,height,scale,
    data,y,x,rtree,
    post_proc_pc,coord,level_increase,ras_resolution,
    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*2**level_increase; y_res = (ym-y0)/height*2**level_increase
    if min(x_res, y_res) > ras_resolution:
        return hv.Points([],vdims=['z','idx'])
    else:
        # slice new data
        idx = rtree.bbox_query((x0, y0, xm, ym), x, y)
        data_ = post_proc_pc(data,idx,i,j); y_ = y[idx]; x_ = x[idx]
        return hv.Points((x_,y_,data_,idx),vdims=['z','idx'])

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

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

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

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

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

# %% ../nbs/API/plot.ipynb 44
def pc_plot(
    data:np.ndarray, # pc dataset
    y:np.ndarray, # y coordinate
    x:np.ndarray, # x coordinate
    ras_resolution:float=20, # minimum resolution of rendered raster data in the pyramid 
    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 on the fly'''    
    if post_proc_pc is None: post_proc_pc = _default_pc_post_proc

    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)
    
    if post_proc_pc is None:
        post_proc_pc = _default_pc_post_proc
    elif post_proc_pc == 'intf_0':
        post_proc_pc = _pc_inf_0_post_proc
        n_kdim = 1
    elif post_proc_pc == 'intf_seq':
        post_proc_pc = _pc_inf_seq_post_proc
        n_kdim = 1
    elif post_proc_pc == 'intf_all':
        post_proc_pc = _pc_inf_all_post_proc
        n_kdim = 2
    elif post_proc_pc == 'phase':
        post_proc_pc = _pc_phase_post_proc
        n_kdim = 1

    if n_kdim is None: n_kdim = data.ndim -1
    assert n_kdim <= 2, 'n_kdim can only be 0 or 1 or2.'
    kdims = ['i','j'][:n_kdim]
    
    if rtree is None:
        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=data,y=y,x=x,rtree=rtree,
                                   post_proc_pc=post_proc_pc,
                                   coord=coord,
                                   level_increase=level_increase,
                                   ras_resolution=ras_resolution,
                                  ),
                           streams=[rangexy,plotsize],kdims=kdims)

    points = hv.DynamicMap(partial(hv_pc_Points_callback,
                                   data=data,y=y,x=x,rtree=rtree,
                                   post_proc_pc=post_proc_pc,
                                   coord=coord,
                                   level_increase=level_increase,
                                   ras_resolution=ras_resolution,
                                  ),
                           streams=[rangexy,plotsize],kdims=kdims)
    return images*points

# %% ../nbs/API/plot.ipynb 70
def bg_alpha(pwr):
    _pwr = np.power(pwr,0.35)
    cv = _pwr.mean()*2.5
    v = (_pwr.clip(0., cv))/cv
    return v
