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

# %% auto 0
__all__ = ['emi', 'ds_temp_coh']

# %% ../nbs/API/pl.ipynb 5
import numpy as np
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/pl.ipynb 8
if is_cuda_available():
    def _emi_cp(
        coh, # complex coherence metrix, dtype, cp.complex64
        ref:int=0,
    ):
        n_points = coh.shape[0]
        max_batch_size = 2**20
        num_batchs = np.ceil(n_points/max_batch_size).astype(int)
        ph = cp.empty(coh.shape[:-1],dtype=coh.dtype)
        emi_quality = cp.empty(coh.shape[:-2], dtype=cp.float32)
        for i in range(num_batchs):
            start = i*max_batch_size
            end = (i+1)*max_batch_size
            if end >= coh.shape[0]: end = coh.shape[0]
            _coh = coh[start:end]
            coh_mag = cp.abs(_coh)
            coh_mag_inv = cp.linalg.inv(coh_mag)
            min_eigval, min_eig = cp.linalg.eigh(coh_mag_inv*_coh)
            min_eigval = min_eigval[...,0]
            min_eig = min_eig[...,0]*min_eig[...,[ref],0].conj()
            ph[start:end] = min_eig/abs(min_eig)
            emi_quality[start:end] = min_eigval
        return ph, emi_quality

# %% ../nbs/API/pl.ipynb 9
@ngpjit
def _emi_numba(
    coh,
    ref:int=0,
):
    n_points = coh.shape[0]
    ph = np.empty(coh.shape[:-1],dtype=coh.dtype)
    emi_quality = np.empty(coh.shape[:-2],dtype=np.float32)
    for i in prange(n_points):
        _coh = coh[i]
        coh_mag = np.abs(_coh)
        coh_mag_inv = np.linalg.inv(coh_mag)
        min_eigval, min_eig = np.linalg.eigh(coh_mag_inv*_coh)
        min_eigval = min_eigval[0]
        min_eig = min_eig[:,0]*np.conj(min_eig[ref,0])
        for j in range(ph.shape[-1]):
            ph[i,j] = min_eig[j]/abs(min_eig[j])
        emi_quality[i] = min_eigval
    return ph, emi_quality

# %% ../nbs/API/pl.ipynb 10
def emi(coh:np.ndarray, #complex coherence metrix,dtype cupy.complex
        ref:int=0, #index of reference image in the phase history output, optional. Default: 0
       )-> tuple[np.ndarray,np.ndarray]: # estimated phase history `ph`, dtype complex; quality (minimum eigvalue, dtype float)
    xp = get_array_module(coh)
    if xp is np:
        return _emi_numba(coh,ref)
    else:
        return _emi_cp(coh,ref)

# %% ../nbs/API/pl.ipynb 22
@ngpjit
def _ds_temp_coh_numba(
    coh:np.ndarray,# complex coherence metrix, dtype cupy.complex
    ph = np.ndarray, # complex phase history, dtype cupy.complex
):
    nimages = ph.shape[-1]
    n_points = ph.shape[0]
    temp_coh = np.empty(n_points,dtype=np.float32)
    for i in prange(n_points):
        _ph = ph[i]
        _coh = coh[i]
        for n in range(nimages):
            _ph[n] = _ph[n]/np.abs(_ph[n])
        for n in range(nimages):
            for k in range(nimages):
                _coh[n,k] = _coh[n,k]/np.abs(_coh[n,k])

        _t_coh = 0
        for n in range(nimages):
            for k in range(nimages):
                if k == n: continue
                int_conj_ph = np.conjugate(_ph[n])*_ph[k]
                diff_ph = _coh[n,k]*int_conj_ph
                _t_coh += diff_ph.real
        _t_coh /= nimages**2-nimages
        temp_coh[i] = _t_coh
    return temp_coh

# %% ../nbs/API/pl.ipynb 23
if is_cuda_available():
    _ds_temp_coh_kernel = cp.ElementwiseKernel(
        'raw T coh, raw T ph, int32 n_points, int32 nimages',
        'raw float32 temp_coh',
        '''
        if (i >= n_points) return;
        float _t_coh = 0;
        int n; int k;
        for (n = 0; n < nimages; n++){
            for (k = 0; k < nimages; k++){
                if (k == n) continue;
                int coh_idx = i*nimages*nimages+n*nimages+k;
                int ph_n_idx = i*nimages+n;
                int ph_k_idx = i*nimages+k;
                _t_coh += real(coh[coh_idx]/sqrt(norm(coh[coh_idx]))*conj(ph[ph_n_idx])/sqrt(norm(ph[ph_n_idx]))*ph[ph_k_idx]/sqrt(norm(ph[ph_k_idx])));
            }
        }
        temp_coh[i] = _t_coh/(nimages*(nimages-1));
        ''',
        name = 'ds_temp_coh_kernel',no_return=True,
    )

# %% ../nbs/API/pl.ipynb 24
def ds_temp_coh(coh:np.ndarray,# complex coherence metrix, np.complex64 or cp.complex64
                ph = np.ndarray, # complex phase history, np.complex64 or cp.complex64
                block_size:int=128, # the CUDA block size, only applied for cuda
            ):
    xp = get_array_module(coh)
    n_points = ph.shape[0]
    nimages = ph.shape[-1]
    assert coh.shape[-2:] == (nimages,nimages), "input dimension not match"
    if xp is np:
        return _ds_temp_coh_numba(coh,ph)
    else:
        temp_coh = cp.empty(n_points, dtype=cp.float32)
        _ds_temp_coh_kernel(coh, ph, cp.int32(n_points),cp.int32(nimages), temp_coh, size=n_points, block_size=block_size)
        return temp_coh
