import SimpleITK as sitk

import numpy as np

def calculate_origin_offset(new_spacing, old_spacing):

    return np.subtract(new_spacing, old_spacing)/2

def sitk_resample_to_spacing(image, new_spacing=(1.0, 1.0, 1.0), interpolator=sitk.sitkLinear, default_value=0.):

    zoom_factor = np.divide(image.GetSpacing(), new_spacing)

    new_size = np.asarray(np.ceil(np.round(np.multiply(zoom_factor, image.GetSize()), decimals=5)), dtype=np.int16)

    offset = calculate_origin_offset(new_spacing, image.GetSpacing())

    reference_image = sitk_new_blank_image(size=new_size, spacing=new_spacing, direction=image.GetDirection(),

                                           origin=image.GetOrigin() + offset, default_value=default_value)

    return sitk_resample_to_image(image, reference_image, interpolator=interpolator, default_value=default_value)

def sitk_resample_to_image(image, reference_image, default_value=0., interpolator=sitk.sitkLinear, transform=None,

                           output_pixel_type=None):

    if transform is None:

        transform = sitk.Transform()

        transform.SetIdentity()

    if output_pixel_type is None:

        output_pixel_type = image.GetPixelID()

    resample_filter = sitk.ResampleImageFilter()

    resample_filter.SetInterpolator(interpolator)

    resample_filter.SetTransform(transform)

    resample_filter.SetOutputPixelType(output_pixel_type)

    resample_filter.SetDefaultPixelValue(default_value)

    resample_filter.SetReferenceImage(reference_image)

    return resample_filter.Execute(image)

def sitk_new_blank_image(size, spacing, direction, origin, default_value=0.):

    image = sitk.GetImageFromArray(np.ones(size, dtype=np.float).T * default_value)

    image.SetSpacing(spacing)

    image.SetDirection(direction)

    image.SetOrigin(origin)

    return image

def resample_to_spacing(data, spacing, target_spacing, interpolation="linear", default_value=0.):

    image = data_to_sitk_image(data, spacing=spacing)

    if interpolation is "linear":

        interpolator = sitk.sitkLinear

    elif interpolation is "nearest":

        interpolator = sitk.sitkNearestNeighbor

    else:

        raise ValueError("'interpolation' must be either 'linear' or 'nearest'. '{}' is not recognized".format(

            interpolation))

    resampled_image = sitk_resample_to_spacing(image, new_spacing=target_spacing, interpolator=interpolator,

                                               default_value=default_value)

    return sitk_image_to_data(resampled_image)

def data_to_sitk_image(data, spacing=(1., 1., 1.)):

    if len(data.shape) == 3:

        data = np.rot90(data, 1, axes=(0, 2))

    image = sitk.GetImageFromArray(data)

    image.SetSpacing(np.asarray(spacing, dtype=np.float))

    return image

def sitk_image_to_data(image):

    data = sitk.GetArrayFromImage(image)

    if len(data.shape) == 3:

        data = np.rot90(data, -1, axes=(0, 2))

    return data