"""Deep learning based operators"""

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

# %% auto 0
__all__ = ['n2f', 'n2ft']

# %% ../../nbs/CLI/dl.ipynb 5
import logging
import zarr
import numpy as np
from numba import prange
from pathlib import Path
import math
import cmath
import importlib
import onnxruntime
from ..utils_ import ngjit, ngpjit

import dask
from dask import array as da
from dask import delayed
from dask.distributed import Client, LocalCluster, progress
from ..utils_ import is_cuda_available
if is_cuda_available():
    import cupy as cp
    from dask_cuda import LocalCUDACluster
    from rmm.allocators.cupy import rmm_cupy_allocator
import moraine as mr
import moraine.cli as mc
from ..utils_ import get_array_module
from ..chunk_ import chunkwise_slicing_mapping, chunkwise_knn_mapping
from ..co import intf as intf_func
from .dask_ import parallel_read_zarr
from ..dl import _ort_session, _infer_n2ft
from .logging import mc_logger
from . import mk_clean_dir, dask_from_zarr, dask_from_zarr_overlap, dask_to_zarr

# %% ../../nbs/CLI/dl.ipynb 6
@ngpjit
def _cli_pre_infer_n2f_numba(
    ref, # reference rslc
    sec, # secoundary rslc
):
    nlines, width = ref.shape
    out = np.empty((1,2,nlines,width),dtype=np.float32)
    mask = np.empty((nlines,width),dtype=np.bool_)
    for i in prange(nlines):
        for j in prange(width):
            intf_i_j = ref[i,j]*sec[i,j].conjugate()
            if math.isnan(intf_i_j.real):
                mask[i,j] = True
                random_phase = np.random.uniform(-math.pi,math.pi)
                out[0,0,i,j] = math.cos(random_phase)
                out[0,1,i,j] = math.sin(random_phase)
            else:
                mask[i,j] = False
                amp = abs(intf_i_j)
                out[0,0,i,j] = intf_i_j.real/amp
                out[0,1,i,j] = intf_i_j.imag/amp
    return out, mask

# %% ../../nbs/CLI/dl.ipynb 7
if is_cuda_available():
    _cli_pre_infer_n2f_kernel = cp.ElementwiseKernel(
            'raw T ref_, raw T sec_, int32 nlines, int32 width',
            'raw float32 out, raw bool mask',
            '''
            int npixels = nlines*width;
            if (i >= npixels) return;

            T intf_i = ref_[i]*conj(sec_[i]);
            if (isnan(intf_i.real())){
                mask[i] = true;
            }
            else{
                mask[i] = false;
                float amp = abs(intf_i);
                out[i] = intf_i.real()/amp;
                out[npixels+i] = intf_i.imag()/amp;
            }
            ''',
            #preamble = '#include "curand.h"',
            # I do not find an easy way to generate random number with cupy kernel
            name = 'cli_pre_infer_n2f_kernel',reduce_dims=False,no_return=True)

# %% ../../nbs/CLI/dl.ipynb 8
if is_cuda_available():
    def _cli_pre_infer_n2f_cp(ref,sec):
        nlines, width = ref.shape
        out = cp.empty((1,2,nlines,width),dtype=cp.float32)
        mask = cp.empty((nlines,width),dtype=bool)
        _cli_pre_infer_n2f_kernel(ref,sec,cp.int32(nlines),cp.int32(width),out,mask,size=nlines*width,block_size=128)

        nan_pos = cp.where(mask)
        random_phase = cp.random.uniform(-cp.pi,cp.pi,len(nan_pos[0]))
        out[0,0,nan_pos[0],nan_pos[1]] = cp.cos(random_phase)
        out[0,1,nan_pos[0],nan_pos[1]] = cp.sin(random_phase)
        return out, mask

# %% ../../nbs/CLI/dl.ipynb 12
def _cli_n2f_cpu(
    ref,
    sec,
    chunks:tuple=None, # chunksize, intf.shape by default 
    depths:tuple=(0,0), # width of the boundary
    model:str=None, # path to the model in onnx format, use the model comes with this package by default
):
    if model is None:
        model = importlib.resources.files('moraine')/'dl_model/n2f.onnx'
    shape = ref.shape
    if chunks is None: chunks = shape
    in_slices, out_slices, map_slices = chunkwise_slicing_mapping(shape,chunks,depths)
    out = np.empty_like(ref)
    
    session = mr.dl._ort_session(model,cuda=False)
    for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
        input_intf_slice, mask_slice = _cli_pre_infer_n2f_numba(ref[in_slice],sec[in_slice])
        infer_out_slice = session.run([session.get_outputs()[0].name,],{session.get_inputs()[0].name: input_intf_slice})[0]
        out[out_slice] = mr.dl._after_infer_n2f_numba(infer_out_slice,mask_slice)[map_slice]
    return out

# %% ../../nbs/CLI/dl.ipynb 13
def _cli_n2f_np_in_gpu(
    ref,
    sec,
    chunks:tuple=None, # chunksize, intf.shape by default 
    depths:tuple=(0,0), # width of the boundary
    model:str=None, # path to the model in onnx format, use the model comes with this package by default
):
    if model is None:
        model = importlib.resources.files('moraine')/'dl_model/n2f.onnx'
    shape = ref.shape
    if chunks is None: chunks = shape
    in_slices, out_slices, map_slices = chunkwise_slicing_mapping(shape,chunks,depths)
    out = np.empty_like(ref)
    
    session = _ort_session(model,cuda=True)
    for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
        ref_slice = cp.asarray(ref[in_slice])
        sec_slice = cp.asarray(sec[in_slice])
        input_intf_slice, mask_slice = _cli_pre_infer_n2f_cp(ref_slice,sec_slice)
        
        input_intf_slice = cp.ascontiguousarray(input_intf_slice)
        output_intf_slice = cp.ascontiguousarray(cp.empty_like(input_intf_slice))
        io_binding = session.io_binding()
        io_binding.bind_input(
                name=session.get_inputs()[0].name,
                device_type='cuda',
                device_id=input_intf_slice.device.id,
                element_type=input_intf_slice.dtype,
                shape=list(input_intf_slice.shape),
                buffer_ptr=input_intf_slice.data.ptr,
        )
        io_binding.bind_output(
                name=session.get_outputs()[0].name,
                device_type='cuda',
                device_id=output_intf_slice.device.id,
                element_type=output_intf_slice.dtype,
                shape=list(output_intf_slice.shape),
                buffer_ptr=output_intf_slice.data.ptr,
        )
        io_binding.synchronize_inputs()
        session.run_with_iobinding(io_binding)
        out[out_slice] = (mr.dl._after_infer_n2f_cp(output_intf_slice,mask_slice)[map_slice]).get()
    return out

# %% ../../nbs/CLI/dl.ipynb 17
@mc_logger
def n2f(
    rslc:str, # input: rslc stack, shape (nlines, width, nimages)
    intf:str, # output: filtered intfergrams stack, shape (nlines, width, nimage_pairs )
    image_pairs:np.ndarray, # input: image pairs
    chunks:tuple=None, # parallel processing azimuth/range chunk size, optional. Default: rslc.chunks[:2]
    out_chunks:tuple=None, # output chunks
    depths:tuple=(0,0), # width of the boundary
    model:str=None, # path to the model in onnx format, use the model comes with this package by default
    cuda:bool=False, # if use cuda for processing, false by default
    processes=None, # use process for dask worker over thread, the default is True for cpu, only applied if cuda==False
    n_workers=None, # number of dask worker, the default is 1 for cpu, number of GPUs for cuda
    threads_per_worker=None, # number of threads per dask worker, the default is 1 for cpu, only applied if cuda==False
    rmm_pool_size=0.9, # set the rmm pool size, only applied when cuda==True
    **dask_cluster_arg, # other dask local/cudalocal cluster args
):
    '''Noise2Fringe interferogram filtering.
    '''
    rslc_path = rslc
    intf_path = intf
    logger = logging.getLogger(__name__)

    rslc_zarr = zarr.open(rslc_path,mode='r')
    logger.zarr_info(rslc_path, rslc_zarr)
    assert rslc_zarr.ndim == 3, "rslc dimentation is not 3."
    nlines, width, nimage = rslc_zarr.shape
    if chunks is None: chunks = rslc_zarr.chunks[:2]
    if out_chunks is None: out_chunks = rslc_zarr.chunks[:2]
    az_chunk, r_chunk = chunks
    logger.info(f'processing azimuth chunk size: {az_chunk}')
    logger.info(f'processing range chunk size: {r_chunk}')

    if cuda:
        Cluster = LocalCUDACluster; cluster_args= {
            'n_workers':n_workers,
            'rmm_pool_size':rmm_pool_size}
        cluster_args.update(dask_cluster_arg)
    else:
        if processes is None: processes = True
        if n_workers is None: n_workers = 1
        if threads_per_worker is None: threads_per_worker = 1
        Cluster = LocalCluster; cluster_args = {'processes':processes, 'n_workers':n_workers, 'threads_per_worker':threads_per_worker}
        cluster_args.update(dask_cluster_arg)

    n_image_pairs = image_pairs.shape[0]
    
    logger.info('starting dask cluster.')
    with Cluster(**cluster_args) as cluster, Client(cluster) as client:
        logger.info('dask cluster started.')
        logger.dask_cluster_info(cluster)
        if cuda:
            client.run(cp.cuda.set_allocator, rmm_cupy_allocator)
            n2f_delayed = delayed(_cli_n2f_np_in_gpu,pure=True,nout=1)
        else:
            n2f_delayed = delayed(_cli_n2f_cpu,pure=True,nout=1)

        ref_cpu_rslc = dask_from_zarr(rslc_path, chunks=(*rslc_zarr.shape[0:2],1))
        # use two different rslc to make dask do not hold too much data 
        sec_cpu_rslc = dask_from_zarr(rslc_path, chunks=(*rslc_zarr.shape[0:2],1))
        logger.darr_info('rslc', ref_cpu_rslc)
        intf = np.empty((1,1,n_image_pairs),dtype=object)
        for i in range(n_image_pairs):
            ref_i, sec_i = image_pairs[i]
            intf[0,0,i] = n2f_delayed(ref_cpu_rslc[:,:,ref_i].to_delayed()[0,0],sec_cpu_rslc[:,:,sec_i].to_delayed()[0,0],chunks=chunks,depths=depths,model=model)
            intf[0,0,i] = da.from_delayed(intf[0,0,i],shape=rslc_zarr.shape[:2],meta=np.array((),dtype=ref_cpu_rslc.dtype)).reshape(*rslc_zarr.shape[:2],1)
        intf = da.block(intf.tolist())
        logger.info('got filtered interferograms.')
        logger.darr_info('intf', intf)

        logger.info('saving filtered interferograms.')
        _intf = dask_to_zarr(intf,intf_path,chunks=(*out_chunks,1))

        logger.info('computing graph setted. doing all the computing.')
        futures = client.persist([_intf,])
        progress(futures,notebook=False)
        da.compute(futures)
        logger.info('computing finished.')
    logger.info('dask cluster closed.')

# %% ../../nbs/CLI/dl.ipynb 24
def _cli_n2ft(
    x:np.ndarray, # x coordinate, e.g., longitude, shape (n,) np.floating
    y:np.ndarray, # y coordinate, e.g., latitude, shape (n,) np.floating
    ref:np.ndarray, # reference slc, shape(n,) np.complex64
    sec:np.ndarray, # secondary slc, shape(n,) np.complex64
    in_indices,
    out_slices,
    map_indices,
    chunks:int=None, # chunksize, intf.shape[0] by default 
    k:int=128, # halo size for chunkwise processing
    model:str=None, # path to the model in onnx format, use the model comes with this package by default
    cuda:bool=False # use gpu for inference
):
    xp = np #mr.utils_.get_array_module(intf) # probability will add support for input and output with cupy array
    if model is None:
        model = importlib.resources.files('moraine')/'dl_model/n2ft.onnx'
    if cuda:
        session = _ort_session(model,cuda=True)
    else:
        session = _ort_session(model,cuda=False)

    intf = intf_func(ref, sec)
    n = intf.shape[0]
    if (chunks is None) or (chunks >= n):
        out = _infer_n2ft(x, y, intf, session)
    else:
        out = xp.empty_like(intf)
        for in_idx, out_slice, map_idx in zip(in_indices, out_slices, map_indices):
            out[out_slice] = _infer_n2ft(x[in_idx],y[in_idx],intf[in_idx],session)[map_idx]

    return out

# %% ../../nbs/CLI/dl.ipynb 25
@mc_logger
def n2ft(
    x:str, # input: x coordinate, e.g., longitude, shape (n,)
    y:str, # input: y coordinate, e.g., latitude, shape (n,)
    rslc:str, # input: rslc stack, shape (n, nimages)
    intf:str, # output: filtered intfergrams stack, shape (n, nimage_pairs )
    image_pairs:np.ndarray, # input: image pairs
    chunks:int=None, # parallel processing point chunk size, optional. Default: rslc.chunks[0]
    out_chunks:int=None, # output point chunk size, Default: rslc.chunks[0]
    k:int=128, # halo size for chunkwise processing
    model:str=None, # path to the model in onnx format, use the model comes with this package by default
    cuda:bool=False, # if use cuda for processing, false by default
    processes=None, # use process for dask worker over thread, the default is True for cpu, only applied if cuda==False
    n_workers=None, # number of dask worker, the default is 1 for cpu, number of GPUs for cuda
    threads_per_worker=None, # number of threads per dask worker, the default is 1 for cpu, only applied if cuda==False
    rmm_pool_size=0.9, # set the rmm pool size, only applied when cuda==True
    **dask_cluster_arg, # other dask local/cudalocal cluster args
):
    '''Noise2Fringe interferogram filtering.
    '''
    logger = logging.getLogger(__name__)
    logger.info('load coordinates')
    pc_x_data = parallel_read_zarr(zarr.open(x,mode='r'),(slice(None),))
    pc_y_data = parallel_read_zarr(zarr.open(y,mode='r'),(slice(None),))
    logger.info('Done')
    rslc_path = rslc
    intf_path = intf

    rslc_zarr = zarr.open(rslc_path,mode='r')
    logger.zarr_info(rslc_path, rslc_zarr)
    assert rslc_zarr.ndim == 2, "rslc dimentation is not 2."
    npoint, nimage = rslc_zarr.shape

    nimage_pairs = image_pairs.shape[0]
    if chunks is None: chunks = rslc_zarr.chunks[0]
    if out_chunks is None: out_chunks = rslc_zarr.chunks[0]

    logger.info(f'processing point chunk size: {chunks}')
    logger.info('distributing every processing chunk with halo data')
    in_indices, out_slices, map_indices = chunkwise_knn_mapping(pc_x_data, pc_y_data, chunks, k=k)
    in_indices_size = [len(in_idx) for in_idx in in_indices]
    logger.info(f'processing chunk size with halo data: {in_indices_size}')

    if cuda:
        Cluster = LocalCUDACluster; cluster_args= {
            'n_workers':n_workers,
            'rmm_pool_size':rmm_pool_size}
        cluster_args.update(dask_cluster_arg)
    else:
        if processes is None: processes = True
        if n_workers is None: n_workers = 1
        if threads_per_worker is None: threads_per_worker = 1
        Cluster = LocalCluster; cluster_args = {'processes':processes, 'n_workers':n_workers, 'threads_per_worker':threads_per_worker}
        cluster_args.update(dask_cluster_arg)

    n_image_pairs = image_pairs.shape[0]
    
    logger.info('starting dask cluster.')
    with Cluster(**cluster_args) as cluster, Client(cluster) as client:
        logger.info('dask cluster started.')
        logger.dask_cluster_info(cluster)
        if cuda: client.run(cp.cuda.set_allocator, rmm_cupy_allocator)
        n2ft_delayed = delayed(_cli_n2ft,pure=True,nout=1)

        ref_cpu_rslc = dask_from_zarr(rslc_path, chunks=(npoint,1))
        # use two different rslc to make dask do not hold too much data 
        sec_cpu_rslc = dask_from_zarr(rslc_path, chunks=(npoint,1))
        logger.darr_info('rslc', ref_cpu_rslc)
        intf = np.empty((1,n_image_pairs),dtype=object)
        for i in range(n_image_pairs):
            ref_i, sec_i = image_pairs[i]
            intf[0,i] = n2ft_delayed(pc_x_data, pc_y_data, ref_cpu_rslc[:,ref_i].to_delayed()[0],sec_cpu_rslc[:,sec_i].to_delayed()[0],
                                     in_indices, out_slices, map_indices,
                                     chunks=chunks, k=k, model=model, cuda=cuda
                                    )
            intf[0,i] = da.from_delayed(intf[0,i],shape=(npoint,),meta=np.array((),dtype=ref_cpu_rslc.dtype)).reshape(npoint,1)
        intf = da.block(intf.tolist())
        logger.info('got filtered interferograms.')
        logger.darr_info('intf', intf)

        logger.info('saving filtered interferograms.')
        _intf = dask_to_zarr(intf,intf_path,chunks=(out_chunks,1))

        logger.info('computing graph setted. doing all the computing.')
        futures = client.persist([_intf,])
        progress(futures,notebook=False)
        da.compute(futures)
        logger.info('computing finished.')
    logger.info('dask cluster closed.')
