# 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

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 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 stack
                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.info('rslc dateset shape: '+str(rslc_zarr.shape))
    logger.info('rslc dataset chunks: '+str(rslc_zarr.chunks))

    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.')
    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 = LocalCUDACluster()
    client = Client(cluster)
    logger.info('dask local CUDA cluster started.')

    cpu_rslc = da.from_zarr(rslc_path,chunks=chunks)
    logger.info('rslc dask array shape: ' + str(cpu_rslc.shape))
    logger.info('rslc dask array chunks: '+ str(cpu_rslc.chunks))

    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.info(f'rslc dask array with overlap shape: {cpu_rslc_overlap.shape}')
    logger.info(f'rslc dask array with overlap chunks: {cpu_rslc_overlap.chunks}')

    rslc_overlap = cpu_rslc_overlap.map_blocks(cp.asarray)
    rmli_overlap = da.abs(rslc_overlap)**2
    logger.info(f'rmli dask array with overlap shape: {rmli_overlap.shape}')
    logger.info(f'rmli dask array with overlap chunks: {rmli_overlap.chunks}')

    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.info(f'p value shape: {p.shape}')
    logger.info(f'p value chunks: {p.chunks}')

    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.info(f'trimmed p value shape: {p.shape}')
    logger.info(f'trimmed p value chunks: {p.chunks}')

    # cpu_dist = da.map_blocks(cp.asnumpy,dist)
    cpu_p = da.map_blocks(cp.asnumpy,p)
    
    # _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.')
    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
                     is_shp:str, # output: bool array indicating the SHPs of every pixel
                     is_ds_can:str, # output: bool array indicating 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
                     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
                     ):
    logger = get_logger(logfile=log)

    p_zarr = zarr.open(pvalue,mode='r')
    logger.info('pvalue dateset shape: '+str(p_zarr.shape))
    logger.info('pvalue dataset chunks: '+str(p_zarr.chunks))

    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.')
    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.info('pvalue dask array shape: ' + str(p.shape))
    logger.info('pvalue dask array chunks: '+ str(p.chunks))

    is_shp_path= is_shp
    is_ds_can_path = is_ds_can

    is_shp = (p < p_max) & (p >= 0)
    logger.info('selecting SHPs based on pvalue threshold: '+str(p_max))
    logger.info(f'is_shp shape: {is_shp.shape}')
    logger.info(f'is_shp chunks: {is_shp.chunks}')

    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.info(f'is_ds_can shape: {is_ds_can.shape}')
    logger.info(f'is_ds_can chunks: {is_ds_can.chunks}')
    
    ds_can_num = is_ds_can.map_blocks(np.count_nonzero,keepdims=True,dtype=int,chunks=(1,1))

    _is_shp = is_shp.to_zarr(is_shp_path,overwrite=True,compute=False)
    logger.info('saving is_shp.')
    _is_ds_can = is_ds_can.to_zarr(is_ds_can_path,overwrite=True,compute=False)
    logger.info('saving is_ds_can.')

    logger.info('computing graph setted. doing all the computing.')
    shp_num_result, is_ds_can_result, ds_can_num_result = da.compute(_is_shp,_is_ds_can,shp_num,is_ds_can,ds_can_num)[2:]
    logger.info('computing finished.')
    logger.info(f'number of ds can in each chunk: {tuple(ds_can_num_result.reshape(-1))}')

    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()
