"""Pixel Quality Metrics"""

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

# %% auto 0
__all__ = ['temp_coh']

# %% ../nbs/API/pqm.ipynb 4
import numpy as np
import moraine as mr
from .utils_ import is_cuda_available, get_array_module
if is_cuda_available():
    import cupy as cp
from numba import prange
from .utils_ import ngpjit

# %% ../nbs/API/pqm.ipynb 5
# @ngpjit
# def _temp_coh_pc_numba(
#     intf:np.ndarray,# complex interferograms/coherence metrix, dtype np.complex64
#     rslc:np.ndarray, # complex rslc/phase history, dtype np.complex64
#     image_pairs:np.ndarray, # image pairs
# ):
#     nimages = rslc.shape[-1]
#     n_points = rslc.shape[0]
#     n_image_pairs = image_pairs.shape[0]
#     temp_coh = np.empty(n_points,dtype=np.float32)
#     for i in prange(n_points):
#         rslc_ = rslc[i]
#         intf_ = intf[i]
#         for j in range(nimages):
#             rslc_[j] = rslc_[j]/abs(rslc_[j])
#         for j in range(n_image_pairs):
#             intf_[j] = intf_[j]/abs(intf_[j])
#         _t_coh = np.float32(0.0)
#         for j in range(n_image_pairs):
#             n, k = image_pairs[j,0],image_pairs[j,1]
#             rslc_intf_ = np.conjugate(rslc_[n])*rslc_[k]
#             diff_ph = intf_[j]*rslc_intf_
#             _t_coh += diff_ph.real
#         _t_coh = _t_coh/n_image_pairs
#         temp_coh[i] = _t_coh
#     return temp_coh

# %% ../nbs/API/pqm.ipynb 6
@ngpjit
def _temp_coh_pc_numba(
    intf:np.ndarray,# complex interferograms/coherence metrix, dtype np.complex64
    rslc:np.ndarray, # complex rslc/phase history, dtype np.complex64
    image_pairs:np.ndarray, # image pairs
):
    nimages = rslc.shape[-1]
    n_points = rslc.shape[0]
    n_image_pairs = image_pairs.shape[0]
    temp_coh = np.empty(n_points,dtype=np.float32)
    for i in prange(n_points):
        rslc_ = rslc[i]
        intf_ = intf[i]
        for j in range(nimages):
            rslc_[j] = rslc_[j]/abs(rslc_[j])
        for j in range(n_image_pairs):
            intf_[j] = intf_[j]/abs(intf_[j])
        _t_coh = np.complex64(0.0)
        for j in range(n_image_pairs):
            n, k = image_pairs[j,0],image_pairs[j,1]
            rslc_intf_ = np.conjugate(rslc_[n])*rslc_[k]
            diff_ph = intf_[j]*rslc_intf_
            _t_coh += diff_ph
        _t_coh = np.abs(_t_coh)/n_image_pairs
        temp_coh[i] = _t_coh
    return temp_coh

# %% ../nbs/API/pqm.ipynb 7
# if is_cuda_available():
#     _temp_coh_pc_kernel = cp.ElementwiseKernel(
#         'raw T intf, raw T rslc, raw I image_pairs, int32 n_points, int32 nimages, int32 n_image_pairs',
#         'raw float32 temp_coh',
#         '''
#         if (i >= n_points) return;
#         float t_coh_ = 0;
#         int j; int n; int k;
#         int intf_idx; int rslc_n_idx; int rslc_k_idx;
#         float norm_factor = 0;
#         for (j = 0; j< n_image_pairs; j++){
#             n = image_pairs[j*2];
#             k = image_pairs[j*2+1];
#             intf_idx = i*n_image_pairs+j;
#             rslc_n_idx = i*nimages+n;
#             rslc_k_idx = i*nimages+k;
#             norm_factor = sqrt(norm(intf[intf_idx])*norm(rslc[rslc_n_idx])*norm(rslc[rslc_k_idx]));
#             t_coh_ += real(intf[intf_idx]*conj(rslc[rslc_n_idx])*rslc[rslc_k_idx]/norm_factor);
#         }
#         temp_coh[i] = t_coh_/n_image_pairs;
#         ''',
#         name = 'temp_coh_pc_kernel',no_return=True,
#     )

# %% ../nbs/API/pqm.ipynb 8
if is_cuda_available():
    _temp_coh_pc_kernel = cp.ElementwiseKernel(
        'raw T intf, raw T rslc, raw I image_pairs, int32 n_points, int32 nimages, int32 n_image_pairs',
        'raw float32 temp_coh',
        '''
        if (i >= n_points) return;
        T t_coh_ = T(0.0,0.0);
        int j; int n; int k;
        int intf_idx; int rslc_n_idx; int rslc_k_idx;
        float norm_factor = 0;
        for (j = 0; j< n_image_pairs; j++){
            n = image_pairs[j*2];
            k = image_pairs[j*2+1];
            intf_idx = i*n_image_pairs+j;
            rslc_n_idx = i*nimages+n;
            rslc_k_idx = i*nimages+k;
            norm_factor = sqrt(norm(intf[intf_idx])*norm(rslc[rslc_n_idx])*norm(rslc[rslc_k_idx]));
            t_coh_ += intf[intf_idx]*conj(rslc[rslc_n_idx])*rslc[rslc_k_idx]/norm_factor;
        }
        temp_coh[i] = abs(t_coh_)/n_image_pairs;
        ''',
        name = 'temp_coh_pc_kernel',no_return=True,
    )

# %% ../nbs/API/pqm.ipynb 9
def temp_coh(
    intf:np.ndarray,# complex interferograms/coherence metrix, dtype cp/np.complex64, shape 2D(pc) or 3D(ras)
    rslc:np.ndarray, # complex rslc/phase history, dtype cp/np.complex64, shape 2D(pc) or 3D(ras)
    image_pairs:np.ndarray=None, # image pairs
    block_size:int=128, # the CUDA block size, only applied for cuda
):
    '''Estimation of temporal coherence.'''
    xp = get_array_module(intf)
    assert intf.ndim == rslc.ndim
    if intf.ndim == 3:
        # convert to pc
        is_ras = True
        height, width = intf.shape[:2]
        n_points =height*width
        intf = intf.reshape((-1,intf.shape[-1]))
        rslc = rslc.reshape((-1,rslc.shape[-1]))
    elif intf.ndim == 2:
        is_ras = False
        n_points = rslc.shape[0]

    nimages = rslc.shape[-1]
    if image_pairs is None:
        image_pairs = mr.TempNet.from_bandwidth(nimages).image_pairs
    image_pairs = image_pairs.astype(np.int32)
    
    if xp is np:
        out_temp_coh = _temp_coh_pc_numba(intf,rslc,image_pairs)
    else:
        image_pairs = cp.asarray(image_pairs)
        n_image_pairs = image_pairs.shape[0]
        out_temp_coh = cp.empty(n_points, dtype=cp.float32)
        _temp_coh_pc_kernel(intf, rslc, image_pairs, cp.int32(n_points),cp.int32(nimages),cp.int32(n_image_pairs), out_temp_coh, size=n_points, block_size=block_size)

    if is_ras:
        out_temp_coh = out_temp_coh.reshape((height,width))
    return out_temp_coh
