"""`math` module contains useful functions for image operations or mathematical calculations in general."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/P4_HPD_cg/92_math.ipynb.

# %% auto 0
__all__ = ['epsilon', 'im2kspace_fft2', 'kspace2im_ifft2', 'rss_rec', 'crop_im', 'zero_pad', 'radial_mask', 'kspace_filter']

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

# %% ../nbs/P4_HPD_cg/92_math.ipynb 4
import numpy as np
from prettytable import PrettyTable

# %% ../nbs/P4_HPD_cg/92_math.ipynb 5
from .plot import *
from .metrics import *

# %% ../nbs/P4_HPD_cg/92_math.ipynb 6
from skimage.metrics import peak_signal_noise_ratio, structural_similarity

# %% ../nbs/P4_HPD_cg/92_math.ipynb 7
epsilon = 1e-10

# %% ../nbs/P4_HPD_cg/92_math.ipynb 8
def im2kspace_fft2(
    img : np.ndarray,   # input image  
    dims=(-2,-1),       # dims for FFT operations
    norm="ortho",       # normalization option
) -> np.ndarray:        # centered-2D FFT of input images
    "Centered-2D-FFT converting images to kspace data"
    return np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(img, dims),norm=norm),dims)

# %% ../nbs/P4_HPD_cg/92_math.ipynb 15
def kspace2im_ifft2(
    ksp : np.ndarray,   # input ksapce data  
    dims=(-2,-1),       # dims for iFFT operations
    norm="ortho",       # normalization option
) -> np.ndarray:        # centered-2D iFFT of input kspace data
    "Centered-2D-iFFT converting kspace data to images"
    return np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(ksp, dims),norm=norm),dims)

# %% ../nbs/P4_HPD_cg/92_math.ipynb 22
def rss_rec(
    ims: np.ndarray,    # coil-by-coil input images
    dim: int=(0),       # axis for RSOS operation
)-> np.ndarray:         # RSOS (root sum-of-squares coil-combined image)
    "root sum-of-squares (RSOS) coil-combined reconstruction"
    return np.array(np.sqrt(np.sum(np.abs(ims)**2, axis=dim)/ims.shape[dim]))

# %% ../nbs/P4_HPD_cg/92_math.ipynb 25
def crop_im(
    ims:np.ndarray,     # input images
    sizes:int=(300,300), # new number of rows,cols
) -> np.ndarray:        # cropped images to (new_row, new_col) size
    "crop images to desired sizes (new_row, new_col)"
    row, col = ims.shape[-2:]   # original dimensions
    nrow, ncol = sizes          # new dimentions
    if nrow > row or ncol > col:
        print(f"Warning: cropped sizes {nrow,ncol} are larger than original sizes {row,col}!")
        return ims
    row1, col1 = (row - nrow)//2, (col - ncol)//2 # starts
    row2, col2 = (row + nrow)//2, (col + ncol)//2 # stops
    return ims[...,row1:row2, col1:col2]

# %% ../nbs/P4_HPD_cg/92_math.ipynb 37
def zero_pad(
    ims:np.ndarray,     # input images
    sizes:int=(512,512),# new number of rows,cols
) -> np.ndarray:        # zero-padded images to (new_row, new_col) size
    "zero-pad images to desired sizes (new_row, new_col)"
    row, col = ims.shape[-2:]   # original dimensions
    nrow, ncol = sizes          # new dimentions
    if nrow < row or ncol < col:
        print(f"Warning: padded sizes {nrow,ncol} are smaller than original sizes {row,col}!")
        return ims
    ims_out = np.zeros(ims.shape[:-2]+sizes, dtype=ims.dtype)
    row1, col1 = (nrow - row)//2, (ncol - col)//2 # starts
    row2, col2 = (row + nrow)//2, (col + ncol)//2 # stops
    ims_out[...,row1:row2, col1:col2] = ims
    return ims_out

# %% ../nbs/P4_HPD_cg/92_math.ipynb 46
def radial_mask(
    Nk:int=256,     # kspace dimension assumed Nkx=Nky=Nk
) -> np.ndarray:    # return radially-binary mask and center point    
    "create radially-binary mask for kspace filltering operation"
    kpoints = np.linspace(-0.5, 0.5, Nk, endpoint=False)
    xx, yy = np.meshgrid(kpoints, kpoints)
    mask = np.sqrt(xx**2 + yy**2)
    mask = np.where(np.abs(mask) > 0.5, 1, 0)
    mask = ~mask.astype(bool)
    center = Nk // 2
    return mask, center

# %% ../nbs/P4_HPD_cg/92_math.ipynb 52
def kspace_filter(
    x:np.ndarray,   # input images
    nreads:int,     # number of readout points
    image_dim:int,  # output image dimension
)->np.ndarray:      # kspace-filtered output images
    "Kspace filtering by zeroing out points outside the acquired trajectory"
    
    # data_type = mriop.operator.DTYPE
    # Nk, Nx = mriop.operator.num_reads, mriop.image_dim
    if x.ndim == 2: x = x[np.newaxis]
    xout = np.zeros_like(x)
    for idx in range(x.shape[0]):
        kspace0 = np.zeros((nreads,nreads),dtype=np.complex128)
        gridmask, gridcenter = radial_mask(nreads)
        img0 = np.zeros_like(kspace0) #
        Nx_range =slice(int(gridcenter-image_dim//2),int(gridcenter+image_dim//2))
        Ny_range = Nx_range
        img0[...,Nx_range, Ny_range] = x[idx]
        kspace1 = im2kspace_fft2(img0)
        masked_kspace1 = kspace1*gridmask
        img1 = kspace2im_ifft2(masked_kspace1)
        xout[idx] = img1[...,Nx_range, Ny_range]
    return np.squeeze(xout)
