"""
This file contains functions for tensorflow neural networks in general.
"""
from operator import mul

import numpy as np
import tensorflow as tf
from typing import Union, Tuple, List, Dict
from tensorflow.python.framework import ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_image_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import math_ops

import vgg


def get_tensor_num_elements(tensor):
    # type: (tf.Tensor) -> int
    tensor_shape = map(lambda i: i.value, tensor.get_shape())
    return reduce(mul, tensor_shape, 1)


def concat_content_img_style_weight_mask_to_input(input_tensor, content_img_style_weight_mask):
    # type: (tf.Tensor, tf.Tensor) -> tf.Tensor
    """

    :param input_tensor: tensor with shape (batch_size, height, width, num_features)
    :param content_img_style_weight_mask: (batch_size, height, width, 1)
    :return: a tensor created by concatenating the input and the mask in their last dimension.
    """
    assert content_img_style_weight_mask is not None
    input_concatenated = tf.concat(3, (input_tensor, content_img_style_weight_mask))
    return input_concatenated


def spatial_batch_norm(input_layer, name='spatial_batch_norm', reuse=False):
    # type: (tf.Tensor, str, bool) -> tf.Tensor
    """
    Batch-normalizes the layer as in http://arxiv.org/abs/1502.03167
    This is important since it allows the different scales to talk to each other when they get joined.
    :param input_layer: tensor with shape (batch_size, height, width, num_features)
    :param name: name of the variable scope
    :param reuse: reuse variables if set to true. Otherwise create new ones.
    :return: a tensor with shape (batch_size, height, width, num_features) with the normalization applied.
    """
    with tf.variable_scope(name, reuse=reuse):
        mean, variance = tf.nn.moments(input_layer, [0, 1, 2])
        # NOTE: Tensorflow norm has some issues when the actual variance is near zero. I have to apply abs on it.
        variance = tf.abs(variance)
        variance_epsilon = 0.001
        num_channels = input_layer.get_shape().as_list()[3]
        scale = tf.get_variable('scale', [num_channels], tf.float32, tf.random_uniform_initializer())
        offset = tf.get_variable('offset', [num_channels], tf.float32, tf.constant_initializer())
        return_val = tf.nn.batch_normalization(input_layer, mean, variance, offset, scale, variance_epsilon, name=name)
        return return_val


def instance_norm(input_layer, name='instance_norm', reuse=False):
    # type: (tf.Tensor, str, bool) -> tf.Tensor
    """
    Instance-normalize the layer as in https://arxiv.org/abs/1607.08022
    :param input_layer: tensor with shape (batch_size, height, width, num_features)
    :param name: name of the variable scope
    :param reuse: reuse variables if set to true. Otherwise create new ones.
    :return: a tensor with shape (batch_size, height, width, num_features) with the normalization applied.
    """
    with tf.variable_scope(name, reuse=reuse):
        input_layers = tf.unpack(input_layer)
        return_val = []
        num_channels = input_layer.get_shape().as_list()[3]
        # The scale and offset variable is reused for all batches in this norm.
        scale = tf.get_variable('scale', [num_channels], tf.float32, tf.random_uniform_initializer())
        offset = tf.get_variable('offset', [num_channels], tf.float32, tf.constant_initializer())
        for l in input_layers:
            l = tf.expand_dims(l, 0)
            # NOTE: Tensorflow norm has some issues when the actual variance is near zero. I have to apply abs on it.
            mean, variance = tf.nn.moments(l, [0, 1, 2])
            variance = tf.abs(variance)
            variance_epsilon = 0.001
            return_val.append(
                tf.squeeze(tf.nn.batch_normalization(l, mean, variance, offset, scale, variance_epsilon, name=name),
                           [0]))
        return_val = tf.pack(return_val)
        return return_val


def conditional_instance_norm(input_layer, input_style_placeholder, name='conditional_instance_norm', reuse=False):
    # type: (tf.Tensor, tf.Tensor, str, bool) -> tf.Tensor
    """
    Instance-normalize the layer conditioned on the style as in https://arxiv.org/abs/1610.07629
    :param input_layer: tensor with shape (batch_size, height, width, num_features)
    :param input_style_placeholder: a one hot vector (1 x N tensor) with length N where N is the number of different
    style images.
    :param name: name of the variable scope
    :param reuse: reuse variables if set to true. Otherwise create new ones.
    :return: a tensor with shape (batch_size, height, width, num_features) with the normalization applied.
    """
    with tf.variable_scope(name, reuse=reuse):
        input_layers = tf.unpack(input_layer)
        return_val = []
        num_styles = input_style_placeholder.get_shape().as_list()[1]
        num_channels = input_layer.get_shape().as_list()[3]
        scale = tf.get_variable('scale', [num_styles, num_channels], tf.float32, tf.random_uniform_initializer())
        offset = tf.get_variable('offset', [num_styles, num_channels], tf.float32, tf.constant_initializer())
        scale_for_current_style = tf.matmul(input_style_placeholder, scale)
        offset_for_current_style = tf.matmul(input_style_placeholder, offset)
        for l in input_layers:
            l = tf.expand_dims(l, 0)
            # NOTE: Tensorflow norm has some issues when the actual variance is near zero. I have to apply abs on it.
            mean, variance = tf.nn.moments(l, [0, 1, 2])
            variance = tf.abs(variance)
            variance_epsilon = 0.001
            return_val.append(tf.squeeze(tf.nn.batch_normalization(
                l, mean, variance, offset_for_current_style, scale_for_current_style, variance_epsilon, name=name),
                [0]))
        return_val = tf.pack(return_val)
        return return_val


def gramian(layer):
    # type: (tf.Tensor) -> tf.Tensor
    """
    :param layer: tensor with shape (batch_size, height, width, num_features)
    :return: The gramian of the layer -- a tensor with dimension gramians of dimension (batches, channels, channels)
    """
    # Instead of iterating over #channels width by height matrices and computing similarity, we vectorize and compute
    # the entire gramian in a single matrix multiplication.
    _, height, width, number = map(lambda i: i.value, layer.get_shape())
    size = height * width * number
    layer_unpacked = tf.unpack(layer)
    grams = []
    for single_layer in layer_unpacked:
        feats = tf.reshape(single_layer, (-1, number))
        # Note: the normalization factor might be wrong. I've seen many different forms of normalization. The current
        # one works though.
        grams.append(tf.matmul(tf.transpose(feats), feats) / size)
    return tf.pack(grams)


def total_variation(image_batch):
    # type: (Union[tf.Tensor,tf.Variable]) -> tf.Tensor
    """
    :param image_batch: A 4D tensor of shape (batch_size, height, width, num_features)
    :return: The variation of the input represented as a scalar tensor.
    """
    batch_shape = image_batch.get_shape().as_list()
    batch_size = batch_shape[0]
    height = batch_shape[1]
    left = tf.slice(image_batch, [0, 0, 0, 0], [-1, height - 1, -1, -1])
    right = tf.slice(image_batch, [0, 1, 0, 0], [-1, -1, -1, -1])

    width = batch_shape[2]
    top = tf.slice(image_batch, [0, 0, 0, 0], [-1, -1, width - 1, -1])
    bottom = tf.slice(image_batch, [0, 0, 1, 0], [-1, -1, -1, -1])

    # left and right are 1 less wide than the original, top and bottom 1 less tall
    # In order to combine them, we take 1 off the height of left-right, and 1 off width of top-bottom
    vertical_diff = tf.slice(tf.sub(left, right), [0, 0, 0, 0], [-1, -1, width - 1, -1])
    horizontal_diff = tf.slice(tf.sub(top, bottom), [0, 0, 0, 0], [-1, height - 1, -1, -1])

    num_pixels_in_vertical_diff = get_tensor_num_elements(vertical_diff)
    num_pixels_in_horizontal_diff = get_tensor_num_elements(horizontal_diff)

    # Why there's a 2 here? I added it according to https://github.com/antlerros/tensorflow-fast-neuralstyle and
    # https://github.com/anishathalye/neural-style
    total_var = 2 * (tf.nn.l2_loss(horizontal_diff) / num_pixels_in_horizontal_diff + tf.nn.l2_loss(
        vertical_diff) / num_pixels_in_vertical_diff)

    return total_var


def decode_image(contents, channels=None, name=None):
    """Convenience function for `decode_gif`, `decode_jpeg`, and `decode_png`.
    Detects whether an image is a GIF, JPEG, or PNG, and performs the appropriate
    operation to convert the input bytes `string` into a `Tensor` of type `uint8`.
    Note: `decode_gif` returns a 4-D array `[num_frames, height, width, 3]`, as
    opposed to `decode_jpeg` and `decode_png`, which return 3-D arrays
    `[height, width, num_channels]`. Make sure to take this into account when
    constructing your graph if you are intermixing GIF files with JPEG and/or PNG
    files.
    Args:
      contents: 0-D `string`. The encoded image bytes.
      channels: An optional `int`. Defaults to `0`. Number of color channels for
        the decoded image.
      name: A name for the operation (optional)
    Returns:
      `Tensor` with type `uint8` with shape `[height, width, num_channels]` for
        JPEG and PNG images and shape `[num_frames, height, width, 3]` for GIF
        images.
    """
    with ops.name_scope(name, 'decode_image') as scope:
        if channels not in (None, 0, 1, 3):
            raise ValueError('channels must be in (None, 0, 1, 3)')
        substr = string_ops.substr(contents, 0, 4)

        def _gif():
            # Create assert op to check that bytes are GIF decodable
            is_gif = math_ops.equal(substr, b'\x47\x49\x46\x38', name='is_gif')
            decode_msg = 'Unable to decode bytes as JPEG, PNG, or GIF'
            assert_decode = control_flow_ops.Assert(is_gif, [decode_msg])
            # Create assert to make sure that channels is not set to 1
            # Already checked above that channels is in (None, 0, 1, 3)
            gif_channels = 0 if channels is None else channels
            good_channels = math_ops.not_equal(gif_channels, 1, name='check_channels')
            channels_msg = 'Channels must be in (None, 0, 3) when decoding GIF images'
            assert_channels = control_flow_ops.Assert(good_channels, [channels_msg])
            with ops.control_dependencies([assert_decode, assert_channels]):
                return gen_image_ops.decode_gif(contents)

        def _png():
            return gen_image_ops.decode_png(contents, channels)

        def check_png():
            is_png = math_ops.equal(substr, b'\211PNG', name='is_png')
            return control_flow_ops.cond(is_png, _png, _gif, name='cond_png')

        def _jpeg():
            return gen_image_ops.decode_jpeg(contents, channels)

        is_jpeg = math_ops.equal(substr, b'\xff\xd8\xff\xe0', name='is_jpeg')
        return control_flow_ops.cond(is_jpeg, _jpeg, check_png, name='cond_jpeg')

def decode_image_with_file_name(contents, file_name, channels=None, name=None):
    # Cannot use the tensorflow decode_image function because sometimes the image file does not contain the proper
    # prefix for some reason, but as long as the extension is correct, my function will work.
    assert len(contents.get_shape().as_list()) == 0 and len(file_name.get_shape().as_list()) == 0
    ext = get_tf_string_extension(file_name)

    def _png():
        return gen_image_ops.decode_png(contents, channels)

    def _jpeg():
        is_jpeg = tf.logical_or(math_ops.equal(ext, 'jpg', name='is_jpg'),
                                math_ops.equal(ext, 'JPG', name='is_jpg_cap'))
        decode_msg = 'Unable to decode bytes as JPEG or PNG.'
        assert_decode = control_flow_ops.Assert(is_jpeg, [decode_msg])
        with ops.control_dependencies([assert_decode, ]):
            return gen_image_ops.decode_jpeg(contents, channels)

    is_png = math_ops.logical_or(math_ops.equal(ext, 'png', name='is_png'), math_ops.equal(ext, 'PNG', name='is_png_cap'))
    return control_flow_ops.cond(is_png, _png, _jpeg, name='cond_png')

def get_tf_string_extension(file_name):
    file_name = tf.identity(file_name)
    if len(file_name.get_shape().as_list()) == 0:
        split_result = tf.string_split([file_name], '.')
        return tf.sparse_tensor_to_dense(split_result, default_value="")[0, -1]
    elif len(file_name.get_shape().as_list()) == 1:
        split_result = tf.string_split(file_name, '.')
        return tf.sparse_tensor_to_dense(split_result, default_value="")[:, -1]
    else:
        raise AttributeError("Incorrect input dimension to get_tf_string_extension. "
                             "The input should be either a 0d or 1d string. The current input shape is %s."
                             %(str(file_name.get_shape().as_list())))
