# Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You may not
# use this file except in compliance with the License. A copy of the License
# is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.

"""
Defines the dynamic source encodings ('coverage' mechanisms) for encoder/decoder networks as used in Tu et al. (2016).
"""
import logging
from typing import Callable

import mxnet as mx

from . import config
from . import constants as C
from . import layers
from . import rnn
from . import utils

logger = logging.getLogger(__name__)


class CoverageConfig(config.Config):
    """
    Coverage configuration.

    :param type: Coverage name.
    :param num_hidden: Number of hidden units for coverage networks.
    :param layer_normalization: Apply layer normalization to coverage networks.
    """
    def __init__(self,
                 type: str,
                 num_hidden: int,
                 layer_normalization: bool) -> None:
        super().__init__()
        self.type = type
        self.num_hidden = num_hidden
        self.layer_normalization = layer_normalization


def get_coverage(config: CoverageConfig) -> 'Coverage':
    """
    Returns a Coverage instance.

    :param config: Coverage configuration.
    :return: Instance of Coverage.
    """
    if config.type == 'count':
        utils.check_condition(config.num_hidden == 1, "Count coverage requires coverage_num_hidden==1")
    if config.type == "gru":
        return GRUCoverage(config.num_hidden, config.layer_normalization)
    elif config.type in {"tanh", "sigmoid", "relu", "softrelu"}:
        return ActivationCoverage(config.num_hidden, config.type, config.layer_normalization)
    elif config.type == "count":
        return CountCoverage()
    else:
        raise ValueError("Unknown coverage type %s" % config.type)


class Coverage:
    """
    Generic coverage class. Similar to Attention classes, a coverage instance returns a callable, update_coverage(),
    function when self.on() is called.
    """
    def __init__(self, prefix=C.COVERAGE_PREFIX):
        self.prefix = prefix

    def on(self, source: mx.sym.Symbol, source_length: mx.sym.Symbol, source_seq_len: int) -> Callable:
        """
        Returns callable to be used for updating coverage vectors in a sequence decoder.

        :param source: Shape: (batch_size, seq_len, encoder_num_hidden).
        :param source_length: Shape: (batch_size,).
        :param source_seq_len: Maximum length of source sequences.
        :return: Coverage callable.
        """

        def update_coverage(prev_hidden: mx.sym.Symbol,
                            attention_prob_scores: mx.sym.Symbol,
                            prev_coverage: mx.sym.Symbol):
            """
            :param prev_hidden: Previous hidden decoder state. Shape: (batch_size, decoder_num_hidden).
            :param attention_prob_scores: Current attention scores. Shape: (batch_size, source_seq_len).
            :param prev_coverage: Shape: (batch_size, source_seq_len, coverage_num_hidden).
            :return: Updated coverage matrix . Shape: (batch_size, source_seq_len, coverage_num_hidden).
            """
            raise NotImplementedError()

        return update_coverage


class CountCoverage(Coverage):
    """
    Coverage class that accumulates the attention weights for each source word.
    """

    def __init__(self) -> None:
        super().__init__()

    def on(self, source: mx.sym.Symbol, source_length: mx.sym.Symbol, source_seq_len: int) -> Callable:
        """
        Returns callable to be used for updating coverage vectors in a sequence decoder.

        :param source: Shape: (batch_size, seq_len, encoder_num_hidden).
        :param source_length: Shape: (batch_size,).
        :param source_seq_len: Maximum length of source sequences.
        :return: Coverage callable.
        """

        def update_coverage(prev_hidden: mx.sym.Symbol,
                            attention_prob_scores: mx.sym.Symbol,
                            prev_coverage: mx.sym.Symbol):
            """
            :param prev_hidden: Previous hidden decoder state. Shape: (batch_size, decoder_num_hidden).
            :param attention_prob_scores: Current attention scores. Shape: (batch_size, source_seq_len).
            :param prev_coverage: Shape: (batch_size, source_seq_len, coverage_num_hidden).
            :return: Updated coverage matrix . Shape: (batch_size, source_seq_len, coverage_num_hidden).
            """
            return prev_coverage + mx.sym.expand_dims(attention_prob_scores, axis=2)

        return update_coverage


class GRUCoverage(Coverage):
    """
    Implements a GRU whose state is the coverage vector.

    TODO: This implementation is slightly inefficient since the source is fed in at every step.
    It would be better to pre-compute the mapping of the source but this will likely mean opening up the GRU.

    :param coverage_num_hidden: Number of hidden units for coverage vectors.
    :param layer_normalization: If true, applies layer normalization for each gate in the GRU cell.
    """

    def __init__(self, coverage_num_hidden: int, layer_normalization: bool) -> None:
        super().__init__()
        self.num_hidden = coverage_num_hidden
        gru_prefix= "%sgru" % self.prefix
        if layer_normalization:
            self.gru = rnn.LayerNormPerGateGRUCell(self.num_hidden, prefix=gru_prefix)
        else:
            self.gru = mx.rnn.GRUCell(self.num_hidden, prefix=gru_prefix)

    def on(self, source: mx.sym.Symbol, source_length: mx.sym.Symbol, source_seq_len: int) -> Callable:
        """
        Returns callable to be used for updating coverage vectors in a sequence decoder.

        :param source: Shape: (batch_size, seq_len, encoder_num_hidden).
        :param source_length: Shape: (batch_size,).
        :param source_seq_len: Maximum length of source sequences.
        :return: Coverage callable.
        """

        def update_coverage(prev_hidden: mx.sym.Symbol,
                            attention_prob_scores: mx.sym.Symbol,
                            prev_coverage: mx.sym.Symbol):
            """
            :param prev_hidden: Previous hidden decoder state. Shape: (batch_size, decoder_num_hidden).
            :param attention_prob_scores: Current attention scores. Shape: (batch_size, source_seq_len).
            :param prev_coverage: Shape: (batch_size, source_seq_len, coverage_num_hidden).
            :return: Updated coverage matrix . Shape: (batch_size, source_seq_len, coverage_num_hidden).
            """

            # (batch_size, source_seq_len, decoder_num_hidden)
            expanded_decoder = mx.sym.broadcast_axis(
                data=mx.sym.expand_dims(data=prev_hidden, axis=1, name="%sexpand_decoder" % self.prefix),
                axis=1, size=source_seq_len, name="%sbroadcast_decoder" % self.prefix)

            expanded_att_scores = mx.sym.expand_dims(data=attention_prob_scores,
                                                     axis=2,
                                                     name="%sexpand_attention_scores" % self.prefix)

            # (batch_size, source_seq_len, encoder_num_hidden + decoder_num_hidden + 1)
            # +1 for the attention_prob_score for the source word
            concat_input = mx.sym.concat(source, expanded_decoder, expanded_att_scores, dim=2,
                                         name="%sconcat_inputs" % self.prefix)

            # (batch_size * source_seq_len, encoder_num_hidden + decoder_num_hidden + 1)
            flat_input = mx.sym.reshape(concat_input, shape=(-3, -1), name="%sflatten_inputs")

            # coverage: (batch_size * seq_len, coverage_num_hidden)
            coverage = mx.sym.reshape(data=prev_coverage, shape=(-3, -1))
            updated_coverage, _ = self.gru(flat_input, states=[coverage])

            # coverage: (batch_size, seq_len, coverage_num_hidden)
            coverage = mx.sym.reshape(updated_coverage, shape=(-1, source_seq_len, self.num_hidden))

            return mask_coverage(coverage, source_length)

        return update_coverage


class ActivationCoverage(Coverage):
    """
    Implements a coverage mechanism whose updates are performed by a Perceptron with
    configurable activation function.

    :param coverage_num_hidden: Number of hidden units for coverage vectors.
    :param activation: Type of activation for Perceptron.
    :param layer_normalization: If true, applies layer normalization before non-linear activation.
    """

    def __init__(self,
                 coverage_num_hidden: int,
                 activation: str,
                 layer_normalization: bool) -> None:
        super().__init__()
        self.activation = activation
        self.num_hidden = coverage_num_hidden
        # input (encoder) to hidden
        self.cov_e2h_weight = mx.sym.Variable("%se2h_weight" % self.prefix)
        # decoder to hidden
        self.cov_dec2h_weight = mx.sym.Variable("%si2h_weight" % self.prefix)
        # previous coverage to hidden
        self.cov_prev2h_weight = mx.sym.Variable("%sprev2h_weight" % self.prefix)
        # attention scores to hidden
        self.cov_a2h_weight = mx.sym.Variable("%sa2h_weight" % self.prefix)
        # optional layer normalization
        self.layer_norm = None
        if layer_normalization and not self.num_hidden != 1:
            self.layer_norm = layers.LayerNormalization(self.num_hidden,
                                                        prefix="%snorm" % self.prefix) if layer_normalization else None

    def on(self, source: mx.sym.Symbol, source_length: mx.sym.Symbol, source_seq_len: int) -> Callable:
        """
        Returns callable to be used for updating coverage vectors in a sequence decoder.

        :param source: Shape: (batch_size, seq_len, encoder_num_hidden).
        :param source_length: Shape: (batch_size,).
        :param source_seq_len: Maximum length of source sequences.
        :return: Coverage callable.
        """

        # (batch_size, seq_len, coverage_hidden_num)
        source_hidden = mx.sym.FullyConnected(data=source,
                                              weight=self.cov_e2h_weight,
                                              no_bias=True,
                                              num_hidden=self.num_hidden,
                                              flatten=False,
                                              name="%ssource_hidden_fc" % self.prefix)

        def update_coverage(prev_hidden: mx.sym.Symbol,
                            attention_prob_scores: mx.sym.Symbol,
                            prev_coverage: mx.sym.Symbol):
            """
            :param prev_hidden: Previous hidden decoder state. Shape: (batch_size, decoder_num_hidden).
            :param attention_prob_scores: Current attention scores. Shape: (batch_size, source_seq_len).
            :param prev_coverage: Shape: (batch_size, source_seq_len, coverage_num_hidden).
            :return: Updated coverage matrix . Shape: (batch_size, source_seq_len, coverage_num_hidden).
            """

            # (batch_size, seq_len, coverage_hidden_num)
            coverage_hidden = mx.sym.FullyConnected(data=prev_coverage,
                                                    weight=self.cov_prev2h_weight,
                                                    no_bias=True,
                                                    num_hidden=self.num_hidden,
                                                    flatten=False,
                                                    name="%sprevious_hidden_fc" % self.prefix)

            # (batch_size, source_seq_len, 1)
            attention_prob_scores = mx.sym.expand_dims(attention_prob_scores, axis=2)

            # (batch_size, source_seq_len, coverage_num_hidden)
            attention_hidden = mx.sym.FullyConnected(data=attention_prob_scores,
                                                     weight=self.cov_a2h_weight,
                                                     no_bias=True,
                                                     num_hidden=self.num_hidden,
                                                     flatten=False,
                                                     name="%sattention_fc" % self.prefix)

            # (batch_size, coverage_num_hidden)
            prev_hidden = mx.sym.FullyConnected(data=prev_hidden, weight=self.cov_dec2h_weight, no_bias=True,
                                                num_hidden=self.num_hidden, name="%sdecoder_hidden")

            # (batch_size, 1, coverage_num_hidden)
            prev_hidden = mx.sym.expand_dims(data=prev_hidden, axis=1,
                                             name="%sinput_decoder_hidden_expanded" % self.prefix)

            # (batch_size, source_seq_len, coverage_num_hidden)
            intermediate = mx.sym.broadcast_add(lhs=source_hidden, rhs=prev_hidden,
                                                name="%ssource_plus_hidden" % self.prefix)

            # (batch_size, source_seq_len, coverage_num_hidden)
            updated_coverage = intermediate + attention_hidden + coverage_hidden

            if self.layer_norm is not None:
                updated_coverage = self.layer_norm.normalize(updated_coverage)

            # (batch_size, seq_len, coverage_num_hidden)
            coverage = mx.sym.Activation(data=updated_coverage,
                                         act_type=self.activation,
                                         name="%sactivation" % self.prefix)

            return mask_coverage(coverage, source_length)

        return update_coverage


def mask_coverage(coverage: mx.sym.Symbol, source_length: mx.sym.Symbol) -> mx.sym.Symbol:
    """
    Masks all coverage scores that are outside the actual sequence.

    :param coverage: Input coverage vector. Shape: (batch_size, seq_len, coverage_num_hidden).
    :param source_length: Source length. Shape: (batch_size,).
    :return: Masked coverage vector. Shape: (batch_size, seq_len, coverage_num_hidden).
    """
    coverage = mx.sym.SwapAxis(data=coverage, dim1=0, dim2=1)
    coverage = mx.sym.SequenceMask(data=coverage, use_sequence_length=True, sequence_length=source_length)
    coverage = mx.sym.SwapAxis(data=coverage, dim1=0, dim2=1)
    return coverage
