# ==============================================================================
# Copyright (C) 2023 Haresh Rengaraj Rajamohan, Tianyu Wang, Kevin Leung, 
# Gregory Chang, Kyunghyun Cho, Richard Kijowski & Cem M. Deniz 
#
# This file is part of OAI-MRI-TKR
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# ==============================================================================
import random
import numpy as np
import utils
import math
import h5py
from scipy import ndimage as nd


class Random_Rotation:


    def __init__(self,image):
        self.image = image

    def RandomRotation(self,output_shape):
        '''
        The function generated a rotation matrix randomly such that the rotation axis is uniformly distributed
        on a unit sphere and the angle is uniformally distributed.
        :param output_shape: shape of the output image
        :return:Randomly rotated 3D image
        '''


        rotation_matrix =utils.generating_random_rotation_matrix()
        Output_image =self.__rotation__(output_shape, rotation_matrix)
        return Output_image

    def __rotation__(self,Output_shape,rotation_matrix):
        '''
        :param Output_shape: shape of the rotated image
        :param rotation_matrix:
        :return: Rotated 3D image based on rotation matrix
        '''

        image = self.image
        h,w,d = image.shape
        coordinate_j, coordinate_i, coordinate_k = np.meshgrid(
            np.array(range(Output_shape[1])), np.array(range(Output_shape[0])),
            np.array(range(Output_shape[2])))

        center_target = np.array([int(sh/2) for sh in list(Output_shape)]).reshape(3,1)
        center_source = np.array([int(h/2),int(w/2),int(d/2)]).reshape(3,1)
        coordinate_init = np.array([coordinate_j.flatten(), coordinate_i.flatten(), coordinate_k.flatten()])
        Rotation_matrix = rotation_matrix
        coordinate = coordinate_init - np.matlib.repmat(center_target,1,coordinate_init.shape[1]) + np.matlib.repmat(np.matmul(Rotation_matrix,center_source),1,coordinate_init.shape[1])


        mapped_to_source_coordinate = np.linalg.solve(Rotation_matrix,coordinate)
        output_coordinate_value = nd.map_coordinates(input=image,coordinates=mapped_to_source_coordinate,cval = -1000,order = 4,mode = 'constant')

        Output_image = -1000*np.ones(shape=Output_shape,dtype=float)

        for k in range(coordinate_init.shape[1]):
            Output_image[coordinate_init[0,k],coordinate_init[1,k],coordinate_init[2,k]] = output_coordinate_value[k]

        return Output_image

    def __random_rotation_matrix_fixing_rotation_axis__(self,axis = 0):

        '''
        :param axis: axis to fix
        :return: randomly rotated 3D image such that the axis of rotation is fixed
        '''

        axis_of_rotation = [0.0, 0.0, 0.0]
        axis_of_rotation[axis] = 1.0
        angle_of_rotation = 2 * math.pi * random.uniform(0, 1)
        Rotation_matrix = utils.angle_axis_to_rotation_matrix(angle=angle_of_rotation,axis=axis_of_rotation)
        return Rotation_matrix

    def RandomRotation_x_axis(self,Output_shape):

        '''
        :param Output_shape: shape of the output image
        :return: rotated 3D image along the x axis
        '''

        h, w, d = self.image.shape
        rotation_matrix = self.__random_rotation_matrix_fixing_rotation_axis__(axis = 0 )
        Output_image = self.__rotation__(Output_shape, rotation_matrix)
        return Output_image

    def RandomRotation_y_axis(self,Output_shape):
        '''
        :param Output_shape: shape of the output image
        :return: rotated 3D image along the y axis
        '''

        h, w, d = self.image.shape
        rotation_matrix = self.__random_rotation_matrix_fixing_rotation_axis__(axis=1)
        Output_image = self.__rotation__(Output_shape, rotation_matrix)
        return Output_image

    def RandomRotation_z_axis(self,Output_shape):
        '''
        :param Output_shape: shape of the output image
        :return: rotated 3D image along the z axis
        '''
        h, w, d = self.image.shape
        rotation_matrix = self.__random_rotation_matrix_fixing_rotation_axis__(axis=2)
        print("rotation_matrix: ",rotation_matrix)
        Output_image = self.__rotation__(Output_shape, rotation_matrix)
        return Output_image


class RandomCrop:

    '''
    Randomly Crop 3D image
    '''

    def __init__(self,image):
        self.image = image
        self.h,self.w,self.d = image.shape

    def __functional__(self,size):
        '''
        :param size: crop size
        :return: cropped 3D image
        '''

        h, w, d = self.image.shape
        crop_h, crop_w, crop_d = size
        i = random.randint(0, h - crop_h)
        j = random.randint(0, w - crop_w)
        k = random.randint(0, d - crop_d)
        crop_image = self.image[i:i + crop_h, j:j + crop_w, k:k + crop_d]
        return crop_image

    def crop_along_hieght_width(self,crop_size):

        crop = (crop_size[0],crop_size[2],self.d)
        self.crop_image = self.__functional__(size=crop)
        return self.crop_image

    def crop_along_hieght_width_depth(self,crop_size):
        self.crop_image = self.__functional__(size=crop_size)
        return self.crop_image


class CenterCrop:

    '''
    CenterCrop Images
    '''

    def __init__(self,image):
        self.image = image
        self.h,self.w,self.d = image.shape

    def __functional__(self,size):
        '''
        :param size: crop Size
        :return: Center Crop Images
        '''
        crop_h = int((self.h-size[0])/2)
        crop_w = int((self.w-size[1])/2)
        crop_d = int((self.d-size[2])/2)
        return self.image[crop_h:crop_h+size[0],crop_w:crop_w+size[1],crop_d:crop_d+size[2]]

    def crop(self,size):
        return self.__functional__(size)


class RandomFlip:

    def __init__(self,image,p=0.5):
        self.image = image
        self.p = p
        self.h,self.w,self.d = image.shape

    def horizontal_flip(self,p=-1):
        '''
        :param p: probability of flip
        :return: randomly horizontaly flipped image
        '''
        if p == -1:
            p = self.p

        integer = random.randint(0, 1)
        if integer <= p:
            output_image = self.image[:, -1:0:-1, :]
        else:
            output_image = self.image

        return output_image

    def vertical_flip(self,p=-1):

        '''
        :param p: probability of flip
        :return: randomly vertically flipped image
        '''

        if p == -1:
            p = self.p

        integer = random.randint(0, 1)
        if integer <= p:
            output_image = np.flipud(self.image)
        else:
            output_image = self.image

        return output_image
    
    def horizontal_flip(self,p=-1):

        '''
        :param p: probability of flip
        :return: randomly vertically flipped image
        '''

        if p == -1:
            p = self.p

        integer = random.randint(0, 1)
        if integer <= p:
            output_image = np.fliplr(self.image)
        else:
            output_image = self.image

        return output_image