"""

This file contains functions to handle keras/tensorflow tensors.
    - blurring_sigma_for_downsampling
    - gaussian_kernel
    - resample_tensor
    - expand_dims


If you use this code, please cite the first SynthSeg paper:
https://github.com/BBillot/lab2im/blob/master/bibtex.bib

Copyright 2020 Benjamin Billot

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in
compliance with the License. You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is
distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied. See the License for the specific language governing permissions and limitations under the
License.

"""

# python imports
import numpy as np
import tensorflow as tf
import keras.layers as KL
import keras.backend as K
from itertools import combinations

# project imports
from lamareg.ext.lab2im import utils

# third-party imports
import lamareg.ext.neuron.layers as nrn_layers
from lamareg.ext.neuron.utils import volshape_to_meshgrid


def blurring_sigma_for_downsampling(
    current_res, downsample_res, mult_coef=None, thickness=None
):
    """Compute standard deviations of 1d gaussian masks for image blurring before downsampling.
    :param downsample_res: resolution to downsample to. Can be a 1d numpy array or list, or a tensor.
    :param current_res: resolution of the volume before downsampling.
    Can be a 1d numpy array or list or tensor of the same length as downsample res.
    :param mult_coef: (optional) multiplicative coefficient for the blurring kernel. Default is 0.75.
    :param thickness: (optional) slice thickness in each dimension. Must be the same type as downsample_res.
    :return: standard deviation of the blurring masks given as the same type as downsample_res (list or tensor).
    """

    if not tf.is_tensor(downsample_res):

        # get blurring resolution (min between downsample_res and thickness)
        current_res = np.array(current_res)
        downsample_res = np.array(downsample_res)
        if thickness is not None:
            downsample_res = np.minimum(downsample_res, np.array(thickness))

        # get std deviation for blurring kernels
        if mult_coef is None:
            sigma = 0.75 * downsample_res / current_res
            sigma[downsample_res == current_res] = 0.5
        else:
            sigma = mult_coef * downsample_res / current_res
        sigma[downsample_res == 0] = 0

    else:

        # reformat data resolution at which we blur
        if thickness is not None:
            down_res = KL.Lambda(lambda x: tf.math.minimum(x[0], x[1]))(
                [downsample_res, thickness]
            )
        else:
            down_res = downsample_res

        # get std deviation for blurring kernels
        if mult_coef is None:
            sigma = KL.Lambda(
                lambda x: tf.where(
                    tf.math.equal(
                        x, tf.convert_to_tensor(current_res, dtype="float32")
                    ),
                    0.5,
                    0.75 * x / tf.convert_to_tensor(current_res, dtype="float32"),
                )
            )(down_res)
        else:
            sigma = KL.Lambda(
                lambda x: mult_coef
                * x
                / tf.convert_to_tensor(current_res, dtype="float32")
            )(down_res)
        sigma = KL.Lambda(lambda x: tf.where(tf.math.equal(x[0], 0.0), 0.0, x[1]))(
            [down_res, sigma]
        )

    return sigma


def gaussian_kernel(sigma, max_sigma=None, blur_range=None, separable=True):
    """Build gaussian kernels of the specified standard deviation. The outputs are given as tensorflow tensors.
    :param sigma: standard deviation of the tensors. Can be given as a list/numpy array or as tensors. In each case,
    sigma must have the same length as the number of dimensions of the volume that will be blurred with the output
    tensors (e.g. sigma must have 3 values for 3D volumes).
    :param max_sigma:
    :param blur_range:
    :param separable:
    :return:
    """
    # convert sigma into a tensor
    if not tf.is_tensor(sigma):
        sigma_tens = tf.convert_to_tensor(
            utils.reformat_to_list(sigma), dtype="float32"
        )
    else:
        assert (
            max_sigma is not None
        ), "max_sigma must be provided when sigma is given as a tensor"
        sigma_tens = sigma
    shape = sigma_tens.get_shape().as_list()

    # get n_dims and batchsize
    if shape[0] is not None:
        n_dims = shape[0]
        batchsize = None
    else:
        n_dims = shape[1]
        batchsize = tf.split(tf.shape(sigma_tens), [1, -1])[0]

    # reformat max_sigma
    if max_sigma is not None:  # dynamic blurring
        max_sigma = np.array(utils.reformat_to_list(max_sigma, length=n_dims))
    else:  # sigma is fixed
        max_sigma = np.array(utils.reformat_to_list(sigma, length=n_dims))

    # randomise the burring std dev and/or split it between dimensions
    if blur_range is not None:
        if blur_range != 1:
            sigma_tens = sigma_tens * tf.random.uniform(
                tf.shape(sigma_tens), minval=1 / blur_range, maxval=blur_range
            )

    # get size of blurring kernels
    windowsize = np.int32(np.ceil(2.5 * max_sigma) / 2) * 2 + 1

    if separable:

        split_sigma = tf.split(sigma_tens, [1] * n_dims, axis=-1)

        kernels = list()
        comb = np.array(list(combinations(list(range(n_dims)), n_dims - 1))[::-1])
        for i, wsize in enumerate(windowsize):

            if wsize > 1:

                # build meshgrid and replicate it along batch dim if dynamic blurring
                locations = tf.cast(tf.range(0, wsize), "float32") - (wsize - 1) / 2
                if batchsize is not None:
                    locations = tf.tile(
                        tf.expand_dims(locations, axis=0),
                        tf.concat(
                            [
                                batchsize,
                                tf.ones(tf.shape(tf.shape(locations)), dtype="int32"),
                            ],
                            axis=0,
                        ),
                    )
                    comb[i] += 1

                # compute gaussians
                exp_term = -K.square(locations) / (2 * split_sigma[i] ** 2)
                g = tf.exp(exp_term - tf.math.log(np.sqrt(2 * np.pi) * split_sigma[i]))
                g = g / tf.reduce_sum(g)

                for axis in comb[i]:
                    g = tf.expand_dims(g, axis=axis)
                kernels.append(tf.expand_dims(tf.expand_dims(g, -1), -1))

            else:
                kernels.append(None)

    else:

        # build meshgrid
        mesh = [
            tf.cast(f, "float32")
            for f in volshape_to_meshgrid(windowsize, indexing="ij")
        ]
        diff = tf.stack(
            [mesh[f] - (windowsize[f] - 1) / 2 for f in range(len(windowsize))], axis=-1
        )

        # replicate meshgrid to batch size and reshape sigma_tens
        if batchsize is not None:
            diff = tf.tile(
                tf.expand_dims(diff, axis=0),
                tf.concat(
                    [batchsize, tf.ones(tf.shape(tf.shape(diff)), dtype="int32")],
                    axis=0,
                ),
            )
            for i in range(n_dims):
                sigma_tens = tf.expand_dims(sigma_tens, axis=1)
        else:
            for i in range(n_dims):
                sigma_tens = tf.expand_dims(sigma_tens, axis=0)

        # compute gaussians
        sigma_is_0 = tf.equal(sigma_tens, 0)
        exp_term = -tf.square(diff) / (
            2 * tf.where(sigma_is_0, tf.ones_like(sigma_tens), sigma_tens) ** 2
        )
        norms = exp_term - tf.math.log(
            tf.where(
                sigma_is_0, tf.ones_like(sigma_tens), np.sqrt(2 * np.pi) * sigma_tens
            )
        )
        kernels = tf.reduce_sum(norms, -1)
        kernels = tf.exp(kernels)
        kernels /= tf.reduce_sum(kernels)
        kernels = tf.expand_dims(tf.expand_dims(kernels, -1), -1)

    return kernels


def sobel_kernels(n_dims):
    """Returns sobel kernels to compute spatial derivative on image of n dimensions."""

    in_dir = tf.convert_to_tensor([1, 0, -1], dtype="float32")
    orthogonal_dir = tf.convert_to_tensor([1, 2, 1], dtype="float32")
    comb = np.array(list(combinations(list(range(n_dims)), n_dims - 1))[::-1])

    list_kernels = list()
    for dim in range(n_dims):

        sublist_kernels = list()
        for axis in range(n_dims):

            kernel = in_dir if axis == dim else orthogonal_dir
            for i in comb[axis]:
                kernel = tf.expand_dims(kernel, axis=i)
            sublist_kernels.append(tf.expand_dims(tf.expand_dims(kernel, -1), -1))

        list_kernels.append(sublist_kernels)

    return list_kernels


def unit_kernel(dist_threshold, n_dims, max_dist_threshold=None):
    """Build kernel with values of 1 for voxel at a distance < dist_threshold from the center, and 0 otherwise.
    The outputs are given as tensorflow tensors.
    :param dist_threshold: maximum distance from the center until voxel will have a value of 1. Can be a tensor of size
    (batch_size, 1), or a float.
    :param n_dims: dimension of the kernel to return (excluding batch and channel dimensions).
    :param max_dist_threshold: if distance_threshold is a tensor, max_dist_threshold must be given. It represents the
    maximum value that will be passed to dist_threshold. Must be a float.
    """

    # convert dist_threshold into a tensor
    if not tf.is_tensor(dist_threshold):
        dist_threshold_tens = tf.convert_to_tensor(
            utils.reformat_to_list(dist_threshold), dtype="float32"
        )
    else:
        assert (
            max_dist_threshold is not None
        ), "max_sigma must be provided when dist_threshold is given as a tensor"
        dist_threshold_tens = tf.cast(dist_threshold, "float32")
    shape = dist_threshold_tens.get_shape().as_list()

    # get batchsize
    batchsize = (
        None
        if shape[0] is not None
        else tf.split(tf.shape(dist_threshold_tens), [1, -1])[0]
    )

    # set max_dist_threshold into an array
    if (
        max_dist_threshold is None
    ):  # dist_threshold is fixed (i.e. dist_threshold will not change at each mini-batch)
        max_dist_threshold = dist_threshold

    # get size of blurring kernels
    windowsize = np.array([max_dist_threshold * 2 + 1] * n_dims, dtype="int32")

    # build tensor representing the distance from the centre
    mesh = [
        tf.cast(f, "float32") for f in volshape_to_meshgrid(windowsize, indexing="ij")
    ]
    dist = tf.stack(
        [mesh[f] - (windowsize[f] - 1) / 2 for f in range(len(windowsize))], axis=-1
    )
    dist = tf.sqrt(tf.reduce_sum(tf.square(dist), axis=-1))

    # replicate distance to batch size and reshape sigma_tens
    if batchsize is not None:
        dist = tf.tile(
            tf.expand_dims(dist, axis=0),
            tf.concat(
                [batchsize, tf.ones(tf.shape(tf.shape(dist)), dtype="int32")], axis=0
            ),
        )
        for i in range(n_dims - 1):
            dist_threshold_tens = tf.expand_dims(dist_threshold_tens, axis=1)
    else:
        for i in range(n_dims - 1):
            dist_threshold_tens = tf.expand_dims(dist_threshold_tens, axis=0)

    # build final kernel by thresholding distance tensor
    kernel = tf.where(
        tf.less_equal(dist, dist_threshold_tens),
        tf.ones_like(dist),
        tf.zeros_like(dist),
    )
    kernel = tf.expand_dims(tf.expand_dims(kernel, -1), -1)

    return kernel


def resample_tensor(
    tensor,
    resample_shape,
    interp_method="linear",
    subsample_res=None,
    volume_res=None,
    build_reliability_map=False,
):
    """This function resamples a volume to resample_shape. It does not apply any pre-filtering.
    A prior downsampling step can be added if subsample_res is specified. In this case, volume_res should also be
    specified, in order to calculate the downsampling ratio. A reliability map can also be returned to indicate which
    slices were interpolated during resampling from the downsampled to final tensor.
    :param tensor: tensor
    :param resample_shape: list or numpy array of size (n_dims,)
    :param interp_method: (optional) interpolation method for resampling, 'linear' (default) or 'nearest'
    :param subsample_res: (optional) if not None, this triggers a downsampling of the volume, prior to the resampling
    step. List or numpy array of size (n_dims,). Default si None.
    :param volume_res: (optional) if subsample_res is not None, this should be provided to compute downsampling ratio.
    list or numpy array of size (n_dims,). Default is None.
    :param build_reliability_map: whether to return reliability map along with the resampled tensor. This map indicates
    which slices of the resampled tensor are interpolated (0=interpolated, 1=real slice, in between=degree of realness).
    :return: resampled volume, with reliability map if necessary.
    """

    # reformat resolutions to lists
    subsample_res = utils.reformat_to_list(subsample_res)
    volume_res = utils.reformat_to_list(volume_res)
    n_dims = len(resample_shape)

    # downsample image
    tensor_shape = tensor.get_shape().as_list()[1:-1]
    downsample_shape = tensor_shape  # will be modified if we actually downsample

    if subsample_res is not None:
        assert (
            volume_res is not None
        ), "volume_res must be given when providing a subsampling resolution."
        assert len(subsample_res) == len(volume_res), (
            "subsample_res and volume_res must have the same length, "
            "had {0}, and {1}".format(len(subsample_res), len(volume_res))
        )
        if subsample_res != volume_res:

            # get shape at which we downsample
            downsample_shape = [
                int(tensor_shape[i] * volume_res[i] / subsample_res[i])
                for i in range(n_dims)
            ]

            # downsample volume
            tensor._keras_shape = tuple(tensor.get_shape().as_list())
            tensor = nrn_layers.Resize(size=downsample_shape, interp_method="nearest")(
                tensor
            )

    # resample image at target resolution
    if (
        resample_shape != downsample_shape
    ):  # if we didn't downsample downsample_shape = tensor_shape
        tensor._keras_shape = tuple(tensor.get_shape().as_list())
        tensor = nrn_layers.Resize(size=resample_shape, interp_method=interp_method)(
            tensor
        )

    # compute reliability maps if necessary and return results
    if build_reliability_map:

        # compute maps only if we downsampled
        if downsample_shape != tensor_shape:

            # compute upsampling factors
            upsampling_factors = np.array(resample_shape) / np.array(downsample_shape)

            # build reliability map
            reliability_map = 1
            for i in range(n_dims):
                loc_float = np.arange(0, resample_shape[i], upsampling_factors[i])
                loc_floor = np.int32(np.floor(loc_float))
                loc_ceil = np.int32(np.clip(loc_floor + 1, 0, resample_shape[i] - 1))
                tmp_reliability_map = np.zeros(resample_shape[i])
                tmp_reliability_map[loc_floor] = 1 - (loc_float - loc_floor)
                tmp_reliability_map[loc_ceil] = tmp_reliability_map[loc_ceil] + (
                    loc_float - loc_floor
                )
                shape = [1, 1, 1]
                shape[i] = resample_shape[i]
                reliability_map = reliability_map * np.reshape(
                    tmp_reliability_map, shape
                )
            shape = KL.Lambda(lambda x: tf.shape(x))(tensor)
            mask = KL.Lambda(
                lambda x: tf.reshape(
                    tf.convert_to_tensor(reliability_map, dtype="float32"), shape=x
                )
            )(shape)

        # otherwise just return an all-one tensor
        else:
            mask = KL.Lambda(lambda x: tf.ones_like(x))(tensor)

        return tensor, mask

    else:
        return tensor


def expand_dims(tensor, axis=0):
    """Expand the dimensions of the input tensor along the provided axes (given as an integer or a list)."""
    axis = utils.reformat_to_list(axis)
    for ax in axis:
        tensor = tf.expand_dims(tensor, axis=ax)
    return tensor
