import numpy as np

import nibabel as nib

import scipy.ndimage

import warnings

import PP

import sys

#---------------------------------------------

#Functions for image augmentations on 3D input

#---------------------------------------------

#img_b, label_b is (batch_num) x 1 x dim1 x dim2 x dim3

#takes in a list of 3D images (1st one is input, 2nd one needs to be label)

def augmentPatchLossy(imgs, rotation=[5,5,5], scale_min=0.9, scale_max=1.1, flip_lvl = 0):

    new_imgs = []

    rot_x = np.random.uniform(-rotation[0], rotation[0]) * np.pi / 180.0

    rot_y = np.random.uniform(-rotation[1], rotation[1]) * np.pi / 180.0

    rot_z = np.random.uniform(-rotation[2], rotation[2]) * np.pi / 180.0

    zoom_val = np.random.uniform(scale_min, scale_max)

    for i in range(len(imgs)):

        l = convertBatchToList(imgs[i])

        if i == 0:

            spline_orders = [3] * len(l)

        else:

            spline_orders = [0] * len(l)

        scaled = applyScale(l, zoom_val, spline_orders)

        rotated = applyRotation(scaled, [rot_x, rot_y, rot_z], spline_orders)

        new_imgs.append(convertListToBatch(rotated))

    return imgs

def convertBatchToList(img):

    l = []

    b, c, d1, d2, d3 = img.shape

    for i in range(img.shape[0]):

        l.append(img[i,:,:,:,:].reshape([1,c,d1,d2,d3]))

    return l

def convertListToBatch(img_list):

    b, c, d1, d2, d3 = img_list[0].shape

    a = np.zeros([len(img_list), c, d1,d2,d3])

    for i in range(len(img_list)):

        a[i,:,:,:,:] = img_list[i]

    return a

def augmentPatchLossLess(imgs):

    new_imgs = []

    p = np.random.rand(3) > 0.5

    locations = np.where(p == 1)[0] + 2

    for i in range(len(imgs)):

        l = convertBatchToList(imgs[i])

        if i == 0:

            spline_orders = [3] * len(l)

        else:

            spline_orders = [0] * len(l)

        flipped = applyFLIPS2(l, locations)

        rot_x = np.random.randint(4) * np.pi / 2.0 # (0,1,2,3)*90/180.0

        rot_y = np.random.randint(4) * np.pi / 2.0 # (0,1,2,3)*90/180.0

        rot_z = np.random.randint(4) * np.pi / 2.0 # (0,1,2,3)*90/180.0

        rotated = applyRotation(flipped, [rot_x, rot_y, rot_z], spline_orders)

        new_imgs.append(convertListToBatch(rotated))

    return new_imgs

def augmentBoth(imgs):

    imgs = augmentPatchLossy(imgs)

    imgs = augmentPatchLessLess(imgs)

    return imgs

def getRotationVal(rotation=[5,5,5]):

    rot_x = np.random.uniform(-rotation[0], rotation[0]) * np.pi / 180.0

    rot_y = np.random.uniform(-rotation[1], rotation[1]) * np.pi / 180.0

    rot_z = np.random.uniform(-rotation[2], rotation[2]) * np.pi / 180.0

    return rot_x, rot_y, rot_z

def getScalingVal(scale_min = 0.9, scale_max = 1.1):

    return np.random.uniform(scale_min, scale_max)

def applyFLIPS(images, flip_lvl = 0):

    if flip_lvl == 0:

        p = np.random.rand(2) > 0.5

    else:

        p = np.random.rand(3) > 0.5

    locations = np.where(p == 1)[0] + 2

    new_imgs = []

    for img in images:

        for i in locations:

            img = np.flip(img, axis=i)

        new_imgs.append(img)

    return new_imgs

def applyFLIPS2(images, locations):

    new_imgs = []

    for img in images:

        for i in locations:

            img = np.flip(img, axis=i)

        new_imgs.append(img)

    return new_imgs

def applyRotation(images, rot, spline_orders):

    transform_x = np.array([[1.0,               0.0,            0.0],

                            [0.0,               np.cos(rot[0]), -np.sin(rot[0])],

                            [0.0,               np.sin(rot[0]), np.cos(rot[0])]])

    transform_y = np.array([[np.cos(rot[1]),    0.0,            np.sin(rot[1])],

                            [0.0,               1.0,            0.0],

                            [-np.sin(rot[1]),   0.0,            np.cos(rot[1])]])

    transform_z = np.array([[np.cos(rot[2]),    -np.sin(rot[2]),    0.0],

                            [np.sin(rot[2]),    np.cos(rot[2]),     0.0],

                            [0.0,               0,                  1]])

    transform = np.dot(transform_z, np.dot(transform_x, transform_y))

    new_imgs = []

    for i, img in enumerate(images):

        mid_index = 0.5 * np.asarray(img.squeeze().shape, dtype=np.int64)

        offset = mid_index - mid_index.dot(np.linalg.inv(transform))

        new_img = scipy.ndimage.affine_transform(

                                            input = img.squeeze(), 

                                            matrix = transform, 

                                            offset = offset, 

                                            order = spline_orders[i],

                                            mode = 'nearest')

        new_img = new_img[np.newaxis,np.newaxis,:]

        new_imgs.append(new_img)

    return new_imgs

def applyScale(images, zoom_val, spline_orders):

    new_imgs = []

    for i, img in enumerate(images):

        with warnings.catch_warnings():

            warnings.simplefilter("ignore")

            try:

                new_img = scipy.ndimage.zoom(img.squeeze(), zoom_val, order = spline_orders[i])

                new_img = new_img[np.newaxis,np.newaxis,:]

                new_imgs.append(new_img)

            except:

                pass

    return new_imgs