"""

tensorflow/keras utilities for the neuron project

If you use this code, please cite 

Dalca AV, Guttag J, Sabuncu MR

Anatomical Priors in Convolutional Networks for Unsupervised Biomedical Segmentation, 

CVPR 2018

or for the transformation/integration functions:

Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration

Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu

MICCAI 2018.

Contact: adalca [at] csail [dot] mit [dot] edu

License: GPLv3

"""

# third party

import tensorflow as tf

from keras import backend as K

from keras.layers import Layer

from copy import deepcopy

# local

from ext.neuron.utils import transform, resize, integrate_vec, affine_to_shift, combine_non_linear_and_aff_to_shift

class SpatialTransformer(Layer):

    """

    N-D Spatial Transformer Tensorflow / Keras Layer

    The Layer can handle both affine and dense transforms. 

    Both transforms are meant to give a 'shift' from the current position.

    Therefore, a dense transform gives displacements (not absolute locations) at each voxel,

    and an affine transform gives the *difference* of the affine matrix from 

    the identity matrix.

    If you find this function useful, please cite:

      Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration

      Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu

      MICCAI 2018.

    Originally, this code was based on voxelmorph code, which 

    was in turn transformed to be dense with the help of (affine) STN code 

    via https://github.com/kevinzakka/spatial-transformer-network

    Since then, we've re-written the code to be generalized to any 

    dimensions, and along the way wrote grid and interpolation functions

    """

    def __init__(self,

                 interp_method='linear',

                 indexing='ij',

                 single_transform=False,

                 **kwargs):

        """

        Parameters: 

            interp_method: 'linear' or 'nearest'

            single_transform: whether a single transform supplied for the whole batch

            indexing (default: 'ij'): 'ij' (matrix) or 'xy' (cartesian)

                'xy' indexing will have the first two entries of the flow 

                (along last axis) flipped compared to 'ij' indexing

        """

        self.interp_method = interp_method

        self.ndims = None

        self.inshape = None

        self.single_transform = single_transform

        self.is_affine = list()

        assert indexing in ['ij', 'xy'], "indexing has to be 'ij' (matrix) or 'xy' (cartesian)"

        self.indexing = indexing

        super(self.__class__, self).__init__(**kwargs)

    def get_config(self):

        config = super().get_config()

        config["interp_method"] = self.interp_method

        config["indexing"] = self.indexing

        config["single_transform"] = self.single_transform

        return config

    def build(self, input_shape):

        """

        input_shape should be a list for two inputs:

        input1: image.

        input2: list of transform Tensors

            if affine:

                should be an N+1 x N+1 matrix

                *or* a N+1*N+1 tensor (which will be reshaped to N x (N+1) and an identity row added)

            if not affine:

                should be a *vol_shape x N

        """

        if len(input_shape) > 3:

            raise Exception('Spatial Transformer must be called on a list of min length 2 and max length 3.'

                            'First argument is the image followed by the affine and non linear transforms.')

        # set up number of dimensions

        self.ndims = len(input_shape[0]) - 2

        self.inshape = input_shape

        trf_shape = [trans_shape[1:] for trans_shape in input_shape[1:]]

        for (i, shape) in enumerate(trf_shape):

            # the transform is an affine iff:

            # it's a 1D Tensor [dense transforms need to be at least ndims + 1]

            # it's a 2D Tensor and shape == [N+1, N+1].

            self.is_affine.append(len(shape) == 1 or

                                  (len(shape) == 2 and all([f == (self.ndims + 1) for f in shape])))

            # check sizes

            if self.is_affine[i] and len(shape) == 1:

                ex = self.ndims * (self.ndims + 1)

                if shape[0] != ex:

                    raise Exception('Expected flattened affine of len %d but got %d' % (ex, shape[0]))

            if not self.is_affine[i]:

                if shape[-1] != self.ndims:

                    raise Exception('Offset flow field size expected: %d, found: %d' % (self.ndims, shape[-1]))

        # confirm built

        self.built = True

    def call(self, inputs, **kwargs):

        """

        Parameters

            inputs: list with several entries: the volume followed by the transforms

        """

        # check shapes

        assert 1 < len(inputs) < 4, "inputs has to be len 2 or 3, found: %d" % len(inputs)

        vol = inputs[0]

        trf = inputs[1:]

        # necessary for multi_gpu models...

        vol = K.reshape(vol, [-1, *self.inshape[0][1:]])

        for i in range(len(trf)):

            trf[i] = K.reshape(trf[i], [-1, *self.inshape[i+1][1:]])

        # reorder transforms, non-linear first and affine second

        ind_nonlinear_linear = [i[0] for i in sorted(enumerate(self.is_affine), key=lambda x:x[1])]

        self.is_affine = [self.is_affine[i] for i in ind_nonlinear_linear]

        self.inshape = [self.inshape[i] for i in ind_nonlinear_linear]

        trf = [trf[i] for i in ind_nonlinear_linear]

        # go from affine to deformation field

        if len(trf) == 1:

            trf = trf[0]

            if self.is_affine[0]:

                trf = tf.map_fn(lambda x: self._single_aff_to_shift(x, vol.shape[1:-1]), trf, dtype=tf.float32)

        # combine non-linear and affine to obtain a single deformation field

        elif len(trf) == 2:

            trf = tf.map_fn(lambda x: self._non_linear_and_aff_to_shift(x, vol.shape[1:-1]), trf, dtype=tf.float32)

        # prepare location shift

        if self.indexing == 'xy':  # shift the first two dimensions

            trf_split = tf.split(trf, trf.shape[-1], axis=-1)

            trf_lst = [trf_split[1], trf_split[0], *trf_split[2:]]

            trf = tf.concat(trf_lst, -1)

        # map transform across batch

        if self.single_transform:

            return tf.map_fn(self._single_transform, [vol, trf[0, :]], dtype=tf.float32)

        else:

            return tf.map_fn(self._single_transform, [vol, trf], dtype=tf.float32)

    def _single_aff_to_shift(self, trf, volshape):

        if len(trf.shape) == 1:  # go from vector to matrix

            trf = tf.reshape(trf, [self.ndims, self.ndims + 1])

        return affine_to_shift(trf, volshape, shift_center=True)

    def _non_linear_and_aff_to_shift(self, trf, volshape):

        if len(trf[1].shape) == 1:  # go from vector to matrix

            trf[1] = tf.reshape(trf[1], [self.ndims, self.ndims + 1])

        return combine_non_linear_and_aff_to_shift(trf, volshape, shift_center=True)

    def _single_transform(self, inputs):

        return transform(inputs[0], inputs[1], interp_method=self.interp_method)

class VecInt(Layer):

    """

    Vector Integration Layer

    Enables vector integration via several methods 

    (ode or quadrature for time-dependent vector fields, 

    scaling and squaring for stationary fields)

    If you find this function useful, please cite:

      Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration

      Adrian V. Dalca, Guha Balakrishnan, John Guttag, Mert R. Sabuncu

      MICCAI 2018.

    """

    def __init__(self, indexing='ij', method='ss', int_steps=7, out_time_pt=1,

                 ode_args=None,

                 odeint_fn=None, **kwargs):

        """        

        Parameters:

            method can be any of the methods in neuron.utils.integrate_vec

            indexing can be 'xy' (switches first two dimensions) or 'ij'

            int_steps is the number of integration steps

            out_time_pt is time point at which to output if using odeint integration

        """

        assert indexing in ['ij', 'xy'], "indexing has to be 'ij' (matrix) or 'xy' (cartesian)"

        self.indexing = indexing

        self.method = method

        self.int_steps = int_steps

        self.inshape = None

        self.out_time_pt = out_time_pt

        self.odeint_fn = odeint_fn  # if none then will use a tensorflow function

        self.ode_args = ode_args

        if ode_args is None:

            self.ode_args = {'rtol': 1e-6, 'atol': 1e-12}

        super(self.__class__, self).__init__(**kwargs)

    def get_config(self):

        config = super().get_config()

        config["indexing"] = self.indexing

        config["method"] = self.method

        config["int_steps"] = self.int_steps

        config["out_time_pt"] = self.out_time_pt

        config["ode_args"] = self.ode_args

        config["odeint_fn"] = self.odeint_fn

        return config

    def build(self, input_shape):

        # confirm built

        self.built = True

        trf_shape = input_shape

        if isinstance(input_shape[0], (list, tuple)):

            trf_shape = input_shape[0]

        self.inshape = trf_shape

        if trf_shape[-1] != len(trf_shape) - 2:

            raise Exception('transform ndims %d does not match expected ndims %d' % (trf_shape[-1], len(trf_shape) - 2))

    def call(self, inputs, **kwargs):

        if not isinstance(inputs, (list, tuple)):

            inputs = [inputs]

        loc_shift = inputs[0]

        # necessary for multi_gpu models...

        loc_shift = K.reshape(loc_shift, [-1, *self.inshape[1:]])

        # prepare location shift

        if self.indexing == 'xy':  # shift the first two dimensions

            loc_shift_split = tf.split(loc_shift, loc_shift.shape[-1], axis=-1)

            loc_shift_lst = [loc_shift_split[1], loc_shift_split[0], *loc_shift_split[2:]]

            loc_shift = tf.concat(loc_shift_lst, -1)

        if len(inputs) > 1:

            assert self.out_time_pt is None, 'out_time_pt should be None if providing batch_based out_time_pt'

        # map transform across batch

        out = tf.map_fn(self._single_int, [loc_shift] + inputs[1:], dtype=tf.float32)

        return out

    def _single_int(self, inputs):

        vel = inputs[0]

        out_time_pt = self.out_time_pt

        if len(inputs) == 2:

            out_time_pt = inputs[1]

        return integrate_vec(vel, method=self.method,

                             nb_steps=self.int_steps,

                             ode_args=self.ode_args,

                             out_time_pt=out_time_pt,

                             odeint_fn=self.odeint_fn)

class Resize(Layer):

    """

    N-D Resize Tensorflow / Keras Layer

    Note: this is not re-shaping an existing volume, but resizing, like scipy's "Zoom"

    If you find this function useful, please cite:

    Anatomical Priors in Convolutional Networks for Unsupervised Biomedical Segmentation,Dalca AV, Guttag J, Sabuncu MR

    CVPR 2018

    Since then, we've re-written the code to be generalized to any 

    dimensions, and along the way wrote grid and interpolation functions

    """

    def __init__(self,

                 zoom_factor=None,

                 size=None,

                 interp_method='linear',

                 **kwargs):

        """

        Parameters: 

            interp_method: 'linear' or 'nearest'

                'xy' indexing will have the first two entries of the flow 

                (along last axis) flipped compared to 'ij' indexing

        """

        self.zoom_factor = zoom_factor

        self.size = list(size)

        self.zoom_factor0 = None

        self.size0 = None

        self.interp_method = interp_method

        self.ndims = None

        self.inshape = None

        super(Resize, self).__init__(**kwargs)

    def get_config(self):

        config = super().get_config()

        config["zoom_factor"] = self.zoom_factor

        config["size"] = self.size

        config["interp_method"] = self.interp_method

        return config

    def build(self, input_shape):

        """

        input_shape should be an element of list of one inputs:

        input1: volume

                should be a *vol_shape x N

        """

        if isinstance(input_shape[0], (list, tuple)) and len(input_shape) > 1:

            raise Exception('Resize must be called on a list of length 1.')

        if isinstance(input_shape[0], (list, tuple)):

            input_shape = input_shape[0]

        # set up number of dimensions

        self.ndims = len(input_shape) - 2

        self.inshape = input_shape

        # check zoom_factor

        if isinstance(self.zoom_factor, float):

            self.zoom_factor0 = [self.zoom_factor] * self.ndims

        elif self.zoom_factor is None:

            self.zoom_factor0 = [0] * self.ndims

        elif isinstance(self.zoom_factor, (list, tuple)):

            self.zoom_factor0 = deepcopy(self.zoom_factor)

            assert len(self.zoom_factor0) == self.ndims, \

                'zoom factor length {} does not match number of dimensions {}'.format(len(self.zoom_factor), self.ndims)

        else:

            raise Exception('zoom_factor should be an int or a list/tuple of int (or None if size is not set to None)')

        # check size

        if isinstance(self.size, int):

            self.size0 = [self.size] * self.ndims

        elif self.size is None:

            self.size0 = [0] * self.ndims

        elif isinstance(self.size, (list, tuple)):

            self.size0 = deepcopy(self.size)

            assert len(self.size0) == self.ndims, \

                'size length {} does not match number of dimensions {}'.format(len(self.size0), self.ndims)

        else:

            raise Exception('size should be an int or a list/tuple of int (or None if zoom_factor is not set to None)')

        # confirm built

        self.built = True

        super(Resize, self).build(input_shape)  # Be sure to call this somewhere!

    def call(self, inputs, **kwargs):

        """

        Parameters

            inputs: volume or list of one volume

        """

        # check shapes

        if isinstance(inputs, (list, tuple)):

            assert len(inputs) == 1, "inputs has to be len 1. found: %d" % len(inputs)

            vol = inputs[0]

        else:

            vol = inputs

        # necessary for multi_gpu models...

        vol = K.reshape(vol, [-1, *self.inshape[1:]])

        # set value of missing size or zoom_factor

        if not any(self.zoom_factor0):

            self.zoom_factor0 = [self.size0[i] / self.inshape[i+1] for i in range(self.ndims)]

        else:

            self.size0 = [int(self.inshape[f+1] * self.zoom_factor0[f]) for f in range(self.ndims)]

        # map transform across batch

        return tf.map_fn(self._single_resize, vol, dtype=vol.dtype)

    def compute_output_shape(self, input_shape):

        output_shape = [input_shape[0]]

        output_shape += [int(input_shape[1:-1][f] * self.zoom_factor0[f]) for f in range(self.ndims)]

        output_shape += [input_shape[-1]]

        return tuple(output_shape)

    def _single_resize(self, inputs):

        return resize(inputs, self.zoom_factor0, self.size0, interp_method=self.interp_method)

# Zoom naming of resize, to match scipy's naming

Zoom = Resize

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

# "Local" layers -- layers with parameters at each voxel

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

class LocalBias(Layer):

    """ 

    Local bias layer: each pixel/voxel has its own bias operation (one parameter)

    out[v] = in[v] + b

    """

    def __init__(self, my_initializer='RandomNormal', biasmult=1.0, **kwargs):

        self.initializer = my_initializer

        self.biasmult = biasmult

        self.kernel = None

        super(LocalBias, self).__init__(**kwargs)

    def get_config(self):

        config = super().get_config()

        config["my_initializer"] = self.initializer

        config["biasmult"] = self.biasmult

        return config

    def build(self, input_shape):

        # Create a trainable weight variable for this layer.

        self.kernel = self.add_weight(name='kernel',

                                      shape=input_shape[1:],

                                      initializer=self.initializer,

                                      trainable=True)

        super(LocalBias, self).build(input_shape)  # Be sure to call this somewhere!

    def call(self, x, **kwargs):

        return x + self.kernel * self.biasmult  # weights are difference from input

    def compute_output_shape(self, input_shape):

        return input_shape