"""Covariance and Coherence Matrix Estimation"""

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

# %% auto 0
__all__ = ['multi_look', 'intf', 'emperical_co', 'emperical_co_pc', 'uncompress_single_coh_numba', 'uncompress_coh', 'ad_intf_pc',
           'isPD', 'nearestPD', 'regularize_spectral']

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

# %% ../nbs/API/co.ipynb 6
@ngpjit
def multi_look(
    rslc1, # reference rslc
    rslc2, # secondary rslc
    looks=(1,1), # # range and azimuth looks
):
    '''multi looked interferogram for raster data.'''
    az_look, r_look = looks
    out_ny = math.floor(rslc1.shape[0]/az_look)
    out_nx = math.floor(rslc1.shape[1]/r_look)
    out = np.empty((out_ny, out_nx),dtype=rslc1.dtype)
    for i in prange(out_ny):
        for j in range(out_nx):
            az_start, az_stop = i*az_look, (i+1)*az_look
            r_start, r_stop = j*r_look, (j+1)*r_look
            rslc1_ = rslc1[az_start:az_stop, r_start:r_stop].flatten()
            rslc2_ = rslc2[az_start:az_stop, r_start:r_stop].flatten()
            dnmnt = math.sqrt(np.sum(rslc1_.real**2+rslc1_.imag**2)*np.sum(rslc2_.real**2+rslc2_.imag**2))
            out[i,j] = np.sum(rslc1_*rslc2_.conj())/dnmnt
    return out

# %% ../nbs/API/co.ipynb 7
@ngpjit
def intf(
    rslc1, # reference rslc, arbitrary dims, e.g., 3D(raster stack), 2D(raster or point cloud stack) or 1D(point cloud)
    rslc2, # secondary rslc
):
    '''1 by 1 look interferogram for both raster and point cloud.'''
    shape = rslc1.shape
    
    rslc1 = rslc1.reshape(-1)
    rslc2 = rslc2.reshape(-1)
    intf = np.empty(rslc1.shape,dtype=rslc1.dtype)
    for i in prange(rslc1.shape[0]):
        intf[i] = rslc1[i]*np.conj(rslc2[i])/(np.abs(rslc1[i])*np.abs(rslc2[i]))
    intf = intf.reshape(shape)
    return intf

# %% ../nbs/API/co.ipynb 9
if is_cuda_available():
    _emperical_co_kernel = cp.ElementwiseKernel(
        'raw T rslc, raw bool is_shp, int32 nlines, int32 width, int32 nimages, int32 az_half_win, int32 r_half_win',
        'raw T cov, raw T coh',
        '''
        if (i >= nlines*width) return;
        int az_win = 2*az_half_win+1;
        int r_win = 2*r_half_win+1;
        int win = az_win*r_win;
        
        int ref_az = i/width;
        int ref_r = i -ref_az*width;
    
        int sec_az, sec_r;
    
        int m,j; // index of each coherence matrix
        int k,l; // index of search window
        T _cov; // covariance
        float _amp2_m; // sum of amplitude square for image i
        float _amp2_j; // sum of amplitude aquare for image j
        int rslc_inx_m, rslc_inx_j;
        int n; // number of shp
    
        for (m = 0; m < nimages; m++) {
            for (j = 0; j < nimages; j++) {
                _cov = T(0.0, 0.0);
                _amp2_m = 0.0;
                _amp2_j = 0.0;
                n = 0;
                for (k = 0; k < az_win; k++) {
                    for (l = 0; l < r_win; l++) {
                        sec_az = ref_az-az_half_win+k;
                        sec_r = ref_r-r_half_win+l;
                        if (is_shp[i*win+k*r_win+l] && sec_az >= 0 && sec_az < nlines && sec_r >= 0 && sec_r < width) {
                            rslc_inx_m = (sec_az*width+sec_r)*nimages+m;
                            rslc_inx_j = (sec_az*width+sec_r)*nimages+j;
                            _amp2_m += norm(rslc[rslc_inx_m]);
                            _amp2_j += norm(rslc[rslc_inx_j]);
                            _cov += rslc[rslc_inx_m]*conj(rslc[rslc_inx_j]);
                            n += 1;
                            //if (i == 0 && m ==3 && j == 1) {
                            //    printf("%f",_cov.real());
                            //}
                        }
                    }
                }
                cov[(i*nimages+m)*nimages+j] = _cov/(float)n;
                //if ( i == 0 && m==3 && j ==1 ) printf("%d",((i*nimages+m)*nimages+j));
                _amp2_m = sqrt(_amp2_m*_amp2_j);
                coh[(i*nimages+m)*nimages+j] = _cov/_amp2_m;
            }
        }
        ''',
        name = 'emperical_co_kernel',reduce_dims = False,no_return=True
    )

# %% ../nbs/API/co.ipynb 11
def emperical_co(rslc:np.ndarray, # rslc stack, dtype: `cupy.complexfloating`
                 is_shp:np.ndarray, # shp bool, dtype: `cupy.bool`
                 block_size:int=128, # the CUDA block size, it only affects the calculation speed
                )-> tuple[np.ndarray,np.ndarray]: # the covariance and coherence matrix `cov` and `coh`
    '''
    Maximum likelihood covariance estimator.
    '''
    xp = get_array_module(rslc)
    if xp is np:
        raise NotImplementedError("Currently only cuda version available.")
    nlines, width, nimages = rslc.shape
    az_win, r_win = is_shp.shape[-2:]
    az_half_win = (az_win-1)//2
    r_half_win = (r_win-1)//2

    cov = cp.empty((nlines,width,nimages,nimages),dtype=rslc.dtype)
    coh = cp.empty((nlines,width,nimages,nimages),dtype=rslc.dtype)

    _emperical_co_kernel(rslc, is_shp, cp.int32(nlines),cp.int32(width),cp.int32(nimages),
                    cp.int32(az_half_win),cp.int32(r_half_win),cov,coh,size = nlines*width,block_size=block_size)
    return cov,coh

# %% ../nbs/API/co.ipynb 21
if is_cuda_available():
    _emperical_co_pc_kernel = cp.ElementwiseKernel(
        'raw T rslc, raw I az_idx, raw I r_idx, raw bool pc_is_shp, raw I image_pair_ref, raw I image_pair_sec, int32 nlines, int32 width, int32 nimages, int32 az_half_win, int32 r_half_win, int32 n_pc, int32 n_image_pair',
        'raw T coh',
        '''
        if (i >= n_pc) return;
        int az_win = 2*az_half_win+1;
        int r_win = 2*r_half_win+1;
        int win = az_win*r_win;
        
        int ref_az = az_idx[i];
        int ref_r = r_idx[i];
    
        int sec_az, sec_r;
    
        int m,j; // index of each coherence matrix
        int pair_i; // index of image pairs
        int k,l; // index of search window
        T _co_nume; // covariance/coherence numerator
        T _coh; // coherence
        float _amp2_m; // sum of amplitude square for image i
        float _amp2_j; // sum of amplitude aquare for image j
        int rslc_inx_m, rslc_inx_j;

        for (pair_i=0; pair_i < n_image_pair; pair_i++){
            m = image_pair_ref[pair_i];
            j = image_pair_sec[pair_i];
            
            _co_nume = T(0.0, 0.0);
            _amp2_m = 0.0;
            _amp2_j = 0.0;
            for (k = 0; k < az_win; k++) {
                for (l = 0; l < r_win; l++) {
                    sec_az = ref_az-az_half_win+k;
                    sec_r = ref_r-r_half_win+l;
                    if (pc_is_shp[i*win+k*r_win+l] && sec_az >= 0 && sec_az < nlines && sec_r >= 0 && sec_r < width) {
                        rslc_inx_m = (sec_az*width+sec_r)*nimages+m;
                        rslc_inx_j = (sec_az*width+sec_r)*nimages+j;
                        _amp2_m += norm(rslc[rslc_inx_m]);
                        _amp2_j += norm(rslc[rslc_inx_j]);
                        _co_nume += rslc[rslc_inx_m]*conj(rslc[rslc_inx_j]);
                    }
                }
            }
            _amp2_m = sqrt(_amp2_m*_amp2_j);
            _coh = _co_nume/_amp2_m;
            coh[i*n_image_pair+pair_i] = _coh;
        }
        ''',
        name = 'emperical_co_pc_kernel',reduce_dims = False,no_return=True
    )

# %% ../nbs/API/co.ipynb 22
@ngpjit
def _emperical_co_pc_numba(
    rslc,
    az_idx,
    r_idx,
    pc_is_shp,
    image_pairs,
):
    nlines, width, nimages = rslc.shape
    n_pc = az_idx.shape[0]
    az_win, r_win = pc_is_shp.shape[1:]
    az_half_win, r_half_win = az_win//2, r_win//2
    npairs = image_pairs.shape[0]

    coh = np.empty((n_pc, npairs),dtype=rslc.dtype)

    for i in prange(n_pc):
        for pair_i in range(npairs):
            m, j = image_pairs[pair_i]
            _co_nume = 0.0 + 0.0j
            _amp2_m = 0.0
            _amp2_j = 0.0
            for k in range(az_win):
                for l in range(r_win):
                    az_idx_ = az_idx[i] - az_half_win + k
                    r_idx_ = r_idx[i] - r_half_win + l
                    if (az_idx_ >= 0) and (az_idx_ < nlines) and (r_idx_ >= 0) and (r_idx_ < width) and pc_is_shp[i,k,l]:
                        rslc_m = rslc[az_idx_, r_idx_, m]
                        rslc_j = rslc[az_idx_, r_idx_, j]
                        _amp2_m += rslc_m.real**2 + rslc_m.imag**2
                        _amp2_j += rslc_j.real**2 + rslc_j.imag**2
                        _co_nume += rslc_m*np.conj(rslc_j)
            _coh = _co_nume/math.sqrt(_amp2_m*_amp2_j)
            coh[i,pair_i] = _coh
    return coh

# %% ../nbs/API/co.ipynb 23
def emperical_co_pc(rslc:np.ndarray, # rslc stack, dtype: `np.complex64`
                    idx:np.ndarray, # index of point target (azimuth_index, range_index), dtype: `np.int32`, shape: (n_pc, 2)
                    pc_is_shp:np.ndarray, # shp bool, dtype: `bool`
                    block_size:int=128, # the CUDA block size, it only affects the calculation speed
                    image_pairs:np.ndarray=None, # only coherence of those image pairs will estimated, dtype: `np.int32`, shape: (n_image_pair, 2)
                   )-> np.ndarray: # `coh`, dtype:`np.complex64`, shape(n_pc, n_image_pair)
    '''
    Maximum likelihood covariance estimator for point cloud data.
    '''
    xp = get_array_module(rslc)
    nlines, width, nimages = rslc.shape
    az_win, r_win = pc_is_shp.shape[-2:]
    az_half_win = (az_win-1)//2
    r_half_win = (r_win-1)//2
    idx = idx.astype(np.int32)
    az_idx = idx[:,0]; r_idx = idx[:,1]
    n_pc = az_idx.shape[0]

    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:
        return _emperical_co_pc_numba(rslc,az_idx,r_idx,pc_is_shp,image_pairs)

    else:
        image_pairs = cp.asarray(image_pairs)
        coh = cp.empty((n_pc,image_pairs.shape[0]),dtype=rslc.dtype)
        _emperical_co_pc_kernel(
            rslc, az_idx, r_idx, pc_is_shp, image_pairs[:,0], image_pairs[:,1],
            cp.int32(nlines),cp.int32(width),cp.int32(nimages),
            cp.int32(az_half_win),cp.int32(r_half_win),
            cp.int32(n_pc), cp.int32(image_pairs.shape[0]),
            coh,
            size = n_pc,block_size=block_size
        )
        return coh

# %% ../nbs/API/co.ipynb 36
@ngjit
def uncompress_single_coh_numba(coh, nimages, image_pairs):
    uncompressed_coh = np.zeros((nimages,nimages),dtype=coh.dtype)
    ref_images, sec_images = image_pairs[:,0], image_pairs[:,1]
    for i in range(ref_images.shape[0]):
        uncompressed_coh[ref_images[i], sec_images[i]] = coh[i]
        uncompressed_coh[sec_images[i], ref_images[i]] = np.conj(coh[i])
    for i in range(nimages):
        uncompressed_coh[i,i] = 1
    return uncompressed_coh

# %% ../nbs/API/co.ipynb 37
def uncompress_coh(
    coh:np.ndarray, # compressed coherence stack, dtype: `np.complex64`, shape(...,n_image_pair)
    image_pairs:np.ndarray=None, # image pairs, dtype: `np.int32`, shape: (n_image_pair, 2), all pairs by default
)-> np.ndarray: # uncompressed, dtype:`np.complex64`, shape(..., n_image_pair)
    '''
    uncompress coh matrix to a hermitian matrix
    '''
    xp = get_array_module(coh)
    if image_pairs is None:
        nimages = mr.nimage_from_npair(coh.shape[-1])
        image_pairs = mr.TempNet.from_bandwidth(nimages).image_pairs
    else:
        nimages = image_pairs[-1,-1]+1
    uncompressed_coh = xp.zeros((*coh.shape[:-1],nimages,nimages),dtype=coh.dtype)
    ref_images, sec_images = image_pairs[:,0], image_pairs[:,1]
    uncompressed_coh[...,ref_images,sec_images] = coh
    uncompressed_coh[...,sec_images,ref_images] = coh.conj()
    uncompressed_coh[...,np.arange(nimages),np.arange(nimages)] = 1
    return uncompressed_coh

# %% ../nbs/API/co.ipynb 40
@ngpjit
def _ad_intf_pc_numba(
    ref_rslc,
    sec_rslc,
    az_idx,
    r_idx,
    pc_is_shp,
):
    nlines, width = ref_rslc.shape
    n_pc = az_idx.shape[0]
    az_win, r_win = pc_is_shp.shape[1:]
    az_half_win, r_half_win = az_win//2, r_win//2

    inf = np.empty(n_pc,dtype=rslc.dtype)

    for i in prange(n_pc):
        _co_nume = 0.0 + 0.0j
        _ref_amp2 = 0.0
        _sec_amp2 = 0.0
        for k in range(az_win):
            for l in range(r_win):
                az_idx_ = az_idx[i] - az_half_win + k
                r_idx_ = r_idx[i] - r_half_win + l
                if (az_idx_ >= 0) and (az_idx_ < nlines) and (r_idx_ >= 0) and (r_idx_ < width) and pc_is_shp[i,k,l]:
                    _ref_rslc = ref_rslc[az_idx_, r_idx_]
                    _sec_rslc = sec_rslc[az_idx_, r_idx_]
                    _ref_amp2 += _ref_rslc.real**2 + _ref_rslc.imag**2
                    _sec_amp2 += _sec_rslc.real**2 + _sec_rslc.imag**2
                    _co_nume += _ref_rslc*np.conj(_sec_rslc)
        _inf = _co_nume/math.sqrt(_ref_amp2*_sec_amp2)
        inf[i] = _inf
    return inf

# %% ../nbs/API/co.ipynb 41
if is_cuda_available():
    _ad_intf_pc_kernel = cp.ElementwiseKernel(
        'raw T ref_rslc, raw T sec_rslc, raw I az_idx, raw I r_idx, raw bool pc_is_shp, int32 nlines, int32 width, int32 az_half_win, int32 r_half_win, int32 n_pc',
        'raw T intf',
        '''
        if (i >= n_pc) return;
        int az_win = 2*az_half_win+1;
        int r_win = 2*r_half_win+1;
        int win = az_win*r_win;
        
        int ref_az = az_idx[i];
        int ref_r = r_idx[i];
    
        int sec_az, sec_r;
    
        int k,l; // index of search window
        T _co_nume = T(0.0, 0.0); // covariance/coherence numerator
        T _intf; // coherence
        float _ref_amp2 = 0.0; // sum of amplitude square for ref image
        float _sec_amp2 = 0.0; // sum of amplitude aquare for sec image
        int _rslc_idx;
        for (k = 0; k < az_win; k++) {
            for (l = 0; l < r_win; l++) {
                sec_az = ref_az-az_half_win+k;
                sec_r = ref_r-r_half_win+l;
                if (pc_is_shp[i*win+k*r_win+l] && sec_az >= 0 && sec_az < nlines && sec_r >= 0 && sec_r < width) {
                    _rslc_idx = sec_az*width+sec_r;
                    _ref_amp2 += norm(ref_rslc[_rslc_idx]);
                    _sec_amp2 += norm(sec_rslc[_rslc_idx]);
                    _co_nume += ref_rslc[_rslc_idx]*conj(sec_rslc[_rslc_idx]);
                }
            }
        }
        _intf = _co_nume/sqrt(_ref_amp2*_sec_amp2);
        intf[i] = _intf;
        ''',
        name = 'ad_intf_pc_kernel',reduce_dims = False,no_return=True
    )

# %% ../nbs/API/co.ipynb 42
def ad_intf_pc(
    ref_rslc:np.ndarray, # reference rslc, dtype: `np.complex64`
    sec_rslc:np.ndarray, # secondary rslc, dtype: `np.complex64`
    idx:np.ndarray, # index of point target (azimuth_index, range_index), dtype: `np.int32`, shape: (n_pc, 2)
    pc_is_shp:np.ndarray, # shp bool, dtype: `bool`
    block_size:int=128, # the CUDA block size, it only affects the calculation speed
)-> np.ndarray: # `coh`, dtype:`np.complex64`, shape(n_pc, n_image_pair)
    '''
    Adaptive multilooking interferogram generation for point cloud data.
    '''
    xp = get_array_module(ref_rslc)
    if xp is np:
        return _ad_intf_pc_numba(ref_rslc, sec_rslc, idx[:,0], idx[:,1], pc_is_shp)
    else:
        nlines, width = ref_rslc.shape
        n_pc = idx.shape[0]
        az_win, r_win = pc_is_shp.shape[-2:]
        az_half_win = (az_win-1)//2
        r_half_win = (r_win-1)//2
        inf = cp.empty(n_pc,dtype=ref_rslc.dtype)
        _ad_intf_pc_kernel(
            ref_rslc, sec_rslc,
            idx[:,0], idx[:,1],
            pc_is_shp,
            cp.int32(nlines), cp.int32(width),
            cp.int32(az_half_win), cp.int32(r_half_win),
            cp.int32(n_pc),
            inf,
            size = n_pc, block_size=block_size
        )
        return inf

# %% ../nbs/API/co.ipynb 49
def isPD(co:np.ndarray, # absolute value of complex coherence/covariance stack
         )-> np.ndarray: # bool array indicating wheather coherence/covariance is positive define
    xp = get_array_module(co)
    L = xp.linalg.cholesky(co)
    is_PD = xp.isfinite(L).all(axis=(-2,-1))
    return is_PD

# %% ../nbs/API/co.ipynb 56
'''
    The method is presented in [1]. John D'Errico implented it in MATLAB [2] under BSD
    Licence and [3] implented it with Python/Numpy based on [2] also under BSD Licence.
    This is a cupy implentation with stack of matrix supported.

    [1] N.J. Higham, "Computing a nearest symmetric positive semidefinite
    matrix" (1988): https://doi.org/10.1016/0024-3795(88)90223-6
    
    [2] https://www.mathworks.com/matlabcentral/fileexchange/42885-nearestspd
    
    [3] https://gist.github.com/fasiha/fdb5cec2054e6f1c6ae35476045a0bbd
'''
def nearestPD(co:np.ndarray, # stack of matrix with shape [...,N,N]
             )-> np.ndarray: # nearest positive definite matrix of input, shape [...,N,N]
    """Find the nearest positive-definite matrix to input matrix."""
    xp = get_array_module(co)
    B = (co + xp.swapaxes(co,-1,-2))/2
    s, V = xp.linalg.svd(co)[1:]
    I = xp.eye(co.shape[-1],dtype=co.dtype)
    S = s[...,None]*I
    del s

    H = xp.matmul(xp.swapaxes(V,-1,-2), xp.matmul(S, V))
    del S, V
    A2 = (B + H) / 2
    del B, H
    A3 = (A2 + xp.swapaxes(A2,-1,-2))/2
    del A2

    if wherePD(A3).all():
        return A3
    
    co_norm = xp.linalg.norm(co,axis=(-2,-1))
    spacing = xp.nextafter(co_norm,co_norm+1.0)-co_norm
    
    k = 0
    while True:
        isPD = wherePD(A3)
        isPD_all = isPD.all()
        if isPD_all or k>=100:
            break
        k+=1
        mineig = xp.amin(xp.linalg.eigvalsh(A3),axis=-1)
        assert xp.isfinite(mineig).all()
        A3 += (~isPD[...,None,None] * I) * (-mineig * k**2 + spacing)[...,None,None]
    #print(k)
    return A3

# %% ../nbs/API/co.ipynb 58
def regularize_spectral(coh:np.ndarray, # stack of matrix with shape [...,N,N]
                        beta:Union[float, np.ndarray], # the regularization parameter, a float number or cupy ndarray with shape [...]
                        )-> np.ndarray: # regularized matrix, shape [...,N,N]
    '''
    Spectral regularizer for coherence matrix.
    '''
    xp = get_array_module(coh)
    I = xp.eye(coh.shape[-1],dtype=coh.dtype)
    beta = xp.asarray(beta)[...,None,None]

    regularized_coh = (1-beta)*coh + beta* I
    return regularized_coh
