# Copyright 2019 Population Health Sciences and Image Analysis, German Center for Neurodegenerative Diseases(DZNE)

#

#    Licensed under the Apache License, Version 2.0 (the "License");

#    you may not use this file except in compliance with the License.

#    You may obtain a copy of the License at

#

#        http://www.apache.org/licenses/LICENSE-2.0

#

#    Unless required by applicable law or agreed to in writing, software

#    distributed under the License is distributed on an "AS IS" BASIS,

#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

#    See the License for the specific language governing permissions and

#    limitations under the License.

import numpy as np

import nibabel as nib

import scipy.ndimage

import os

def calculated_new_ornt(iornt,base_ornt):

    new_iornt=iornt[:]

    for axno, direction in np.asarray(base_ornt):

        idx=np.where(iornt[:,0] == axno)

        idirection=iornt[int(idx[0][0]),1]

        if direction == idirection:

            new_iornt[int(idx[0][0]), 1] = 1.0

        else:

            new_iornt[int(idx[0][0]), 1] = -1.0

    return new_iornt

def check_orientation(img,base_ornt=np.array([[0,-1],[1,1],[2,1]])):

    iornt=nib.io_orientation(img.affine)

    if not np.array_equal(iornt,base_ornt):

        img = img.as_reoriented(calculated_new_ornt(iornt,base_ornt))

    return img

def resample(image, spacing, new_spacing=[1, 1, 1],order=1,prefilter=True):

    # Determine current pixel spacing

    resize_factor = spacing / new_spacing

    new_real_shape = image.shape * resize_factor

    new_shape = np.round(new_real_shape)

    real_resize_factor = new_shape / image.shape

    new_spacing = spacing / real_resize_factor

    image = scipy.ndimage.interpolation.zoom(image,real_resize_factor,order=order,prefilter=prefilter)

    return image, new_spacing

def define_size(mov_dim,ref_dim):

    new_dim=np.zeros(len(mov_dim),dtype=np.int)

    borders=np.zeros((len(mov_dim),2),dtype=int)

    padd = [int(mov_dim[0] // 2), int(mov_dim[1] // 2), int(mov_dim[2] // 2)]

    for i in range(len(mov_dim)):

        new_dim[i]=int(max(2*mov_dim[i],2*ref_dim[i]))

        borders[i,0]= int(new_dim[i] // 2) -padd [i]

        borders[i,1]= borders[i,0] +mov_dim[i]

    return list(new_dim),borders

def map_size(arr,base_shape,axial):

    if axial:

        base_shape[2]=arr.shape[2]

    print('Volume will be resize from %s to %s ' % (arr.shape, base_shape))

    new_shape,borders=define_size(np.array(arr.shape),np.array(base_shape))

    new_arr=np.zeros(new_shape)

    final_arr=np.zeros(base_shape)

    new_arr[borders[0,0]:borders[0,1],borders[1,0]:borders[1,1],borders[2,0]:borders[2,1]]= arr[:]

    middle_point = [int(new_arr.shape[0] // 2), int(new_arr.shape[1] // 2), int(new_arr.shape[2] // 2)]

    padd = [int(base_shape[0]/2), int(base_shape[1]/2), int(base_shape[2]/2)]

    low_border=np.array((np.array(middle_point)-np.array(padd)),dtype=int)

    high_border=np.array(np.array(low_border)+np.array(base_shape),dtype=int)

    final_arr[:,:,:]= new_arr[low_border[0]:high_border[0],

                   low_border[1]:high_border[1],

                   low_border[2]:high_border[2]]

    return final_arr

def map_image(img_arr,base_zoom,izoom,order,axial):

    if axial:

        base_zoom[2] = izoom[2]

    print('Volume will be sample from %s to %s ' % (izoom, base_zoom))

    resample_arr, izoom= resample(img_arr, spacing=np.array(izoom),

                                         new_spacing=np.array(base_zoom), order=order)

    resample_arr[resample_arr < 0] = 0

    return resample_arr,izoom

def conform(img,flags,order,save_path,mod,axial=False):

    """

    Args:

        img: nibabel img: Loaded source image

        flags: dict : Dictionary containing the image size, spacing and orientation

        order: int : interpolation order (0=nearest,1=linear(default),2=quadratic,3=cubic)

    Returns:

        new_img: nibabel img : conformed nibabel image

    """

    save=False

    # check orientation LAS

    img=check_orientation(img,base_ornt=flags['base_ornt'])

    img_arr=img.get_data()

    img_header = img.header

    # check voxel sizer

    i_zoom=img.header.get_zooms()

    #check the spacing idx for interpolation

    if axial:

        idx=2

    else:

        idx=3

    if not np.allclose(np.array(i_zoom)[:idx],np.array(flags['spacing'])[:idx],rtol=0.3):

        img_arr,i_zoom= map_image(img_arr,flags['spacing'],i_zoom,order,axial)

        save=True

    ishape = img_arr.shape

    # check dimensions

    if int(ishape[0]) != int(flags['imgSize'][0]) or int(ishape[1]) != int(flags['imgSize'][1]) or int(ishape[2]) != int(flags['imgSize'][2]):

        img_arr=map_size(img_arr,flags['imgSize'],axial)

        save = True

    img_header.set_data_shape(img_arr.shape)

    img_header.set_zooms(i_zoom)

    affine = img_header.get_qform()

    affine[0][3] += ((flags['imgSize'][0] - ishape[0]) / 2 * i_zoom[0])

    affine[1][3] -= ((flags['imgSize'][1] - ishape[1]) / 2 * i_zoom[1])

    affine[2][3] -= ((flags['imgSize'][2] - ishape[2]) / 2 * i_zoom[2])

    img_header.set_qform(affine)

    new_img = nib.Nifti1Image(img_arr, affine, img_header)

    #save images if modified

    if save:

        if not os.path.isdir(os.path.join(save_path, 'MRI')):

            os.mkdir(os.path.join(save_path, 'MRI'))

        mri_path = os.path.join(save_path, 'MRI')

        if mod == 'fat':

            new_img_path = os.path.join(mri_path, 'FatImaging_F.nii.gz')

        else:

            new_img_path = os.path.join(mri_path, 'FatImaging_W.nii.gz')

        nib.save(new_img, new_img_path)

    return new_img