# ============================================================================================= #
#  Author:       Filip Lux                                                                      #
#  Copyright:    Filip Lux          lux.filip@gmail.com                                         #
#                                                                                               #
#  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.                                                                                    #
# ============================================================================================= #

from typing import Sequence, Union
from ..typing import TypeSextetFloat, TypeTripletFloat, TypePairFloat
import numpy as np
import SimpleITK as sitk

from collections.abc import Iterable


def parse_limits(input_limit: Union[float, TypePairFloat, TypeTripletFloat, TypeSextetFloat],
                 identity_element: float = 1) -> TypeSextetFloat:

    # input_limit = None
    # returns (ie, ie, ie, ie, ie, ie)
    if input_limit is None:
        return tuple((identity_element, ) * 6)

    # input_limit = x : float
    # returns (ie-x, ie+x, ie-x, ie+x, ie-x, ie+x)
    elif len(input_limit) == 1:
        return tuple((identity_element - input_limit, identity_element + input_limit) * 3)

    # input_limit = (a, b) : TypePairFloat
    # returns (a, b, a, b, a, b)
    elif len(input_limit) == 2:
        a, b = input_limit
        return a, b, a, b, a, b

    # input_limit = (a, b, c) : TypeTripletFloat
    # returns (ie-a, ie+a, ie-b, ie+b, ie-c, ie+c)
    elif len(input_limit) == 3:
        res = []
        for item in input_limit:
            # input_limit = ((a, b), (c, d), (e, f))
            # return (a, b, c, d, e, f)
            if isinstance(item, Iterable):
                for val in item:
                    res.append(float(val))
            # input_limit = (a, b, c)
            # return (ie-a, ie+a, ie-b, ie+b, ie-c, ie+c)
            else:
                res.append(float(identity_element - item))
                res.append(float(identity_element + item))
        return tuple(res)

    # input_limit = (a, b, c, d, e, f)
    # returns (a, b, c, d, e, f)
    elif len(input_limit) == 6:
        return input_limit


def parse_coefs(coefs: Union[float, TypeTripletFloat],
                identity_element: float = 1) -> TypeTripletFloat:

    # input_limit = None
    # return (ie, ie, ie)
    if coefs is None:
        return tuple((identity_element, ) * 3)
    # return (a, a, a)
    elif isinstance(coefs, (int, float)):
        return coefs, coefs, coefs
    # return (a, b, c)
    elif len(coefs) == 3:
        return coefs


def get_image_center(image: np.array,
                     spacing: TypeTripletFloat = (1., 1., 1.),
                     lps: bool = False) -> TypeTripletFloat:

    shape = np.array(image.shape)
    if len(shape) == 3:
        center = (shape - 1) / 2
    elif len(shape) == 4:
        center = (shape[1:4] - 1) / 2
    else:
        center = np.array((0, 0, 0))

    if lps:
        center = ras_to_lps(center)

    return center * np.array(spacing)


def to_tuple(param, low=None, bias=None):
    """Convert input argument to min-max tuple

        Args:
            param (scalar, tuple or list of 2+ elements): Input value.
                If value is scalar, return value would be (offset - value, offset + value).
                If value is tuple, return value would be value + offset (broadcasted).
            low:  Second element of tuple can be passed as optional argument
            bias: An offset factor added to each element
    """
    if low is not None and bias is not None:
        raise ValueError("Arguments low and bias are mutually exclusive")

    if param is None:
        return param

    if isinstance(param, (int, float)):
        if low is None:
            param = -param, +param
        else:
            param = (low, param) if low < param else (param, low)
    elif isinstance(param, Sequence):
        param = tuple(param)
    else:
        raise ValueError("Argument param must be either scalar (int, float) or tuple")

    if bias is not None:
        return tuple(bias + x for x in param)

    return tuple(param)


# Simple ITK uses LPS coordinates format
def ras_to_lps(triplet: Sequence[float]):
    return np.array((-1, -1, 1), dtype=float) * np.asarray(triplet)


def np_to_sitk(img: np.array) -> sitk.Image:

    # image in format (c, s1, s2, s3, [t])
    assert len(img.shape) == 5
    channels, w, h, d, frames = img.shape

    sample = np.moveaxis(img, 0, 3)
    sample = sample.reshape((w, h, d, channels * frames))

    # TODO: rather swap axis of parameters than data
    sample = np.swapaxes(sample, 0, 2)

    return sitk.GetImageFromArray(sample)


def sitk_to_np(sitk_img: sitk.Image,
               channels,
               frames=1) -> np.array:

    # shape (d, w, h, c*f)
    img = sitk.GetArrayFromImage(sitk_img)

    if len(img.shape) == 3:
        img = np.expand_dims(img, 3)

    assert channels * frames == img.shape[-1], (f'Number of channels ({channels}) and frames ({frames})'
                                                f'do not correspond to the sitk vector size {img.shape[-1]}')

    # split channels and frames
    w, h, d = img.shape[:3]
    img = img.reshape((w, h, d, channels, frames))

    img = np.swapaxes(img, 0, 2)
    img = np.moveaxis(img, 3, 0)

    # shape (c, w, h, d, f)
    return img
