import tensorflow as tf
import numpy as np
from .pointer_components import Encoder, Decoder, Attention, PointerSwitch
from .decode import PointerDecoder
from .eval import rouge_score

tf.enable_eager_execution()


class PointerNetwork(tf.keras.Model):
    def __init__(self, enc_units, dec_units, voc_size, att_units, switch_units,
                 max_len, start_token, end_token, padding_char):
        super().__init__()
        self.encoder = Encoder(enc_units)
        self.decoder = Decoder(dec_units, voc_size)
        self.attention = Attention(att_units)
        self.pointer_switch = PointerSwitch(switch_units)
        self.embeddings = False
        self.max_len = max_len
        self.start_token = start_token
        self.end_token = end_token
        self.voc_size = voc_size
        self.padding_char = padding_char

        self.optimizer = tf.train.AdamOptimizer()
        self._loss = 0  # used during training

    def set_embeddings_layer(self, embeddings_layer):
        self.embeddings = embeddings_layer
        self.pointer_decoder = PointerDecoder(self.embeddings, self.encoder,
                                              self.decoder, self.attention, self.pointer_switch,
                                              self.max_len, self.start_token, self.end_token)

    def predict_batch(self, X):
        assert self.embeddings, "Call self.set_embeddings_layer first"
        return self.pointer_decoder.predict_batch(X)

    def pointer_batch_loss(self, gen, y, d_prob, p_prob, s_prob):
        # Compute the mask to ignore the padding in the loss
        mask = 1-tf.cast(tf.equal(gen[:, None],
                                  tf.ones(gen[:, None].shape) *
                                  self.padding_char
                                  ), tf.float32)

        # Compute pointer loss across all values of y for the pointer and generated probs
        pointer_mat = (p_prob + (1 - s_prob)) * mask
        generator_mat = (d_prob + s_prob) * mask

        # Add the expected loss in terms of likelihood
        batch_loss = 0
        for i, g in enumerate(gen):
            # Add if the word was taken from the input
            if g == 0:
                batch_loss += pointer_mat[i, y[i]]
            # Add if the word was generated by the network
            else:
                batch_loss += generator_mat[i, y[i]]

        # Reduce to scalar, dont forget to include minus sign (its a loss not a likelihood)
        return -batch_loss

    def __train_batch(self, X, y, gen):
        assert self.embeddings, "Call self.load_embeddings first"

        X = tf.convert_to_tensor(X)
        y = tf.convert_to_tensor(y, dtype='int32')
        gen = tf.convert_to_tensor(gen, dtype='float32')

        enc_inp = self.embeddings(X)
        enc_states, h1, h2 = self.encoder(enc_inp)
        c_vec = h1
        input_tokens = y[:, 0]
        loss = 0
        for t in range(1, y.shape[1]):
            # Get embeddings
            dec_input = self.embeddings(input_tokens)

            # Get decoder output
            decoded_state, h1, h2, decoded_probs = self.decoder(
                dec_input, c_vec, [h1, h2])

            # Get context vector for the next step, and pointer probabilities
            c_vec, pointer_probs = self.attention(enc_states, decoded_state)

            # Get switch probability (BS*1)
            switch_probs = self.pointer_switch(h1, c_vec)

            # Is target generated or extracted from the input (BS * 1)
            batch_gen = tf.convert_to_tensor(gen[:, t])

            # Compute Pointer Network batch loss at timestep t
            loss += self.pointer_batch_loss(batch_gen, y[:, t], decoded_probs,
                                            pointer_probs, switch_probs)

            # Get next decoder input tokens
            input_tokens = y[:, t]

        # Dont forget to divide by summary lenght N, since we lose the /N component n by calling
        # N times softmax cross entropy
        loss = loss / int(y.shape[1]-1)
        self._loss = loss
        return loss

    def train_on_batch(self, X, y, gen):
        self.optimizer.minimize(lambda: self.__train_batch(X, y, gen))
        return [self._loss]

    def evaluate(self, X, y, verbose=0):
        y_ = self.predict_batch(X)
        return rouge_score(y, y_)
