# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_manipulate.ipynb (unless otherwise specified).

__all__ = ['crop', 'pad', 'resample', 'resample_by', 'resize']

# Cell
from fastscript import call_parse, Param, bool_arg
from scipy import ndimage

import math
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import pydicom

# Cell
def crop(image_array, final_dims_in_pixels, zero_fill_mode=False):
    """
    Crop image data to final_dim_in_pixels

    Attributes:
        image_array (float, np.array): 3D numpy array containing image data
        final_dim_in_pixels    (list): Final number of pixels in xyz dimensions. Example: [256, 256, 80]
        zero_fill_mode         (bool): If True, returns array filled with zeros

    Returns:
        cropped_image_array (arr): Resized array containing image data
    """

    dims = len(final_dims_in_pixels)
    original_dims_in_pixels = [image_array.shape[d]
                               for d in range(len(image_array.shape))]

    # test if input and output dimensions match
    if dims != len(original_dims_in_pixels):
        raise ValueError("Dimensions of the input (" + str(len(image_array.shape)) +
                         ") do not match those of output (" + str(len(final_dims_in_pixels)) + ")")

    # test if desired final image is smaller than original
    if any(final_dims_in_pixels[d] > original_dims_in_pixels[d] for d in range(dims)):
        raise ValueError(
            "Final dimensions are larger than original. Did you mean to `pad`?")

    cropped_image_array = np.zeros(image_array.shape)
    new_first_pixel = [0 for i in range(dims)]
    new_last_pixel = [0 for i in range(dims)]

    for dim in range(dims):
        new_first_pixel[dim] = int(math.floor(
            (original_dims_in_pixels[dim] - final_dims_in_pixels[dim]) / 2))
        new_last_pixel[dim] = new_first_pixel[dim] + final_dims_in_pixels[dim]

    # for 2D:
    if dims == 2:
        cropped_image_array = image_array[new_first_pixel[0]: new_last_pixel[0],
                                          new_first_pixel[1]: new_last_pixel[1]]
    elif dims == 3:
        cropped_image_array = image_array[new_first_pixel[0]: new_last_pixel[0],
                                          new_first_pixel[1]: new_last_pixel[1],
                                          new_first_pixel[2]: new_last_pixel[2]]
    if zero_fill_mode:
        cropped_image_array = cropped_image_array*0.

    return(cropped_image_array)

# Cell
def pad(image_array, final_dims_in_pixels, zero_fill_mode=False):
    """
    Pad image data to final_dim_in_pixels

    Attributes:
        image_array (float, np.array): 3D numpy array containing image data
        final_dim_in_pixels    (list): Final number of pixels in xyz dimensions. Example: [256, 256, 80]
        zero_fill_mode         (bool): If True, returns array filled with zeros

    Returns:
        padded_image_array (arr): Resized array containing image data
    """

    dims = len(final_dims_in_pixels)
    original_dims_in_pixels = [image_array.shape[d]
                               for d in range(len(image_array.shape))]

    # test if input and output dimensions match
    if dims != len(original_dims_in_pixels):
        raise ValueError("Dimensions of the input (" + str(len(image_array.shape)) +
                         ") do not match those of output (" + str(len(final_dims_in_pixels)) + ")")

    # test if desired final image is larger than original
    if any(final_dims_in_pixels[d] < original_dims_in_pixels[d] for d in range(dims)):
        raise ValueError(
            "Final dimensions are smaller than original. Did you mean to `crop`?")

    padded_image_array = np.zeros(final_dims_in_pixels)
    new_first_image_pixel = [0 for i in range(dims)]
    new_last_image_pixel = [0 for i in range(dims)]

    for dim in range(dims):
        new_first_image_pixel[dim] = int(math.floor(
            (final_dims_in_pixels[dim] - original_dims_in_pixels[dim]) / 2))
        new_last_image_pixel[dim] = new_first_image_pixel[dim] + \
            original_dims_in_pixels[dim]

    # for 2D:
    if dims == 2:
        padded_image_array[new_first_image_pixel[0]: new_last_image_pixel[0],
                           new_first_image_pixel[1]: new_last_image_pixel[1]] = image_array
    elif dims == 3:
        padded_image_array[new_first_image_pixel[0]: new_last_image_pixel[0],
                           new_first_image_pixel[1]: new_last_image_pixel[1],
                           new_first_image_pixel[2]: new_last_image_pixel[2]] = image_array
    if zero_fill_mode:
        padded_image_array = padded_image_array*0.

    return(padded_image_array)

# Cell
def resample(image_array, target_pixel_dims_list, is_seg=False):
    """
    Resamples `image` to new resolution according to a compression factor using scipy's `ndimage`

    Attributes:
        image_array (float, np.array) : Array of voxel values for image
        target_pixel_dims_list (list) : Dimensions of final image in pixels (e.g. [256, 256, 128])
        is_seg (bool)                 : Whether or not the image is a segmentation (i.e. pixel values are int)

    Returns:
        resampled_image_array (arr): Resampled array containing image data
    """
    # if dealing with segmentation images, do not interpolate
    # this is done by choosing interpolation order == 0
    order = 0 if is_seg == True else 3

    original_dims_in_pixels = [image_array.shape[d]
                               for d in range(len(image_array.shape))]
    compression_list = [target_pixel_dims_list[d] / original_dims_in_pixels[d]
                        for d in range(len(image_array.shape))]

    if (is_seg):
        resized_image = ndimage.interpolation.zoom(
            image_array, zoom=compression_list, order=order, cval=0)
        return resized_image
    else:
        resized_image = ndimage.interpolation.zoom(
            image_array, zoom=compression_list, order=order, cval=0)
        return resized_image

# Cell
def resample_by(image_array, compression_factor_list, is_seg=False):
    """
    Resamples `image` to new resolution according to a compression factor using scipy's `ndimage`

    Attributes:
        image_array (float, np.array)  : Array of voxel values for image
        compression_factor_list (list) : Compression or expansion factor (e.g. [0.1, 10.0, 1.0])
        is_seg (bool)                  : Whether or not the image is a segmentation (i.e. pixel values are int)

    Returns:
        resampled_image_array (arr): Resampled array containing image data
    """
    original_dims_in_pixels = [image_array.shape[d]
                               for d in range(len(image_array.shape))]
    target_pixel_dims_list = [int(math.floor(compression_factor_list[d] * original_dims_in_pixels[d]))
                            for d in range(len(image_array.shape))]

    resized_image = resample(image_array, target_pixel_dims_list, is_seg)

    return(resized_image)

# Cell
def resize(image_array, original_pixel_spacing_list, target_field_of_view_list=[140,140,80],
           target_pixel_spacing_list=[0.5,0.5,0.5], is_seg=False):
    """
    Resizes `image_array` to desired field of view by first resampling and then cropping / padding.
    The returned resampled image has isotropic pixel spacing = 0.5mm

    Args:
        image_array: np.array of voxel values
        pixel_spacing_list: list of pixel spacing values prior to resizing in mm
        target_field_of_view_list: list of field of values in mm (default: 140, 140, 80)
        target_pixel_spacing_list: list of pixel spacing values after resizing in mm (default: 0.5, 0.5, 0.5)
        is_seg (bool): whether image_array is a segmentation or not (default: False)

    Returns:
        image_cropped_resampled: np.array with voxels values after cropping and re-sampling

    """
    import math

    num_dim = len(image_array.shape)
    compression_factor_list = [original_pixel_spacing_list[i]/target_pixel_spacing_list[i] for i in range(num_dim)]
    image_resampled = resample_by(image_array=image_array, compression_factor_list= compression_factor_list)
    current_dim_in_pixels_list = list(image_resampled.shape)
    target_dim_in_pixels_list = [int(target_field_of_view_list[i]/target_pixel_spacing_list[i]) for i in range(num_dim)]

    #pad or crop in xy (while ignoring z)
    if current_dim_in_pixels_list[0] >= target_dim_in_pixels_list[0]:
        print('Cropping image in xy')
        image_resampled_cropped = crop(image_array=image_resampled,
                                       final_dims_in_pixels=[target_dim_in_pixels_list[0],
                                                             target_dim_in_pixels_list[1],
                                                             current_dim_in_pixels_list[2]])
    else:
        print('Padding image in xy')
        image_resampled_cropped = pad(image_array=image_resampled,
                                      final_dims_in_pixels=[target_dim_in_pixels_list[0],
                                                            target_dim_in_pixels_list[1],
                                                            current_dim_in_pixels_list[2]])

    #pad or crop in z (while ignoring xy)
    current_dim_in_pixels_list = list(image_resampled_cropped.shape)

    if current_dim_in_pixels_list[2] >= target_dim_in_pixels_list[2]:
        print('Cropping image in z')
        image_resampled_cropped = crop(image_array=image_resampled_cropped,
                                       final_dims_in_pixels=[current_dim_in_pixels_list[0],
                                                             current_dim_in_pixels_list[1],
                                                             target_dim_in_pixels_list[2]])
    else:
        print('Padding image in z')
        image_resampled_cropped = pad(image_array=image_resampled_cropped,
                                      final_dims_in_pixels=[current_dim_in_pixels_list[0],
                                                            current_dim_in_pixels_list[1],
                                                            target_dim_in_pixels_list[2]])

    return(image_resampled_cropped)