# Allen Institute Software License - This software license is the 2-clause BSD
# license plus a third clause that prohibits redistribution for commercial
# purposes without further permission.
#
# Copyright 2015-2018. Allen Institute. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Redistributions for commercial purposes are not permitted without the
# Allen Institute's written permission.
# For purposes of this license, commercial purposes is the incorporation of the
# Allen Institute's software into anything for which you will charge fees or
# other compensation. Contact terms@alleninstitute.org for commercial licensing
# opportunities.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#

import warnings

import SimpleITK as sitk
import numpy as np


def get_sitk_image_information(image):
    ''' Extract information about a SimpleITK image

    Parameters
    ----------
    image : sitk.Image
        Extract information about this image.

    Returns
    -------
    dict : 
        Extracted information. Includes spacing, origin, size, direction, and 
        number of components per pixel

    '''

    return {'spacing': image.GetSpacing(),
            'origin': image.GetOrigin(),
            'size': image.GetSize(),
            'direction': image.GetDirection(),
            'ncomponents': image.GetNumberOfComponentsPerPixel()}


def set_sitk_image_information(image, information):
    ''' Set information on a SimpleITK image

    Parameters
    ----------
    image : sitk.Image
        Set information on this image.
    information : dict
        Stores information to be set. Supports spacing, origin, direction. Also 
        checks (but cannot set) size and number of components per pixel

    '''

    if 'spacing' in information:
        image.SetSpacing(information.pop('spacing'))
    if 'origin' in information:
        image.SetOrigin(information.pop('origin'))
    if 'direction' in information:
        image.SetDirection(information.pop('direction'))

    if 'size' in information:
        assert(np.array_equal( information.pop('size'), image.GetSize() ))
    if 'ncomponents' in information:
        assert( information.pop('ncomponents') == image.GetNumberOfComponentsPerPixel() )
    
    if not len(information) == 0:
        warnings.warn('unwritten keys: {}'.format(','.join(information.keys())))


def fix_array_dimensions(array, ncomponents=1):
    ''' Convenience function that reorders ndarray dimensions for io with SimpleITK

    Parameters
    ----------
    array : np.ndarray
        The array to be reordered
    ncomponents : int, optional
        Number of components per pixel, default 1. 

    Returns
    -------
    np.ndarray : 
        Reordered array

    '''

    act_size = list(array.shape)
    ndims = len(act_size)
    multicomponent = ncomponents > 1

    from_order = list(range( ndims - multicomponent ))
    to_order = list(range( ndims - multicomponent ))[::-1]

    if multicomponent:
        from_order += [-1]
        to_order += [-1]

    return np.ascontiguousarray(np.moveaxis(array, from_order, to_order))


def read_ndarray_with_sitk(path):
    ''' Read a numpy array from a file using SimpleITK

    Parameters
    ----------
    path : str
        Read from this path
    
    Returns
    -------
    image : np.ndarray
        Obtained array
    information : dict
        Additional information about the array

    '''

    image = sitk.ReadImage(str(path))
    information = get_sitk_image_information(image)
    image = sitk.GetArrayFromImage(image)

    image = fix_array_dimensions(image, information['ncomponents'])
    return image, information


def write_ndarray_with_sitk(array, path, **information):
    ''' Write a numpy array to a file using SimpleITK

    Parameters
    ----------
    array : np.ndarray
        Array to be written.
    path : str
        Write to here
    **information : dict
        Contains additional information to be stored in the image file. 
        See set_sitk_image_information for more information.

    '''

    if not 'ncomponents' in information:
        information['ncomponents'] = 1
    ncomponents = information.pop('ncomponents')

    array = fix_array_dimensions(array, ncomponents)
    
    array = sitk.GetImageFromArray(array, ncomponents > 1)
    set_sitk_image_information(array, information)

    sitk.WriteImage(array, str(path))
