#utils.py

#Copyright (c) 2020 Rachel Lea Ballantyne Draelos

#MIT License

#Permission is hereby granted, free of charge, to any person obtaining a copy

#of this software and associated documentation files (the "Software"), to deal

#in the Software without restriction, including without limitation the rights

#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

#copies of the Software, and to permit persons to whom the Software is

#furnished to do so, subject to the following conditions:

#The above copyright notice and this permission notice shall be included in all

#copies or substantial portions of the Software.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

#SOFTWARE

import torch

import numpy as np

"""CT volume preprocessing functions"""

#############################################

# Pixel Values (on torch Tensors for speed) #-----------------------------------

#############################################

def normalize(ctvol, lower_bound, upper_bound): #Done testing

    """Clip images and normalize"""

    #formula https://stats.stackexchange.com/questions/70801/how-to-normalize-data-to-0-1-range

    ctvol = torch.clamp(ctvol, lower_bound, upper_bound)

    ctvol = (ctvol - lower_bound) / (upper_bound - lower_bound)

    return ctvol

def torchify_pixelnorm_pixelcenter(ctvol, pixel_bounds):

    """Normalize using specified pixel_bounds and then center on the ImageNet

    mean. Used in 2019_10 dataset preparation"""

    #Cast to torch Tensor

    #use torch Tensor instead of numpy array because addition, subtraction,

    #multiplication, and division are faster in torch Tensors than np arrays

    ctvol = torch.from_numpy(ctvol).type(torch.float)

    #Clip Hounsfield units and normalize pixel values

    ctvol = normalize(ctvol, pixel_bounds[0], pixel_bounds[1])

    #Center on the ImageNet mean since you are using an ImageNet pretrained

    #feature extractor:

    ctvol = ctvol - 0.449

    return ctvol

###########

# Padding #---------------------------------------------------------------------

###########

def pad_slices(ctvol, max_slices): #Done testing

    """For <ctvol> of shape (slices, side, side) pad the slices to shape

    max_slices for output of shape (max_slices, side, side)"""

    padding_needed = max_slices - ctvol.shape[0]

    assert (padding_needed >= 0), 'Image slices exceed max_slices by'+str(-1*padding_needed)

    if padding_needed > 0:

        before_padding = int(padding_needed/2.0)

        after_padding = padding_needed - before_padding

        ctvol = np.pad(ctvol, pad_width = ((before_padding, after_padding), (0,0), (0,0)),

                     mode = 'constant', constant_values = np.amin(ctvol))

        assert ctvol.shape[0]==max_slices

    return ctvol

def pad_sides(ctvol, max_side_length): #Done testing

    """For <ctvol> of shape (slices, side, side) pad the sides to shape

    max_side_length for output of shape (slices, max_side_length,

    max_side_length)"""

    needed_padding = 0

    for side in [1,2]:

        padding_needed = max_side_length - ctvol.shape[side]

        if padding_needed > 0:

            before_padding = int(padding_needed/2.0)

            after_padding = padding_needed - before_padding

            if side == 1:

                ctvol = np.pad(ctvol, pad_width = ((0,0), (before_padding, after_padding), (0,0)),

                         mode = 'constant', constant_values = np.amin(ctvol))

                needed_padding += 1

            elif side == 2:

                ctvol = np.pad(ctvol, pad_width = ((0,0), (0,0), (before_padding, after_padding)),

                         mode = 'constant', constant_values = np.amin(ctvol))

                needed_padding += 1

    if needed_padding == 2: #if both sides needed to be padded, then they

        #should be equal (but it's possible one side or both were too large

        #in which case we wouldn't expect them to be equal)

        assert ctvol.shape[1]==ctvol.shape[2]==max_side_length

    return ctvol

def pad_volume(ctvol, max_slices, max_side_length):

    """Pad <ctvol> to a minimum size of

    [max_slices, max_side_length, max_side_length], e.g. [402, 308, 308]

    Used in 2019_10 dataset preparation"""

    if ctvol.shape[0] < max_slices:

        ctvol = pad_slices(ctvol, max_slices)

    if ctvol.shape[1] < max_side_length:

        ctvol = pad_sides(ctvol, max_side_length)

    return ctvol

###########################

# Reshaping to 3 Channels #-----------------------------------------------------

###########################

def sliceify(ctvol): #Done testing

    """Given a numpy array <ctvol> with shape [slices, square, square]

    reshape to 'RGB' [max_slices/3, 3, square, square]"""

    return np.reshape(ctvol, newshape=[int(ctvol.shape[0]/3), 3, ctvol.shape[1], ctvol.shape[2]])

def reshape_3_channels(ctvol):

    """Reshape grayscale <ctvol> to a 3-channel image

    Used in 2019_10 dataset preparation"""

    if ctvol.shape[0]%3 == 0:

        ctvol = sliceify(ctvol)

    else:

        if (ctvol.shape[0]-1)%3 == 0:

            ctvol = sliceify(ctvol[:-1,:,:])

        elif (ctvol.shape[0]-2)%3 == 0:

            ctvol = sliceify(ctvol[:-2,:,:])

    return ctvol

##################################

# Cropping and Data Augmentation #----------------------------------------------

##################################

def crop_specified_axis(ctvol, max_dim, axis): #Done testing

    """Crop 3D volume <ctvol> to <max_dim> along <axis>"""

    dim = ctvol.shape[axis]

    if dim > max_dim:

        amount_to_crop = dim - max_dim

        part_one = int(amount_to_crop/2.0)

        part_two = dim - (amount_to_crop - part_one)

        if axis == 0:

            return ctvol[part_one:part_two, :, :]

        elif axis == 1:

            return ctvol[:, part_one:part_two, :]

        elif axis == 2:

            return ctvol[:, :, part_one:part_two]

    else:

        return ctvol

def single_crop_3d_fixed(ctvol, max_slices, max_side_length):

    """Crop a single 3D volume to shape [max_slices, max_side_length,

    max_side_length]"""

    ctvol = crop_specified_axis(ctvol, max_slices, 0)

    ctvol = crop_specified_axis(ctvol, max_side_length, 1)

    ctvol = crop_specified_axis(ctvol, max_side_length, 2)

    return ctvol

def single_crop_3d_augment(ctvol, max_slices, max_side_length):

    """Crop a single 3D volume to shape [max_slices, max_side_length,

    max_side_length] with randomness in the centering and random

    flips or rotations"""

    #Introduce random padding so that the centered crop will be slightly random

    ctvol = rand_pad(ctvol)

    #Obtain the center crop

    ctvol = single_crop_3d_fixed(ctvol, max_slices, max_side_length)

    #Flip and rotate

    ctvol = rand_flip(ctvol)

    ctvol = rand_rotate(ctvol)

    #Make contiguous array to avoid Pytorch error

    return np.ascontiguousarray(ctvol)

def rand_pad(ctvol):

    """Introduce random padding between 0 and 15 pixels on each of the 6 sides

    of the <ctvol>"""

    randpad = np.random.randint(low=0,high=15,size=(6))

    ctvol = np.pad(ctvol, pad_width = ((randpad[0],randpad[1]), (randpad[2],randpad[3]), (randpad[4], randpad[5])),

                         mode = 'constant', constant_values = np.amin(ctvol))

    return ctvol

def rand_flip(ctvol):

    """Flip <ctvol> along a random axis with 50% probability"""

    if np.random.randint(low=0,high=100) < 50:

        chosen_axis = np.random.randint(low=0,high=3) #0, 1, and 2 are axis options

        ctvol =  np.flip(ctvol, axis=chosen_axis)

    return ctvol

def rand_rotate(ctvol):

    """Rotate <ctvol> some random amount axially with 50% probability"""

    if np.random.randint(low=0,high=100) < 50:

        chosen_k = np.random.randint(low=0,high=4)

        ctvol = np.rot90(ctvol, k=chosen_k, axes=(1,2))

    return ctvol

###########################################

# 2019_10 Dataset Preprocessing Sequences #-------------------------------------

###########################################

def prepare_ctvol_2019_10_dataset(ctvol, pixel_bounds, data_augment, num_channels,

                                  crop_type):

    """Pad, crop, possibly augment, reshape to correct

    number of channels, cast to torch tensor (to speed up subsequent operations),

    Clip Hounsfield units, normalize pixel values, center on the

    ImageNet mean, and return as a torch tensor (for crop_type='single')

    <pixel_bounds> is a list of ints e.g. [-1000,200] Hounsfield units. Used for

        pixel value clipping and normalization.

    <data_augment> is True to employ data augmentation, and False otherwise

    <num_channels> is an int, e.g. 3 to reshape the grayscale volume into

        a volume of 3-channel images

    <crop_type>: if 'single' then return the volume as one 3D numpy array."""

    max_slices = 402

    max_side_length = 420

    assert num_channels == 3 or num_channels == 1

    assert crop_type == 'single'

    #Padding to minimum size [max_slices, max_side_length, max_side_length]

    ctvol = pad_volume(ctvol, max_slices, max_side_length)

    #Cropping, and data augmentation if indicated

    if crop_type == 'single':

        if data_augment is True:

            ctvol = single_crop_3d_augment(ctvol, max_slices, max_side_length)

        else:

            ctvol = single_crop_3d_fixed(ctvol, max_slices, max_side_length)

        #Reshape to 3 channels if indicated

        if num_channels == 3:

            ctvol = reshape_3_channels(ctvol)

        #Cast to torch tensor and deal with pixel values

        output = torchify_pixelnorm_pixelcenter(ctvol, pixel_bounds)

    return output