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

# %% auto 0
__all__ = ['de_shp_test', 'de_select_ds_can']

# %% ../../nbs/CLI/shp.ipynb 3
from itertools import product
import math

import zarr
import cupy as cp
import numpy as np
from matplotlib import pyplot as plt
import colorcet

import dask
from dask import array as da
from dask import delayed
from dask.distributed import Client, LocalCluster
from dask_cuda import LocalCUDACluster

from ..shp import ks_test
from .utils.logging import get_logger, log_args
from .utils.dask import get_cuda_cluster
from .utils.chunk_size import get_pc_chunk_size_from_n_az_chunk
from fastcore.script import call_parse

# %% ../../nbs/CLI/shp.ipynb 4
@call_parse
@log_args
def de_shp_test(rslc:str, # input: rslc stack
                pvalue:str, # output: the p value of the test
                az_half_win:int, # azimuth half window size
                r_half_win:int, # range half window size
                method:str=None, # SHP identification method,optional. Default: ks
                az_chunk_size:int=None, # azimuth chunk size, optional. Default: the azimuth chunk size in rslc
                log:str=None, # log file, optional. Default: no log file
               ):
    '''SHP identification through hypothetic test.'''
    rslc_path = rslc
    pvalue_path = pvalue

    logger = get_logger(logfile=log,level='debug')
    if not method: method = 'ks'
    logger.info(f'hypothetic test method: {method}')
    if method != 'ks':
        logger.warning('Currently only KS test is implented. Switching to it.')
        method = 'ks'

    rslc_zarr = zarr.open(rslc_path,mode='r')
    logger.zarr_info(rslc_path,rslc_zarr)

    assert rslc_zarr.ndim == 3, " rslcs dimentation is not 3."

    if not az_chunk_size:
        az_chunk_size = rslc_zarr.chunks[0]
        logger.info('using default parallel processing azimuth chunk size from input rslc.')
    logger.info('parallel processing azimuth chunk size: '+str(az_chunk_size))

    chunks=(az_chunk_size,*rslc_zarr.shape[1:])
    
    logger.info('starting dask CUDA local cluster.')
    cluster, client = get_cuda_cluster()
    # client = Client(cluster)
    logger.info('dask local CUDA cluster started.')

    cpu_rslc = da.from_zarr(rslc_path,chunks=chunks)
    logger.darr_info('rslc',cpu_rslc)

    az_win = 2*az_half_win+1
    logger.info(f'azimuth half window size: {az_half_win}; azimuth window size: {az_win}')
    r_win = 2*r_half_win+1
    logger.info(f'range half window size: {r_half_win}; range window size: {r_win}')

    depth = {0:az_half_win, 1:r_half_win, 2:0}; boundary = {0:'none',1:'none',2:'none'}
    cpu_rslc_overlap = da.overlap.overlap(cpu_rslc,depth=depth, boundary=boundary)
    logger.info('setting shared boundaries between rlsc chunks.')
    logger.darr_info('rslc with overlap', cpu_rslc_overlap)

    rslc_overlap = cpu_rslc_overlap.map_blocks(cp.asarray)
    rmli_overlap = da.abs(rslc_overlap)**2
    logger.darr_info('rmli with overlap', rmli_overlap)

    sorted_rmli_overlap = rmli_overlap.map_blocks(cp.sort,axis=-1)

    delayed_ks_test = delayed(ks_test,pure=True,nout=2)
    rmli_delayed = sorted_rmli_overlap.to_delayed()
    p_delayed = np.empty_like(rmli_delayed,dtype=object)
    dist_delayed = np.empty_like(rmli_delayed,dtype=object)

    logger.info('applying test on sorted rmli stack.')
    with np.nditer(p_delayed,flags=['multi_index','refs_ok'], op_flags=['readwrite']) as p_it:
        for p_block in p_it:
            idx = p_it.multi_index
            dist_delayed[idx],p_delayed[idx] = delayed_ks_test(rmli_delayed[idx],az_half_win=az_half_win,r_half_win=r_half_win)

            chunk_shape = (*sorted_rmli_overlap.blocks[idx].shape[:-1],az_win,r_win)
            dtype = sorted_rmli_overlap.dtype
            # dist_delayed[idx] = da.from_delayed(dist_delayed[idx],shape=chunk_shape,meta=cp.array((),dtype=dtype))
            p_delayed[idx] = da.from_delayed(p_delayed[idx],shape=chunk_shape,meta=cp.array((),dtype=dtype))
    
    p = da.block(p_delayed.reshape(*p_delayed.shape,1).tolist())
    # dist = da.block(dist_delayed.reshape(*dist_delayed.shape,1).tolist())
    logger.info('p value generated')
    logger.darr_info('p value', p)

    depth = {0:az_half_win, 1:r_half_win, 2:0, 3:0}; boundary = {0:'none',1:'none',2:'none',3:'none'}
    # dist = da.overlap.trim_overlap(dist,depth=depth,boundary=boundary)
    p = da.overlap.trim_overlap(p,depth=depth,boundary=boundary)
    logger.info('trim shared boundaries between p value chunks')
    logger.darr_info('trimmed p value', p)

    # cpu_dist = da.map_blocks(cp.asnumpy,dist)
    cpu_p = da.map_blocks(cp.asnumpy,p)
    # cpu_p = cpu_p.rechunk((*cpu_p.chunksize[:2],1,1))
    
    # _cpu_dist = cpu_dist.to_zarr(statistic,overwrite=True,compute=False)
    _cpu_p = cpu_p.to_zarr(pvalue_path,overwrite=True,compute=False)
    logger.info('saving p value.')
    logger.info('computing graph setted. doing all the computing.')
    da.compute(_cpu_p)
    logger.info('computing finished.')
    # pdb.set_trace()
    cluster.close()
    logger.info('dask cluster closed.')

# %% ../../nbs/CLI/shp.ipynb 10
@call_parse
@log_args
def de_select_ds_can(pvalue:str, # input: pvalue of hypothetic test
                     ds_can_idx:str, # output: index array of DS candidate
                     ds_can_is_shp:str, # output: bool array indicating the SHPs of DS candidate
                     p_max:float=0.05, # threshold of p value to select SHP,optional. Default: 0.05
                     shp_num_min:int=50, # threshold of number of SHPs to select DS candidate,optional. Default: 50
                     az_chunk_size:int=None, # azimuth chunk size, optional. Default: the azimuth chunk size in pvalue
                     n_pc_chunk:int=None, # number of point target chunks, optional.
                     pc_chunk_size:int=None, # point target chunk size, optional. 
                     shp_num_fig:str=None, # path to the plot of number of SHPs, optional. Default: no plot
                     is_ds_can_fig:str=None, # path to the plot of DSs candidate distribution, optional. Default: no plot
                     log=None, # log file. Default: no log file
                     ):
    '''
    Select DS candidate based on pvalue of SHP test.
    Only one of `n_pc_chunk` and `pc_chunk_size` needs to be setted. The other one is automatically determined.
    If all of them are not setted, the `n_pc_chunk` will be setted as the number of azimuth chunks.
    '''

    logger = get_logger(logfile=log)

    p_zarr = zarr.open(pvalue,mode='r')
    logger.zarr_info(pvalue, p_zarr)

    assert p_zarr.ndim == 4, " pvalue dimentation is not 4."

    if not az_chunk_size:
        az_chunk_size = p_zarr.chunks[0]
        logger.info('using default parallel processing azimuth chunk size from input pvalue zarr.')
    logger.info('parallel processing azimuth chunk size: '+str(az_chunk_size))

    chunks=(az_chunk_size,*p_zarr.shape[1:])

    logger.info('starting dask local cluster.')
    cluster = LocalCluster()
    client = Client(cluster)
    logger.info('dask local cluster started.')

    p = da.from_zarr(pvalue,chunks=chunks)
    logger.darr_info('pvalue', p)

    ds_can_is_shp_path= ds_can_is_shp
    ds_can_idx_path = ds_can_idx

    is_shp = (p < p_max) & (p >= 0)
    logger.info('selecting SHPs based on pvalue threshold: '+str(p_max))
    logger.darr_info('is_shp', is_shp)

    shp_num = da.count_nonzero(is_shp,axis=(-2,-1))
    is_ds_can = shp_num >= shp_num_min
    logger.info('selecting DS candidates based on minimum of number of SHPs: '+str(shp_num_min))
    logger.darr_info('is_ds_can', is_ds_can)
    
    logger.info('calculate ds_can index:')
    ds_can_idx = da.nonzero(is_ds_can)
    ds_can_idx = da.stack(ds_can_idx,allow_unknown_chunksizes=True)
    logger.info('slice is_shp on ds_can index:')
    with dask.config.set(**{'array.slicing.split_large_chunks': False}):
        ds_can_is_shp = is_shp.reshape(-1,*is_shp.shape[2:])[is_ds_can.reshape(-1)]
    # ds_can_is_shp, ds_can_idx, is_ds_can, shp_num = client.persist((ds_can_is_shp,ds_can_idx,is_ds_can,shp_num))

    ds_can_is_shp.compute_chunk_sizes()
    ds_can_idx.compute_chunk_sizes()
    
    logger.darr_info('ds_can_idx',ds_can_idx)
    logger.darr_info('ds_can_is_shp',ds_can_is_shp)
    
    pc_chunk_size = get_pc_chunk_size_from_n_az_chunk('p_value','ds_can_idx', p.shape[0], az_chunk_size, ds_can_idx.shape[1], logger, n_pc_chunk=n_pc_chunk,pc_chunk_size=pc_chunk_size)

    logger.info(f'rechunk ds_can_idx and ds_can_is_shp:')
    ds_can_is_shp = ds_can_is_shp.rechunk((pc_chunk_size,*ds_can_is_shp.shape[1:]))
    ds_can_idx = ds_can_idx.rechunk((2,pc_chunk_size))
    
    logger.darr_info('ds_can_idx',ds_can_idx)
    logger.darr_info('ds_can_is_shp',ds_can_is_shp)

    _ds_can_is_shp = ds_can_is_shp.to_zarr(ds_can_is_shp_path,overwrite=True,compute=False)
    logger.info('saving ds_can_is_shp.')
    _ds_can_idx = ds_can_idx.to_zarr(ds_can_idx_path,overwrite=True,compute=False)
    logger.info('saving ds_can_idx.')

    logger.info('computing graph setted. doing all the computing.')
    shp_num_result, is_ds_can_result = da.compute(_ds_can_is_shp,_ds_can_idx,shp_num,is_ds_can)[2:]
    logger.info('computing finished.')
    cluster.close()
    logger.info('dask cluster closed.')

    if shp_num_fig:
        logger.info('plotting number of SHPs.')
        fig, ax = plt.subplots(1,1,figsize=(10,10))
        pcm = ax.imshow(shp_num_result,cmap=colorcet.cm.fire)
        ax.set(title='Number of SHPs',xlabel='Range Index',ylabel='Azimuth Index')
        fig.colorbar(pcm)
        fig.savefig(shp_num_fig)
        fig.show()
    
    if is_ds_can_fig:
        logger.info('plotting DS candidate distribution.')
        fig, ax = plt.subplots(1,1,figsize=(10,10))
        pcm = ax.imshow(is_ds_can_result,cmap=colorcet.cm.fire)
        ax.set(title='DS Candidate distribution',xlabel='Range Index',ylabel='Azimuth Index')
        fig.colorbar(pcm)
        fig.savefig(is_ds_can_fig)
        fig.show()
