# -*- coding: utf-8 -*-
from __future__ import absolute_import, print_function

import numpy as np
import tensorflow as tf
from tensorflow.contrib.layers.python.layers import regularizers

from niftynet.engine.application_initializer import GlorotUniform
from niftynet.layer.convolution import ConvolutionalLayer as Conv
from niftynet.layer.downsample_res_block import DownBlock as DownRes
from niftynet.layer.grid_warper import _create_affine_features
from niftynet.layer.layer_util import infer_spatial_rank, check_spatial_dims
from niftynet.layer.linear_resize import LinearResizeLayer as Resize
from niftynet.layer.spatial_gradient import SpatialGradientLayer as ImgGrad
from niftynet.layer.upsample_res_block import UpBlock as UpRes
from niftynet.network.base_net import BaseNet


class INetDense(BaseNet):
    def __init__(self,
                 decay=0.0,
                 smoothing=0,
                 disp_w_initializer=None,
                 disp_b_initializer=None,
                 acti_func='relu',
                 multi_scale_fusion=True,
                 name='inet-dense'):
        """
        The network estimates dense displacement fields from a pair
        of moving and fixed images:

            Hu et al., Label-driven weakly-supervised learning for
            multimodal deformable image registration, arXiv:1711.01666
            https://arxiv.org/abs/1711.01666

            Hu et al., Weakly-Supervised Convolutional Neural Networks for
            Multimodal Image Registration, Medical Image Analysis (2018)
            https://doi.org/10.1016/j.media.2018.07.002

        see also:
            https://github.com/YipengHu/label-reg

        :param decay:
        :param smoothing:
        :param disp_w_initializer: initialisation of the displacement fields
        :param disp_b_initializer: initialisation of the dis
        :param acti_func:
        :param multi_scale_fusion: True/False indicating whether to use
            multiscale feature fusion.
        :param name:
        """
        BaseNet.__init__(self, name=name)

        # self.fea = [40, 80, 160, 320, 640]
        # self.fea = [32, 64, 128, 256, 512]
        self.fea = [30, 60, 120, 240, 480]
        # self.fea = [16, 32, 64, 128, 256]
        self.k_conv = 3
        self.multi_scale_fusion = multi_scale_fusion

        self.down_res_param = {
            'w_initializer': GlorotUniform.get_instance(''),
            'w_regularizer': regularizers.l2_regularizer(decay),
            'acti_func': acti_func}

        self.up_res_param = {
            'acti_func': acti_func,
            'w_initializer': GlorotUniform.get_instance(''),
            'w_regularizer': regularizers.l2_regularizer(decay),
            'is_residual_upsampling': True,
            'type_string': 'bn_acti_conv'}

        # displacement initialiser & regulariser
        if disp_w_initializer is None:
            disp_b_initializer = tf.constant_initializer(0.0)
            #disp_w_initializer = tf.random_normal_initializer(0, 1e-4)
        if disp_b_initializer is None:
            disp_b_initializer = tf.constant_initializer(0.0)
            #disp_w_initializer = tf.random_normal_initializer(0, 0.0)
        self.disp_param = {
            'w_initializer': disp_w_initializer,
            'w_regularizer': regularizers.l2_regularizer(decay),
            'b_initializer': disp_b_initializer,
            'b_regularizer': None}

        if smoothing > 0:
            self.smoothing_func = _smoothing_func(smoothing)
        else:
            self.smoothing_func = None

    def layer_op(self,
                 fixed_image,
                 moving_image,
                 base_grid=None,
                 is_training=True,
                 **unused_kwargs):
        """

        :param fixed_image:
        :param moving_image:
        :param base_grid:
        :param is_training:
        :return: estimated dense displacement fields
        """

        spatial_rank = infer_spatial_rank(fixed_image)
        spatial_shape = fixed_image.get_shape().as_list()[1:-1]
        check_spatial_dims(fixed_image, lambda x: x % 16 == 0)

        #  resize the moving image to match the fixed
        moving_image = Resize(spatial_shape)(moving_image)
        img = tf.concat([moving_image, fixed_image], axis=-1)
        down_res_0, conv_0_0, _ = \
            DownRes(self.fea[0], kernel_size=7, **self.down_res_param)(img, is_training)
        down_res_1, conv_0_1, _ = \
            DownRes(self.fea[1], **self.down_res_param)(down_res_0, is_training)
        down_res_2, conv_0_2, _ = \
            DownRes(self.fea[2], **self.down_res_param)(down_res_1, is_training)
        down_res_3, conv_0_3, _ = \
            DownRes(self.fea[3], **self.down_res_param)(down_res_2, is_training)

        conv_4 = Conv(n_output_chns=self.fea[4],
                      kernel_size=self.k_conv,
                      **self.down_res_param)(down_res_3, is_training)

        up_res_0 = UpRes(self.fea[3], **self.up_res_param)(
            conv_4, conv_0_3, is_training)
        up_res_1 = UpRes(self.fea[2], **self.up_res_param)(
            up_res_0, conv_0_2, is_training)
        up_res_2 = UpRes(self.fea[1], **self.up_res_param)(
            up_res_1, conv_0_1, is_training)
        up_res_3 = UpRes(self.fea[0], **self.up_res_param)(
            up_res_2, conv_0_0, is_training)

        if self.multi_scale_fusion:
            output_list = [up_res_3, up_res_2, up_res_1, up_res_0, conv_4]
        else:
            output_list = [up_res_3]

        # converting all output layers to displacement fields 
        dense_fields = []
        for scale_out in output_list:
            field = Conv(n_output_chns=spatial_rank,
                         kernel_size=self.k_conv,
                         with_bias=True,
                         with_bn=False,
                         acti_func=None,
                         **self.disp_param)(scale_out)
            resized_field = Resize(new_size=spatial_shape)(field)
            dense_fields.append(resized_field)

        if base_grid is None:
            # adding a reference grid if it doesn't exist
            in_spatial_size = [None] * spatial_rank
            base_grid = _create_affine_features(output_shape=spatial_shape,
                                                source_shape=in_spatial_size)
            base_grid = np.asarray(base_grid[:-1])
            base_grid = np.reshape(
                base_grid.T, [-1] + spatial_shape + [spatial_rank])
            base_grid = tf.constant(base_grid, dtype=resized_field.dtype)

        if self.multi_scale_fusion and len(dense_fields) > 1:
            dense_field = tf.reduce_sum(dense_fields, axis=0)
        else:
            dense_field = dense_fields[0]

        # TODO filtering
        if self.smoothing_func is not None:
            dense_field = self.smoothing_func(dense_field, spatial_rank)

        tf.add_to_collection('bending_energy',
                             _computing_bending_energy(dense_field))
        tf.add_to_collection('gradient_norm',
                             _computing_gradient_norm(dense_field))

        dense_field = dense_field + base_grid
        return dense_field


def _get_smoothing_kernel(sigma, spatial_rank):
    # sigma defined in voxel not in freeform deformation grid
    if sigma <= 0:
        raise NotImplementedError
    tail = int(sigma * 2)
    if spatial_rank == 2:
        x, y = np.mgrid[-tail:tail + 1, -tail:tail + 1]
        g = np.exp(-0.5 * (x * x + y * y) / sigma * sigma)
    elif spatial_rank == 3:
        x, y, z = np.mgrid[-tail:tail + 1, -tail:tail + 1, -tail:tail + 1]
        g = np.exp(-0.5 * (x * x + y * y + z * z) / sigma * sigma)
    else:
        raise NotImplementedError
    return g / g.sum()


def _smoothing_func(sigma):
    def smoothing(dense_field, spatial_rank):
        kernel = _get_smoothing_kernel(sigma, spatial_rank)
        kernel = tf.constant(kernel, dtype=dense_field.dtype)
        kernel = tf.expand_dims(kernel, axis=-1)
        kernel = tf.expand_dims(kernel, axis=-1)
        smoothed = [
            tf.nn.convolution(tf.expand_dims(coord, axis=-1), kernel, 'SAME')
            for coord in tf.unstack(dense_field, axis=-1)]
        return tf.concat(smoothed, axis=-1)

    return smoothing


def _computing_bending_energy(displacement):
    spatial_rank = infer_spatial_rank(displacement)
    if spatial_rank == 2:
        return _computing_bending_energy_2d(displacement)
    if spatial_rank == 3:
        return _computing_bending_energy_3d(displacement)
    raise NotImplementedError(
        "Not implmented: bending energy for {}-d input".format(spatial_rank))


def _computing_bending_energy_2d(displacement):
    dTdx = ImgGrad(spatial_axis=0)(displacement)
    dTdy = ImgGrad(spatial_axis=1)(displacement)

    dTdxx = ImgGrad(spatial_axis=0)(dTdx)
    dTdyy = ImgGrad(spatial_axis=1)(dTdy)
    dTdxy = ImgGrad(spatial_axis=1)(dTdx)

    energy = tf.reduce_mean([dTdxx * dTdxx, dTdyy * dTdyy, 2 * dTdxy * dTdxy])
    return energy


def _computing_bending_energy_3d(displacement):
    dTdx = ImgGrad(spatial_axis=0)(displacement)
    dTdy = ImgGrad(spatial_axis=1)(displacement)
    dTdz = ImgGrad(spatial_axis=2)(displacement)

    dTdxx = ImgGrad(spatial_axis=0)(dTdx)
    dTdyy = ImgGrad(spatial_axis=1)(dTdy)
    dTdzz = ImgGrad(spatial_axis=2)(dTdz)

    dTdxy = ImgGrad(spatial_axis=1)(dTdx)
    dTdyz = ImgGrad(spatial_axis=2)(dTdy)
    dTdxz = ImgGrad(spatial_axis=2)(dTdx)

    energy = tf.reduce_mean(
        [dTdxx * dTdxx, dTdyy * dTdyy, dTdzz * dTdzz,
         2 * dTdxy * dTdxy, 2 * dTdxz * dTdxz, 2 * dTdyz * dTdyz])
    return energy


def _computing_gradient_norm(displacement, flag_L1=False):
    norms = []
    for spatial_ind in range(infer_spatial_rank(displacement)):
        dTdt = ImgGrad(spatial_axis=spatial_ind)(displacement)
        if flag_L1:
            norms.append(tf.abs(dTdt))
        else:
            norms.append(dTdt * dTdt)
    return tf.reduce_mean(norms)
