''' This file contains the recipe for data preprocessing used to generate the combined coil images from the fastmri multicoil brain and knee datasets.
    These combined coil images were then used to train the autoencoder. The combined coil images are generated by combining the coil images using
    the sensitivity maps calculated with bart. To run this recipe, the bart toolbox needs to be installed and then follow the steps outlined in
    the preprocess_recipe function.'''


# bart toolbox installation instructions - https://mrirecon.github.io/bart/installation.html
_BART_TOOLBOX_PATH = ''

import numpy as np
import h5py
from tqdm import tqdm
import sys, os

os.environ["TOOLBOX_PATH"] = _BART_TOOLBOX_PATH
sys.path.append(os.path.join(_BART_TOOLBOX_PATH, 'python'))
from bart import bart
os.environ["OMP_NUM_THREADS"] = "1"

def fftc(input,  axes=None, norm='ortho'):
    """
    Perform a Fast Fourier Transform on the input array.

    Parameters:
    input (numpy.ndarray): The input array to transform.
    axes (tuple, optional): Axes over which to compute the FFT. If not specified, compute over all axes.
    norm (str, optional): Normalization mode. Default is 'ortho' for orthonormal transform.

    Returns:
    numpy.ndarray: The transformed output array.
    """
    tmp = np.fft.ifftshift(input, axes=axes)
    tmp = np.fft.fftn(tmp, axes=axes, norm=norm)
    output = np.fft.fftshift(tmp, axes=axes)
    return output

def ifftc(input, axes=None, norm='ortho'):
    """
    Perform an Inverse Fast Fourier Transform on the input array.

    Parameters:
    input (numpy.ndarray): The input array to transform.
    axes (tuple, optional): Axes over which to compute the inverse FFT. If not specified, compute over all axes.
    norm (str, optional): Normalization mode. Default is 'ortho' for orthonormal transform.

    Returns:
    numpy.ndarray: The transformed output array.
    """
    tmp = np.fft.ifftshift(input, axes=axes)
    tmp = np.fft.ifftn(tmp, axes=axes, norm=norm)
    output = np.fft.fftshift(tmp, axes=axes)
    return output

def adjoint(ksp, maps, mask):
    """
    Perform the adjoint operation on k-space data with coil sensitivity maps and a mask.

    Parameters:
    ksp (numpy.ndarray): The input k-space data, shape: [1, C, H, W].
    maps (numpy.ndarray): The coil sensitivity maps, shape: [1, C, H, W].
    mask (numpy.ndarray): The mask to apply on the k-space data, shape: [1, 1, H, W].

    Returns:
    numpy.ndarray: The output image after applying the adjoint operation, shape: [1, 1, H, W].
    """
    masked_ksp = ksp*mask
    coil_imgs = ifftc(masked_ksp,axes=(-2,-1))
    img_out = np.sum(coil_imgs*np.conj(maps),axis=1)[:,None,...]
    return img_out

def _expand_shapes(*shapes):
    """
    Expand the dimensions of the given shapes to match the maximum dimension.

    This function prepends 1s to the shapes with fewer dimensions to match the maximum number of dimensions.

    Parameters:
    *shapes (tuple): A variable length tuple containing shapes (as lists or tuples of integers).

    Returns:
    tuple: A tuple of expanded shapes, where each shape is a list of integers.
    """

    shapes = [list(shape) for shape in shapes]
    max_ndim = max(len(shape) for shape in shapes)
    shapes_exp = [[1] * (max_ndim - len(shape)) + shape
                  for shape in shapes]

    return tuple(shapes_exp)

def resize(input, oshape, ishift=None, oshift=None):
    """
    Resize with zero-padding or cropping.

    Parameters:
    input (array): Input array.
    oshape (tuple of ints): Output shape.
    ishift (None or tuple of ints): Input shift.
    oshift (None or tuple of ints): Output shift.

    Returns:
    array: Zero-padded or cropped result.
    """

    ishape1, oshape1 = _expand_shapes(input.shape, oshape)

    if ishape1 == oshape1:
        return input.reshape(oshape)

    if ishift is None:
        ishift = [max(i // 2 - o // 2, 0) for i, o in zip(ishape1, oshape1)]

    if oshift is None:
        oshift = [max(o // 2 - i // 2, 0) for i, o in zip(ishape1, oshape1)]

    copy_shape = [min(i - si, o - so)
                  for i, si, o, so in zip(ishape1, ishift, oshape1, oshift)]
    islice = tuple([slice(si, si + c) for si, c in zip(ishift, copy_shape)])
    oslice = tuple([slice(so, so + c) for so, c in zip(oshift, copy_shape)])

    output = np.zeros(oshape1, dtype=input.dtype)
    input = input.reshape(ishape1)
    output[oslice] = input[islice]

    return output.reshape(oshape)

def shape_data(ksp, final_res):
    """
    Reshape coil k-space data to output coil images with isotropic pixels and correct FOV = origional image width and the correct square image size given by "final_res".

    This function assumes that the k-space data has already been padded to make the corresponding images have isotropic pixels.

    Parameters:
    ksp (numpy.ndarray): The input coil k-space data, shape: [S, C, H, W].
    final_res (int): The final resolution for the output image.

    Returns:
    numpy.ndarray: The output image after reshaping, shape: [S, C, final_res, final_res].
    """
    H = ksp.shape[-2]
    W = ksp.shape[-1]
    S = ksp.shape[0]
    C = ksp.shape[1]
    # bring the coil ksp into coil image space
    img1 = ifftc(ksp,axes=(-2,-1))
    img1_cropped = resize(img1, oshape=(S,C,W,W))
    # FOV is now the same in both directions without modifying the resolution
    ksp1 = fftc(img1_cropped,axes=(-2,-1))
    # crop or pad the ksp isotropically in fourier space to the correct image size while mainting the same field of view (in width direction) in the original image
    ksp1_cropped = resize(ksp1, oshape=(S,C,final_res,final_res))
    img_out = ifftc(ksp1_cropped,axes=(-2,-1))

    return img_out

def read_fastmri_data(file_path):
    """
    This function reads k-space data from a .h5 file.

    Parameters:
    file_path (str): The path to the .h5 file containing FastMRI data.

    Returns:
    numpy.ndarray: The k-space data as a numpy array.
    """
    hf = h5py.File(file_path, 'r')
    ksp = np.asarray(hf['kspace'])
    return ksp

def combine_coils(ksp):
    """
    Combine multi-coil k-space data into a single coil image.

    This function reshapes the raw multi-coil k-space data, calculates sensitivity maps for the reshaped data using the BART tool's 'ecalib' command, and then uses these maps to create a single coil image via a fully sampled adjoint operation.

    Parameters:
    ksp (numpy.ndarray): The input multi-coil k-space data, shape: [B, C, H, W].

    Returns:
    numpy.ndarray: The output single coil image, shape: [B, 1, H, W].
    """
    # reshape raw multi-coil kspace to desired shape (ex [B,C,256,256])
    coil_img_rs = shape_data(ksp, final_res=256)
    coil_ksp_rs = fftc(coil_img_rs, axes=(-2,-1))

    # calculate sensitivity maps for reshaped coil ksp
    ksp_rs = coil_ksp_rs.transpose((2,3,0,1))
    maps = np.array(ksp_rs)
    #calculate Espirit maps with bart
    for j in tqdm(range(ksp_rs.shape[2])):
        sens = bart(1,'ecalib -m1 -W -c0', ksp_rs[:,:,j,None,:])#requires data of the form (Row,Column,None,Coil)<-output of ecalib too, this should then be saved (slice, coil, rows, columns)
        maps[:,:,j,:] = sens[:,:,0,:]

    maps_rs = maps.transpose((2,3,0,1))
    # use new maps to create single coil image via fully sampled adjoint operation
    single_coil_rs_img = adjoint(ksp=coil_ksp_rs, maps = maps_rs, mask = np.ones_like(coil_ksp_rs))
    return single_coil_rs_img

def preprocess_data_recipe():
    # for each file in the fastMRI dataset
    # call read_fastmri_data to get the kspace data
    # call combine_coils to create the combined coil image
    pass