"""`mri` module contains useful functions and classes for MRI forward and adjoint operations."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/P4_HPD_cg/94_mri.ipynb.

# %% auto 0
__all__ = ['kb_func', 'kb_kernel', 'calc_deapodizer', 'BaseOperator', 'NuFFT', 'hanning_window', 'coils_norm', 'coil_sens',
           'prewhiten', 'dcf', 'MriImagingModel', 'MriOperator', 'init_params', 'setup_params']

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 2
from nbdev.showdoc import *
from fastcore.test import *
from fastcore.utils import *

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 3
import h5py
import numpy as np
import scipy
import skimage
from matplotlib import pyplot as plt
import time
from abc import ABC, abstractmethod
import itertools

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 4
from .plot import *
from .metrics import *
from .math import *
from .io import *
import cgsense2023 as cgs2003

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 6
def kb_func(
  u:np.ndarray,   # kernel radius 
  width:int,      # kernel width 
  beta:float,     # scale params (Jackson '91.)
) -> np.ndarray:  # Kaiser-Bessel window function
    "Kaiser-Bessel window function y = I0 [beta*sqrt(1-(2u/w)^2)]/w"
    assert np.size(width) == 1
    x = beta * np.sqrt(1 - (2 * u / width) ** 2)
    y = scipy.special.i0(x) / width
    return y

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 8
def kb_kernel(
    kwidth:int,     # kernel width
    overgrid:float, # over-sampled factor 
    gridsize:int,   # grid size of over-sampled grid
    klength:int=32, # kernal look-up-table (lut) length
    ) -> tuple:     # return (kernel values and FT of kernel)
    "Calculate the appropriate Kaiser-Bessel gridding kernel and its FFT for deapodization."
    # From Beatty et al. equation [5]
    beta = np.pi*np.sqrt((kwidth/overgrid)**2 * (overgrid - 0.5) ** 2 - 0.8)
    # Kernel radii.
    u = np.linspace(0, (kwidth/2), int(np.ceil(klength*kwidth/2)))
    kvs = kb_func(u, kwidth, beta)
    kvs /= kvs.max()

    fft_y = np.flip(kvs) #TODO
    if np.mod(kwidth, 2): fft_y = np.concatenate((fft_y[:-1], kvs))
    else: fft_y = np.concatenate((fft_y, kvs))

    fft_y = np.abs(np.fft.fftshift(np.fft.ifft(fft_y, gridsize*klength))*fft_y.size)
    x = np.linspace(-int((gridsize//(2*overgrid))),int((gridsize//(2*overgrid)))-1,int((gridsize//(overgrid))))

    fft_y = fft_y[(fft_y.size / 2 + x).astype(int)]
    h = np.sinc(x / (gridsize * klength)) ** 2

    kvs_fft = fft_y * h
    kvs_fft /= kvs_fft.max()
    return kvs, kvs_fft

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 10
def calc_deapodizer(
    ktable_fft:np.ndarray,  # FFT of kaiser bessel func
) -> np.ndarray:            # deapodization func
    "calculate deapodizatoin func based on FFT of kaiser bessel func."
    deapodizer1D = 1.0 /ktable_fft
    deapodizer2D = np.outer(deapodizer1D, deapodizer1D)
    return deapodizer2D

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 17
class BaseOperator(ABC):
    "Abstract Base Class (ABC) for linear Operators."

    def __init__(self, data_par):
        "data_par (dict):(image_dim),(num_coils),(num_reads),(num_projs)"
        self.image_dim = data_par["image_dim"]
        self.grid_size = data_par["grid_size"]
        self.num_reads = data_par["num_reads"]
        self.num_coils = data_par["num_coils"]
        self.num_projs = data_par["num_projs"]

    @abstractmethod
    def forward(self, image):
        raise NotImplementedError

    @abstractmethod
    def adjoint(self, kdata):
        raise NotImplementedError

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 20
class NuFFT(BaseOperator):
    "Non-uniform FFT (NuFFT) object"
    def __init__(self, par, trajectory, fft_dim=(-2, -1)):
        super().__init__(par["Data"])
        
        self.trajectory = trajectory
        self.fft_dim = fft_dim
        self.overgrid = par["Data"]["overgrid"]
        self.klength = par["FFT"]["klength"]
        self.kwidth = par["FFT"]["kwidth"]
    
        self.ktable, ktable_FFT = kb_kernel(self.kwidth, self.overgrid, self.grid_size, self.klength)
        self.deapodizer = calc_deapodizer(ktable_FFT.astype(np.float64))
        
        self.num_kpoints = self.ktable.size    
        self.kwidth2 = par["FFT"]["kwidth"] / 2 #kernel radius i.e., half of the width
        
        if "density_cor" in par["FFT"].keys():
            self.density_cor = par["FFT"]["density_cor"]
        else:
            self.density_cor = np.ones(trajectory.shape[:-1], dtype=np.float64)
        
        if "grid_matrix" in par["FFT"].keys():
            self.grid_matrix = par["FFT"]["grid_matrix"]
        else:
            self.grid_matrix = self._generate_grid_matrix()
        
        self.grid_matrix_adj = self.grid_matrix.transpose()

    def adjoint(self, kdata):
        # "sqrt" density compensation
        denscor_inp = kdata * self.density_cor
        # gridding
        ogkspace = np.zeros((self.num_coils, self.grid_size, self.grid_size),dtype=np.complex64)
        for nc in range(self.num_coils):
            ogkspace[nc] = self.grid_matrix_adj.dot(denscor_inp[nc].flatten()).reshape(self.grid_size, self.grid_size)
        # iFFT
        image_grid = kspace2im_ifft2(ogkspace)
        # deapodization
        return self._deapodizer_adj(image_grid)

    def forward(self, image):
        # deapodization
        image_deapodizer = self._deapodizer_fwd(image)
        # FFT
        ogkspace = im2kspace_fft2(image_deapodizer)
        # Resample on radial trajectories
        kspace = np.zeros((self.num_coils, self.num_projs, self.num_reads),dtype=np.complex64)
        for nc in range(self.num_coils):
            kspace[nc] = self.grid_matrix.dot(ogkspace[nc].flatten()).reshape(self.num_projs, self.num_reads)
        # "sqrt" density compensation
        return kspace*self.density_cor

    def _deapodizer_adj(self, image_grid):
        # crop the grid_size_image to desired size image then multiply with deapodizer
        return crop_im(image_grid,(self.image_dim, self.image_dim))*self.deapodizer

    def _deapodizer_fwd(self, image):
        # multiply the desired image with deapodizer then zeropad to the grid_size_image 
        return zero_pad(image * self.deapodizer, (self.grid_size, self.grid_size))

    def _grid_lut(self, s, return_mapping=False):
        gridcenter = self.grid_size // 2
        sg = np.zeros((self.num_coils,self.grid_size,self.grid_size),dtype=np.complex64)
        grid_point_mapping = []
        kdat = s * self.density_cor
        for iproj, iread in itertools.product(range(self.num_projs),range(self.num_reads)):
            temp_mapping = []
            temp_mapping.append((iread, iproj))
            kx = self.trajectory[iproj, iread, 1]*self.overgrid
            ky = self.trajectory[iproj, iread, 0]*self.overgrid

            ixmin = int((kx - self.kwidth2) + gridcenter)
            ixmax = int((kx + self.kwidth2) + gridcenter) + 1
            iymin = int((ky - self.kwidth2) + gridcenter)
            iymax = int((ky + self.kwidth2) + gridcenter) + 1

            for gcount1 in np.arange(ixmin, ixmax+1):
                dkx = (gcount1 - gridcenter) - kx

                for gcount2 in np.arange(iymin, iymax+1):
                    dky = (gcount2 - gridcenter) - ky
                    dk = np.sqrt(dkx ** 2 + dky ** 2)

                    if dk < self.kwidth2:
                        fracind = dk / self.kwidth2 * (self.num_kpoints - 1)
                        kidx = int(fracind)
                        fracdk = fracind - kidx

                        kern = (self.ktable[kidx] * (1 - fracdk) + self.ktable[kidx + 1] * fracdk)

                        indx = gcount1
                        indy = gcount2

                        indx = np.mod(indx, self.grid_size)
                        indy = np.mod(indy, self.grid_size)

                        temp_mapping.append((indx, indy))

                        sg[:, indy, indx] += (kern * kdat[:,iproj,iread])
            
            grid_point_mapping.append(temp_mapping)

        if return_mapping:
            return sg, grid_point_mapping
        else:
            return sg

    def _invgrid_lut(self, sg):
        gridcenter = self.grid_size / 2

        s = np.zeros((self.num_coils,self.num_projs,self.num_reads),dtype=np.complex64)

        for iproj, ismpl in itertools.product(range(self.num_projs),range(self.num_reads)):

            kx = self.trajectory[iproj, ismpl, 1]*self.overgrid
            ky = self.trajectory[iproj, ismpl, 0]*self.overgrid

            ixmin = int((kx - self.kwidth2) + gridcenter)
            ixmax = int((kx + self.kwidth2) + gridcenter) + 1
            iymin = int((ky - self.kwidth2) + gridcenter)
            iymax = int((ky + self.kwidth2) + gridcenter) + 1

            for gcount1 in np.arange(ixmin, ixmax+1):
                dkx = (gcount1 - gridcenter) - kx

                for gcount2 in np.arange(iymin, iymax+1):
                    dky = (gcount2 - gridcenter) - ky
                    dk = np.sqrt(dkx ** 2 + dky ** 2)
                    if dk < self.kwidth2:

                        fracind = dk / self.kwidth2 * (self.num_kpoints - 1)
                        kidx = int(fracind)
                        fracdk = fracind - kidx

                        kern = (self.ktable[kidx] * (1 - fracdk) + self.ktable[kidx + 1] * fracdk)
                        indx = gcount1
                        indy = gcount2

                        if gcount1 < 0:
                            indx += self.grid_size
                            indy = self.grid_size - indy

                        if gcount1 >= self.grid_size:
                            indx -= self.grid_size
                            indy = self.grid_size - indy

                        if gcount2 < 0:
                            indy += self.grid_size
                            indx = self.grid_size - indx

                        if gcount2 >= self.grid_size:
                            indy -= self.grid_size
                            indx = self.grid_size - indx

                        s[:, iproj, ismpl] += kern*sg[:, indy, indx]
        return s * self.density_cor

    def _generate_grid_matrix(self):
        gridcenter = self.grid_size / 2
        rowind, colind, value = [], [], []
        for iproj, ismpl in itertools.product(range(self.num_projs),range(self.num_reads)):

            kx = self.trajectory[iproj, ismpl, 1]*self.overgrid
            ky = self.trajectory[iproj, ismpl, 0]*self.overgrid

            ixmin = int((kx - self.kwidth2) + gridcenter)
            ixmax = int((kx + self.kwidth2) + gridcenter) + 1
            iymin = int((ky - self.kwidth2) + gridcenter)
            iymax = int((ky + self.kwidth2) + gridcenter) + 1

            for gcount1 in np.arange(ixmin, ixmax+1):
                dkx = (gcount1 - gridcenter) - kx
                for gcount2 in np.arange(iymin, iymax+1):
                    dky = (gcount2 - gridcenter) - ky
                    dk = np.sqrt(dkx ** 2 + dky ** 2)
                    if dk < self.kwidth2:
                        fracind = (dk / self.kwidth2 * (self.num_kpoints - 1))
                        kidx = int(fracind)
                        fracdk = fracind - kidx
                        kern = (self.ktable[kidx] * (1 - fracdk) + self.ktable[kidx + 1] * fracdk)
                        indx = gcount1
                        indy = gcount2

                        if gcount1 < 0:
                            indx += self.grid_size
                            indy = self.grid_size - indy
                        if gcount1 >= self.grid_size:
                            indx -= self.grid_size
                            indy = self.grid_size - indy
                        if gcount2 < 0:
                            indy += self.grid_size
                            indx = self.grid_size - indx
                        if gcount2 >= self.grid_size:
                            indy -= self.grid_size
                            indx = self.grid_size - indx

                        rowind.append(iproj*self.num_reads+ismpl)
                        colind.append(indy*self.grid_size+indx)
                        value.append(kern)

        gridmat = scipy.sparse.coo_matrix((value, (rowind, colind)), shape=(self.num_projs*self.num_reads, self.grid_size**2))
        gridmat = gridmat.tocsc()
        return gridmat

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 27
def hanning_window(
    N:int=64,           # window size
) -> np.ndarray:    # window function values
    "1D symetric hanning window function"
    return 0.5*(1 - np.cos(2*np.pi*np.arange(N)/(N - 1)))

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 30
def coils_norm(
    coils:np.ndarray,   # raw coil sensitivities 
) -> tuple:             # tuple of (normalized coils, intensity_cor)
    "coil sensitivities normalization"
    # sosqrt = np.sqrt(np.sum(np.abs(coils)**2, axis=0)) # np.sum(np.sqrt((coils * np.conj(coils))),0)
    sosqrt = np.sum(np.sqrt((coils * np.conj(coils))),0)
    sosqrt[sosqrt == 0] = 1 
    intensity_cor = sosqrt # intensity correction
    return coils/sosqrt, intensity_cor

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 32
def coil_sens(
    kdata:np.ndarray,   # raw-uncompensated kspace
    traj:np.ndarray,    # kspace trajectories
    par:dict,           # dictionary contains paramss
) -> tuple:             # tuple of (coils, mask, intensity_cor)
    "Estimate complex coil sensitivities using SOS"
    FFT = NuFFT(par=par, trajectory=traj)
    wsize = 64 # size of selected low-res portion of kdata
    nspokes, nreads = kdata.shape[-2:]
    kdata_c = crop_im(kdata, (nspokes, wsize))
    hwin = hanning_window(N=wsize)
    kdata_cw = kdata_c*hwin
    kdata_w = zero_pad(kdata_cw, (nspokes,nreads))
    coil_images = FFT.adjoint(kdata_w * par["FFT"]["density_cor"])
    rss_image = rss_rec(coil_images)
    coils = coil_images/rss_image
    thresh = skimage.filters.threshold_otsu(rss_image)
    mask = rss_image > thresh*0.3
    mask = scipy.ndimage.binary_dilation(mask, iterations=10)
    coils, intensity_cor = coils_norm(coils)
    return coils, mask, intensity_cor

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 35
def prewhiten(
    kdata:np.ndarray,       # kspace data (Nc,Nsp,Nr)
    noise:np.ndarray,       # noises (Nc,Nsp,Nr)
    coils:np.ndarray=None,  # sensitivity maps (Nc,Nsp,Nr)
) -> tuple:                 # pre-whitened (kspace,coils)
    "Noise pre-whitenning"
    
    cov = np.cov(noise.reshape(noise.shape[0], -1))
    L = np.linalg.cholesky(cov)
    
    kdata_rsh = kdata.reshape([kdata.shape[0], -1])
    kdata_pwh = np.reshape(np.linalg.solve(L, kdata_rsh), kdata.shape)

    if coils is not None:
        coils_rsh = coils.reshape(coils.shape[0], -1)
        coils_pwh = np.reshape(np.linalg.solve(L, coils_rsh), coils.shape)
    
    return kdata_pwh, coils_pwh

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 38
def dcf(
    grid_matrix:np.ndarray, # gridding matrix
    num_projs:int=96,       # number of radial spokes
    num_reads:int=512,      # number of readout samples
) -> np.ndarray:            # density compensation function
    "density compensation function (DCF)"
    dens = grid_matrix.transpose()@(np.ones(grid_matrix.shape[0]))
    dens = 1/(grid_matrix@dens)
    dens = np.reshape(dens,(num_projs,num_reads))
    # sqrt because of NUFFT's Forward and Adjoint are adjoint with each other
    return np.sqrt(dens.astype(np.float32)) 

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 41
class MriImagingModel(BaseOperator):
    "The MRI imaging model including Coils."

    def __init__(self, par):
        super().__init__(par["Data"])

        self.NuFFT = NuFFT(par=par, trajectory=par["Data"]["trajectory"])
        self.coils = par["Data"]["coils"]
        self.conj_coils = np.conj(self.coils)

    def adjoint(self, kdata):
        return np.sum(self.NuFFT.adjoint(kdata) * self.conj_coils, 0)

    def forward(self, image):
        return self.NuFFT.forward(image * self.coils)

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 46
class MriOperator:
    "MRI Operator Class"
    def __init__(self, data_par, optimizer_par):
        self.image_dim = data_par["image_dim"]
        self.num_coils = data_par["num_coils"]
        self.mask = data_par["mask"]
        self.do_incor = data_par["do_intensity_cor"]
        self.incor = data_par["intensity_cor"].astype(np.complex64)
        self.maxit = optimizer_par["max_iter"]
        self.lambd = optimizer_par["lambda"]
        self.tol = optimizer_par["tolerance"]
        self.operator = None

    def set_operator(self, op):
        "Mandatory to set an operator (MRImagingOperator) prior to minimization."
        self.operator = op

    def operator_lhs(self, image):
        "Compute the LHS (A^T A x) of the normal equation."
        assert self.operator is not None, \
            "Please set an operator with the set_operation method"
        return self.operator_rhs(self.operator.forward(image))

    def operator_rhs(self, kspace):
        "Compute the RHS (A^T b) of the normal equation."
        assert self.operator is not None, \
            "Please set an operator with the set_operation method"
        return self.operator.adjoint(kspace)

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 56
def init_params()->dict: 
    "initialize dictionary of parameters"
    params = {}
    params['Optimizer'] = {} 
    params['FFT'] = {} 
    params['Data'] = {} 
    return params

# %% ../nbs/P4_HPD_cg/94_mri.ipynb 58
def setup_params(
    fpath:str,  # path to the data file
    R:int,      # acceleration factor (1,2,3,...)
)->dict:        # return dictionary contains all params and Data
    "One-stop-shop for setting up Parameters and Data"
    # Read input data
    kdata, trajectory, InScale, coils = read_data(fpath, R)
    noise = None
    
    # extract parameters    
    ncoils, nspokes, nreads = kdata.shape

    params = init_params()
    
    # Parameters: Optimizer
    params['Optimizer']['tolerance'] = 0
    params['Optimizer']['lambda'] = 0
    params['Optimizer']['max_iter'] = 10
    
    # Parameters: FFT
    params['FFT']['kwidth'] = 5
    params['FFT']['klength'] = 10_000
    
    # Parameters: Data
    params['Data']['noise'] = noise
    params['Data']['do_intensity_cor'] = True
    params['Data']['do_density_cor'] = True
    params['Data']['overgrid'] = 1.7033398310591292
    params['Data']['image_dim'] = 300
    params['Data']['mask'] = 1
    params['Data']['num_coils'] = ncoils
    params['Data']['num_reads'] = nreads
    params['Data']['num_projs'] = nspokes
    params['Data']['grid_size'] = 512
    params['Data']['intensity_cor'] = np.array([1.0])
    
    
    FFT = NuFFT(par=params, trajectory=trajectory)
    grid_matrix = FFT.grid_matrix
    params['FFT']['grid_matrix'] = grid_matrix
    density_cor = dcf(grid_matrix, num_projs=nspokes, num_reads=nreads)
    params['FFT']['density_cor'] = density_cor
    
    ## calculated mask and intensity correction - similar results except the scaling differences
    # _, mask, intensity_cor = coil_sens(kdata, trajectory, params)
    # params['Data']['intensity_cor'] = intensity_cor #np.array([1.0])
    # params['Data']['mask'] = mask
    
    # noise pre-whitening if noise is not None
    noise = None
    if noise is not None:
        kdata, coils = prewhiten(kdata, noise, coils)
    
    # coil sensitivity estimation if coils = None
    # coils = None
    if coils is None:
        coils, mask, intensity_cor = coil_sens(kdata, trajectory, params)
            
    params['Data']['coils'] = coils    
    # coils = params['Data']['coils']
    params['Data']['rawdata'] = kdata
    params['Data']['trajectory'] = trajectory
    params['Data']['rawdata_density_cor'] = kdata * density_cor
    return params
