"""Deep learning based operators"""

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

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

# %% ../nbs/API/dl.ipynb 3
import numpy as np
from numba import prange
from scipy.spatial import KDTree
import fpsample
import os

import onnxruntime
import importlib
from pathlib import Path
import requests
import math
import random
import moraine as mr
from .utils_ import ngjit, ngpjit
from .utils_ import is_cuda_available, get_array_module, get_n_cpus_avail
from .chunk_ import chunkwise_slicing_mapping, chunkwise_knn_mapping
if is_cuda_available():
    import cupy as cp

# %% ../nbs/API/dl.ipynb 5
def _fetch_github_file(
    gid:str, # github id
    repo:str, # github repo
    branch:str, # repo branch
    file_path:str, # path to the file
    out_path:str, # output path
):
    url = '/'.join(['https://raw.githubusercontent.com', gid, repo, 'refs','heads', branch, file_path])
    print(' '.join(['Downloading', url,'to',out_path]))
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()  # Raise an exception for bad status codes (4xx or 5xx)

        with open(out_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"File '{out_path}' downloaded successfully.")
    except requests.exceptions.RequestException as e:
        print(f"Error downloading file: {e}")

# %% ../nbs/API/dl.ipynb 6
def download_dl_model(
    models:str|list=None, # deep learning models to be downloaded, all available models by default
    path:str=None, # directory to save these models, inside installed Moraine package by default
):
    if models is None:
        models = ['n2f','n2fs3d','n2ft']
    else:
        if isinstance(models,str):
            models = [models,]

    if path is None: 
        path = importlib.resources.files('moraine')/'dl_model'
    else:
        path = Path(path)

    if 'n2f' in models:
        _fetch_github_file('kanglcn','n2f','main','n2f.onnx',str(path/'n2f.onnx'))
    if 'n2fs3d' in models:
        _fetch_github_file('kanglcn','n2f','main','n2fs3d.onnx',str(path/'n2fs3d.onnx'))
    if 'n2ft' in models:
        _fetch_github_file('kanglcn','n2ft','main','n2ft.onnx',str(path/'n2ft.onnx'))
        _fetch_github_file('kanglcn','n2ft','main','n2ft.onnx.data',str(path/'n2ft.onnx.data'))

# %% ../nbs/API/dl.ipynb 9
def _ort_session(
    path:str, # path to the model in onnx format
    cuda:bool=False, # if use cuda or not
):
    session_options = onnxruntime.SessionOptions()
    n_cpu = get_n_cpus_avail()
    
    session_options.intra_op_num_threads = n_cpu
    session_options.inter_op_num_threads = n_cpu

    if cuda:
        import cupy as cp
        providers = [("CUDAExecutionProvider", {"device_id": cp.cuda.runtime.getDevice()}),'CPUExecutionProvider']
    else:
        providers = ['CPUExecutionProvider']
    ort_session = onnxruntime.InferenceSession(path, sess_options=session_options, providers=providers)
    return ort_session

# %% ../nbs/API/dl.ipynb 10
@ngpjit
def _pre_infer_n2f_numba(intf):
    nlines, width = intf.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 = intf[i,j]
            if math.isnan(intf_i_j.real):
                mask[i,j] = True
                random_phase = 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/API/dl.ipynb 11
if is_cuda_available():
    _pre_infer_n2f_kernel = cp.ElementwiseKernel(
            'raw T intf, int32 nlines, int32 width',
            'raw float32 out, raw bool mask',
            '''
            int npixels = nlines*width;
            if (i >= npixels) return;

            //int az = i/width;
            //int r = i-az*width;
            T intf_i = intf[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 = 'pre_infer_n2f_kernel',no_return=True)

# %% ../nbs/API/dl.ipynb 12
if is_cuda_available():
    def _pre_infer_n2f_cp(intf):
        nlines, width = intf.shape
        out = cp.empty((1,2,nlines,width),dtype=cp.float32)
        mask = cp.empty((nlines,width),dtype=cp.bool_)
        _pre_infer_n2f_kernel(intf,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/API/dl.ipynb 16
@ngpjit
def _after_infer_n2f_numba(
    infer_out,
    mask,
):
    nlines, width = mask.shape
    out = np.empty((nlines,width),dtype=np.complex64)
    for i in prange(nlines):
        for j in prange(width):
            if mask[i,j]:
                out[i,j] = np.nan+1j*np.nan
            else:
                out[i,j] = infer_out[0,0,i,j]+1j*infer_out[0,1,i,j]
    return out

# %% ../nbs/API/dl.ipynb 17
if is_cuda_available():
    _after_infer_n2f_kernel = cp.ElementwiseKernel(
            'raw float32 infer_out, raw bool mask, int32 nlines, int32 width',
            'raw T out',
            '''
            int npixels = nlines*width;
            if (i >= npixels) return;

            if (mask[i]){
                out[i] = T(CUDART_NAN,CUDART_NAN);
            }
            else{
                out[i] = T(infer_out[i],infer_out[npixels+i]);
            }
            ''',
            preamble = '#include <cupy/math_constants.h>',
            name = 'after_infer_n2f_kernel',no_return=True)

# %% ../nbs/API/dl.ipynb 18
if is_cuda_available():
    def _after_infer_n2f_cp(infer_out,mask):
        nlines, width = mask.shape
        out = cp.empty((nlines,width),dtype=cp.complex64)
        _after_infer_n2f_kernel(infer_out,mask,cp.int32(nlines),cp.int32(width),out,size=nlines*width,block_size=128)
        return out

# %% ../nbs/API/dl.ipynb 23
def _infer_n2f_cpu(
    intf,
    session,
):
    input_intf, mask = _pre_infer_n2f_numba(intf)
    infer_out = session.run([session.get_outputs()[0].name,],{session.get_inputs()[0].name: input_intf})[0]
    return _after_infer_n2f_numba(infer_out,mask)

# %% ../nbs/API/dl.ipynb 24
def _infer_n2f_gpu(
    intf,
    session,
):
    input_intf, mask = _pre_infer_n2f_cp(intf)
    input_intf = cp.ascontiguousarray(input_intf)
    out = cp.ascontiguousarray(cp.empty_like(input_intf))
    io_binding = session.io_binding()
    io_binding.bind_input(
            name=session.get_inputs()[0].name,
            device_type='cuda',
            device_id=input_intf.device.id,
            element_type=input_intf.dtype,
            shape=list(input_intf.shape),
            buffer_ptr=input_intf.data.ptr,
    )
    io_binding.bind_output(
            name=session.get_outputs()[0].name,
            device_type='cuda',
            device_id=out.device.id,
            element_type=out.dtype,
            shape=list(out.shape),
            buffer_ptr=out.data.ptr,
    )
    io_binding.synchronize_inputs()
    session.run_with_iobinding(io_binding)
    return _after_infer_n2f_cp(out,mask)

# %% ../nbs/API/dl.ipynb 25
def n2f(
    intf:np.ndarray, # interferogram, 2d np.complex64 or cp.complex64
    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
):
    xp = mr.utils_.get_array_module(intf)
    if model is None:
        model = importlib.resources.files('moraine')/'dl_model/n2f.onnx'
    shape = intf.shape
    if chunks is None: chunks = shape
    in_slices, out_slices, map_slices = chunkwise_slicing_mapping(shape,chunks,depths)
    out = xp.empty_like(intf)

    if xp is np:
        session = _ort_session(model,cuda=False)
        for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
            out[out_slice] = _infer_n2f_cpu(intf[in_slice],session)[map_slice]
    else:
        session = _ort_session(model,cuda=True)
        for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
            out[out_slice] = _infer_n2f_gpu(intf[in_slice],session)[map_slice]

    return out

# %% ../nbs/API/dl.ipynb 33
def _n2f_np_in_gpu(
    intf:np.ndarray, # interferogram, 2d np.complex64 or cp.complex64
    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
):
    '''compare with n2f, input and output are np.ndarray but use gpu for inference'''
    if model is None:
        model = importlib.resources.files('moraine')/'dl_model/n2f.onnx'
    shape = intf.shape
    if chunks is None: chunks = shape
    in_slices, out_slices, map_slices = chunkwise_slicing_mapping(shape,chunks,depths)
    out = np.empty_like(intf)

    session = _ort_session(model,cuda=True)
    for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
        #out[out_slice] = _infer_n2f_cpu(intf[in_slice],session)[map_slice]

        out[out_slice] = _infer_n2f_gpu(cp.asarray(intf[in_slice]),session)[map_slice].get()
    return out

# %% ../nbs/API/dl.ipynb 35
@ngpjit
def _pre_infer_n2fs3d_numba(adi,intf):
    nlines, width = intf.shape
    out = np.empty((1,3,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 = intf[i,j]
            adi_i_j = adi[i,j]
            if math.isnan(intf_i_j.real) or math.isnan(adi_i_j):
                mask[i,j] = True
                out[0,0,i,j] = 1.0
                random_phase = random.uniform(-math.pi,math.pi)
                out[0,1,i,j] = math.cos(random_phase)
                out[0,2,i,j] = math.sin(random_phase)
            else:
                mask[i,j] = False
                out[0,0,i,j] = adi_i_j
                amp = abs(intf_i_j)
                out[0,1,i,j] = intf_i_j.real/amp
                out[0,2,i,j] = intf_i_j.imag/amp
    return out, mask

# %% ../nbs/API/dl.ipynb 36
if is_cuda_available():
    _pre_infer_n2fs3d_kernel = cp.ElementwiseKernel(
            'raw float32 adi, raw T intf, int32 nlines, int32 width',
            'raw float32 out, raw bool mask',
            '''
            int npixels = nlines*width;
            if (i >= npixels) return;

            //int az = i/width;
            //int r = i-az*width;
            T intf_i = intf[i];
            float adi_i = adi[i];
            if (isnan(intf_i.real())||isnan(adi_i) ){
                mask[i] = true;
                out[i] = 1.0;
            }
            else{
                mask[i] = false;
                float amp = abs(intf_i);
                out[i] = adi_i;
                out[npixels+i] = intf_i.real()/amp;
                out[2*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 = 'pre_infer_n2f_kernel',no_return=True)

# %% ../nbs/API/dl.ipynb 37
if is_cuda_available():
    def _pre_infer_n2fs3d_cp(adi,intf):
        nlines, width = intf.shape
        out = cp.empty((1,3,nlines,width),dtype=cp.float32)
        mask = cp.empty((nlines,width),dtype=cp.bool_)
        _pre_infer_n2fs3d_kernel(adi,intf,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,1,nan_pos[0],nan_pos[1]] = cp.cos(random_phase)
        out[0,2,nan_pos[0],nan_pos[1]] = cp.sin(random_phase)
        return out, mask

# %% ../nbs/API/dl.ipynb 45
def _infer_n2fs3d_cpu(
    adi,
    intf,
    session,
):
    input_intf, mask = _pre_infer_n2fs3d_numba(adi,intf)
    infer_out = session.run([session.get_outputs()[0].name,],{session.get_inputs()[0].name: input_intf})[0]
    return _after_infer_n2f_numba(infer_out,mask)

# %% ../nbs/API/dl.ipynb 46
def _infer_n2fs3d_gpu(
    adi,
    intf,
    session,
):
    input_intf, mask = _pre_infer_n2fs3d_cp(adi,intf)
    input_intf = cp.ascontiguousarray(input_intf)
    out = cp.ascontiguousarray(cp.empty((1,2,*mask.shape),dtype=input_intf.dtype))
    io_binding = session.io_binding()
    io_binding.bind_input(
            name=session.get_inputs()[0].name,
            device_type='cuda',
            device_id=input_intf.device.id,
            element_type=input_intf.dtype,
            shape=list(input_intf.shape),
            buffer_ptr=input_intf.data.ptr,
    )
    io_binding.bind_output(
            name=session.get_outputs()[0].name,
            device_type='cuda',
            device_id=out.device.id,
            element_type=out.dtype,
            shape=list(out.shape),
            buffer_ptr=out.data.ptr,
    )
    io_binding.synchronize_inputs()
    session.run_with_iobinding(io_binding)
    return _after_infer_n2f_cp(out,mask)

# %% ../nbs/API/dl.ipynb 47
def n2fs3d(
    adi:np.ndarray, # amplitude dispersion index, 2d np.float32 or cp.float32
    intf:np.ndarray, # interferogram, 2d np.complex64 or cp.complex64
    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
):
    xp = mr.utils_.get_array_module(intf)
    if model is None:
        model = importlib.resources.files('moraine')/'dl_model/n2fs3d.onnx'
    shape = intf.shape
    if chunks is None: chunks = shape
    in_slices, out_slices, map_slices = chunkwise_slicing_mapping(shape,chunks,depths)
    out = xp.empty_like(intf)

    if xp is np:
        session = _ort_session(model,cuda=False)
        for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
            out[out_slice] = _infer_n2fs3d_cpu(adi[in_slice],intf[in_slice],session)[map_slice]
    else:
        session = _ort_session(model,cuda=True)
        for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
            out[out_slice] = _infer_n2fs3d_gpu(adi[in_slice],intf[in_slice],session)[map_slice]

    return out

# %% ../nbs/API/dl.ipynb 55
def _n2fs3d_np_in_gpu(
    adi:np.ndarray, # adi, 2d np.float32
    intf:np.ndarray, # interferogram, 2d np.complex64 or cp.complex64
    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
):
    '''compare with n2f, input and output are np.ndarray but use gpu for inference'''
    if model is None:
        model = importlib.resources.files('moraine')/'dl_model/n2fs3d.onnx'
    shape = intf.shape
    if chunks is None: chunks = shape
    in_slices, out_slices, map_slices = chunkwise_slicing_mapping(shape,chunks,depths)
    out = np.empty_like(intf)

    session = _ort_session(model,cuda=True)
    for in_slice, out_slice, map_slice in zip(in_slices, out_slices, map_slices):
        out[out_slice] = _infer_n2fs3d_gpu(cp.asarray(adi[in_slice]),cp.asarray(intf[in_slice]),session)[map_slice].get()
    return out

# %% ../nbs/API/dl.ipynb 57
@ngpjit
def _weights(dd):
    N, K = dd.shape #(N, 3)
    out = np.empty_like(dd)
    for i in prange(N):  # OpenMP-style parallel loop
        denom = 0.0
        for j in range(K):
            out[i, j] = 1.0 / (dd[i, j] + 1e-8)
            denom += out[i, j]
        for j in range(K):
            out[i, j] /= denom
    return out

# %% ../nbs/API/dl.ipynb 58
def _sample_and_knn(pos, k=16, workers=-1):  # (N, 2)
    N0 = pos.shape[0]
    N1 = N0 // 4
    N2 = N1 // 4
    N3 = N2 // 4

    p0 = pos  # [N0, 2]
    idx01 = fpsample.bucket_fps_kdtree_sampling(p0, N1)
    p1 = p0[idx01]
    idx12 = fpsample.bucket_fps_kdtree_sampling(p1, N2)
    p2 = p1[idx12]
    idx23 = fpsample.bucket_fps_kdtree_sampling(p2, N3)
    p3 = p2[idx23]

    tree0 = KDTree(p0)
    tree1 = KDTree(p1)
    tree2 = KDTree(p2)
    tree3 = KDTree(p3)

    # KNN queries
    _, ii_00 = tree0.query(p0, k=k, workers=workers) #[N0, k]
    _, ii_01 = tree0.query(p1, k=k, workers=workers) #[N1, k]
    _, ii_11 = tree1.query(p1, k=k, workers=workers)
    _, ii_12 = tree1.query(p2, k=k, workers=workers)
    _, ii_22 = tree2.query(p2, k=k, workers=workers)
    _, ii_23 = tree2.query(p3, k=k, workers=workers)
    _, ii_33 = tree3.query(p3, k=k, workers=workers)

    dd_10, ii_10 = tree1.query(p0, k=3, workers=workers)
    dd_21, ii_21 = tree2.query(p1, k=3, workers=workers)
    dd_32, ii_32 = tree3.query(p2, k=3, workers=workers)

    ww_10 = _weights(dd_10).astype(np.float32)
    ww_21 = _weights(dd_21).astype(np.float32)
    ww_32 = _weights(dd_32).astype(np.float32)

    return (
        ii_00, ii_01, ii_11, ii_12, ii_22, ii_23, ii_33,
        ii_10, ii_21, ii_32,
        ww_10, ww_21, ww_32
    )

# %% ../nbs/API/dl.ipynb 59
@ngpjit
def _pos_norm(x, y): # (N,), (N,)
    '''return normalized pos, (N, 2)'''
    y_min, y_max = np.min(y), np.max(y)
    x_min, x_max = np.min(x), np.max(x)
    y_norm = (y - y_min)/(y_max - y_min)
    x_norm = (x - x_min)/(x_max - x_min)

    return np.stack((x_norm, y_norm),axis=-1).astype(np.float32)

# %% ../nbs/API/dl.ipynb 60
@ngpjit
def _complex2channel(intf):
    # convert 1d complex array to 2 column
    n = intf.shape[0]
    out = np.empty((n,2),dtype=np.float32)
    for i in prange(n):
        amp_i = abs(intf[i])
        out[i,0] = intf[i].real/amp_i
        out[i,1] = intf[i].imag/amp_i
    return out

# %% ../nbs/API/dl.ipynb 61
@ngpjit
def _2channel2complex(intf):
    # convert (n,2) to (n,) complex array
    n = intf.shape[0]
    out = np.empty(n,dtype=np.complex64)
    for i in prange(n):
        out[i] = intf[i,0]+intf[i,1]*1j
    return out

# %% ../nbs/API/dl.ipynb 62
def _infer_n2ft(
    x,
    y,
    intf,
    session,
):
    pos = _pos_norm(x,y)
    keys = _sample_and_knn(pos)

    pos = pos[None,...]
    keys = tuple(key[None,...] for key in keys)
    intf = _complex2channel(intf)[None,...]

    inputs = (pos, intf) + keys
    input_names = [input.name for input in session.get_inputs()]
    inputs = dict(zip(input_names, inputs))
    infer_out = session.run([session.get_outputs()[0].name,], inputs)[0][0]

    filt_intf = _2channel2complex(infer_out)
    return filt_intf

# %% ../nbs/API/dl.ipynb 63
def 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
    intf:np.ndarray, # interferogram, shape(n,) np.complex64
    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)
    
    n = intf.shape[0]
    if (chunks is None) or (chunks >= n):
        out = _infer_n2ft(x, y, intf, session)
    else:
        out = xp.empty_like(intf)
        in_indices, out_slices, map_indices = chunkwise_knn_mapping(x, y, chunks, k=k)
        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
