import tensorflow as tf
import numpy as np
import dio
import seq2seq


OPTIMIZERS = {"rmsprop": tf.train.RMSPropOptimizer,
              "adadelta": tf.train.AdadeltaOptimizer,
              "adagrad": tf.train.AdagradOptimizer,
              "momentum": tf.train.MomentumOptimizer,
              "adam": tf.train.AdamOptimizer,
              "sgd": tf.train.GradientDescentOptimizer}

ACTIVATIONS = {"sigmoid": tf.sigmoid,
               "relu": tf.nn.relu,
               "tanh": tf.nn.tanh,
               "none": lambda x: x}


def bilstm(inputs, n_hidden, sequence_lengths, activation, dropout_rate=None):
    """
    Creates intermediate biLSTM layer with specified properties. In Tensorflow
    lingo, it creates the ops needed to run input through this layer (especially
    the `outputs` and `dropouts` ops
    :param inputs:
    :param n_hidden:
    :param sequence_lengths:
    :param activation:
    :param dropout_rate:
    :return:
    """
    if type(inputs) == tuple:  # output from lower biLSTM
        inputs = tf.concat(inputs, axis=2, name="concatenated_input")

    with tf.variable_scope('forward'):
        cell_fw = tf.contrib.rnn.LSTMCell(n_hidden, state_is_tuple=True,
                                          activation=activation)
    with tf.variable_scope('backward'):
        cell_bw = tf.contrib.rnn.LSTMCell(n_hidden, state_is_tuple=True,
                                          activation=activation)
    # get layer output op
    outputs, states = tf.nn.bidirectional_dynamic_rnn(
                                        cell_fw, cell_bw, inputs,
                                        dtype=tf.float32,
                                        sequence_length=sequence_lengths)
    dropouts = outputs
    # wrap dropout operation around
    if dropout_rate is not None:
        if type(outputs) == tuple:  # in case of bidirectional LSTM
            dropouts = tuple([tf.nn.dropout(o, keep_prob=1-dropout_rate,
                                            name="dropouts") for o in outputs])
        else:
            dropouts = tf.nn.dropout(outputs, keep_prob=1-dropout_rate,
                                     name="dropouts")
    return dropouts, states



def projection(inputs, y, sequence_lengths, class_weights):
    """
    Used for sequence labeling tasks, makes a dense projection of inputs to
    labels
    :param inputs:
    :param y:
    :param sequence_lengths:
    :param class_weights:
    :return:
    """
    if type(inputs) == tuple:  # concat fw and bw states
        inputs = tf.concat(inputs, axis=2, name="input_concat")
    batch_size, timesteps, n_hidden = [int(d) for d in inputs.get_shape()]
    n_classes = len(class_weights)
    # output projection, projects from hidden states to outputs
    w = tf.get_variable("weights", [n_hidden, n_classes],
                        initializer=xavier_init(n_hidden, n_classes))
    b = tf.get_variable("biases", [n_classes],
                        initializer=xavier_init(1, n_classes))

    # Generate predictions
    inputs_flat = tf.reshape(inputs, [-1, n_hidden])
    logits_flat = tf.nn.xw_plus_b(inputs_flat, w, b)
    preds_flat = tf.nn.softmax(logits_flat, name="preds_flat")
    # Return predictions as [batch, timesteps, classes]
    preds = tf.reshape(preds_flat, [batch_size, timesteps, n_classes],
                       name="preds")

    # Computing masked cross-entropy loss
    y_flat = tf.reshape(y, [-1], name="y_flat")
    y_onehot = tf.one_hot(y_flat, depth=n_classes, name="y_flat_onehot")
    losses = tf.nn.softmax_cross_entropy_with_logits(
        logits=logits_flat, labels=y_onehot)

    # preds, losses = tf.contrib.learn.ops.softmax_classifier(
    #     inputs, y_onehot, w, b, class_weight=class_weights)

    mask = tf.sign(tf.to_float(y_flat))  # 0 for PAD (PAD = 0), 1 for all others
    masked_losses = mask * losses  # make loss for all PAD zero
    # Bring back to [batch_size, timesteps] shape
    masked_losses = tf.reshape(masked_losses, [batch_size, timesteps],
                               name="masked_losses")
    # Calculate mean loss per example in batch
    mean_loss_by_example = tf.truediv(  # divide timestep-loss sum by seq_lens
        tf.reduce_sum(masked_losses, axis=1),
        tf.cast(sequence_lengths, tf.float32), name="mean_loss_by_ex"
    )
    # Finally, average loss over examples in batch
    loss = tf.reduce_mean(mean_loss_by_example, name="batch_mean_loss")
    return preds, loss



class SeqMtlModel:
    """
    Model for sequence multitask learning.
    """
    def __init__(self, tasks, input_size, num_layers, hid_dims, optimizer,
                 actfunc, timesteps=50, dropout_rate=0.1, learning_rate=0.1,
                 batch_size=64, embed_size=0, device='/cpu:0',
                 pretrained_embeddings=None, weighted_costs=False):
        """
        :param tasks:
        :param input_size: only relevant if not embed
        :param num_layers:
        :param hid_dims:
        :param optimizer:
        :param actfunc:
        :param timesteps:
        :param dropout_rate:
        :param learning_rate:
        :param batch_size:
        :param embed_size:
        """
        # The input placeholder, expects either IDs (if embed==True) or
        # vectors (pre-trained embeddings or one-hots)
        if embed_size > 0 or pretrained_embeddings is not None:
            self.X_in = tf.placeholder(tf.int32, (batch_size, timesteps),
                                       name="X_in")
        else:
            self.X_in = tf.placeholder(
                tf.float32, (batch_size, timesteps, input_size), name="X_in")
        # The following are registries for different operations per task
        # Each dict maps tasks to appropriate tf operations as registered below
        self.preds = {}
        self.losses = {}
        self.ys = {}

        # Some utilities
        self.optimizers = {}
        self.task2layer = {}
        self.task2tasktype = {}
        self.layer_states = []  # holds states for each intermediate layer
        self.layer_outputs = []  # holds output ops for each intermediate layer
        self.global_step = tf.Variable(0, trainable=False, name="global_step")
        self.dropout_rate_val = dropout_rate
        self.dropout_rate_op = tf.placeholder(tf.float32, name="dropout_rate")
        self.sequence_lengths = tf.placeholder(tf.int32, batch_size,
                                               name="sequence_lengths")

        print("=====================================")
        with tf.device(device):
            # # # INPUT LAYER # # #
            if pretrained_embeddings is not None:
                print("Creating trainable embedding layer initialized with "
                      "provided pretrained embeddings.")
                emb_matrix = tf.get_variable(
                    name="embed_matrix",
                    shape=pretrained_embeddings.shape,
                    initializer=tf.constant_initializer(pretrained_embeddings),
                    trainable=True)  # allow updating embeddings task-specifically
                output = tf.nn.embedding_lookup(emb_matrix, self.X_in,
                                                name="X_embedded")
            elif embed_size > 0:  # learn embed matrix from indices to embeddings
                print("Creating embedding layer mapping inputs to embeddings "
                      "of length {}.".format(embed_size))
                emb_matrix = tf.get_variable(
                    name="embed_matrix",
                    shape=[input_size, embed_size], dtype=tf.float32,
                    initializer=xavier_init(input_size, embed_size))
                output = tf.nn.embedding_lookup(emb_matrix, self.X_in,
                                                name="X_embedded")
            else:  # use provided vectors (pre-trained embeddings or one-hots)
                output = self.X_in
            self.layer_states.append(None)
            self.layer_outputs.append(output)

            # # # INTERMEDIATE LAYERS # # #
            print("Creating stacked biLSTM of {} layers with {} hidden units."
                  .format(num_layers, hid_dims))
            for i in range(num_layers):
                activation = ACTIVATIONS.get(actfunc)
                layer_input = self.layer_outputs[-1]
                with tf.variable_scope("layer_{}".format(i+1), reuse=None):
                    output, states = bilstm(layer_input, hid_dims,
                                            self.sequence_lengths, activation,
                                            self.dropout_rate_op)
                self.layer_states.append(states)
                self.layer_outputs.append(output)
                # dropouts from previous layer are input to next layer
                # layer_input = dropouts

            # # # TASK OUTPUTS # # #
            # Register tasks at layers
            for task_id, task in tasks.items():
                layer = task.layer
                if layer > num_layers:
                    print("Warning! Task was to be registered at layer {}, but "
                          "model only has {} layers. Using top layer instead.".
                          format(layer, num_layers))
                    layer = num_layers
                print("Registering task '{}' (type '{}', {} classes) at LSTM "
                      "layer {}.".format(task_id, task.task_type,
                                         task.get_num_labels(), layer))
                try:
                    layer_input = self.layer_outputs[layer]
                except IndexError:
                    import sys
                    sys.stderr.write("Error: Trying trying to register task, not "
                                     "enough hidden layers (expected at least {}, "
                                     "only {} present.)".format(layer, num_layers))
                layer_y, pred, loss, opt = None, None, None, None
                loss = None
                n_classes = task.get_num_labels()
                with tf.variable_scope("output_{}".format(task_id), reuse=None):
                    if task.task_type == "lbl":
                        class_weights = task.get_class_weights() if weighted_costs \
                            else np.ones(n_classes)
                        # print("  Class weights: ", class_weights)
                        layer_y = tf.placeholder(tf.int32,
                                                 (batch_size, timesteps),
                                                 name="labels")
                        pred, loss = projection(layer_input, layer_y,
                                                self.sequence_lengths,
                                                class_weights)
                    elif task.task_type == "seq":
                        init_state = self.layer_states[layer]
                        layer_y = tf.placeholder(tf.int32,
                                                 (batch_size, timesteps),
                                                 name="labels")
                        pred_tr, pred, loss = decoder_seq2seq(
                            layer_input, init_state, layer_y, hid_dims, n_classes,
                            self.sequence_lengths, batch_size, None)
                    elif task.task_type == "cls":
                        layer_y = tf.placeholder(tf.int32, (batch_size, 1),
                                                 name="labels")
                        pred, loss = classifier(layer_input, layer_y, n_classes)
                    else:
                        raise ValueError("Illegal task type '{}'. Only 'lbl' and "
                                         "'seq' supported.".format(task.task_type))
                optmzr = OPTIMIZERS.get(optimizer)(learning_rate=learning_rate)
                opt = optmzr.minimize(loss)
                # Register operations for tasks
                self.ys[task_id] = layer_y
                self.preds[task_id] = pred
                self.losses[task_id] = loss
                self.optimizers[task_id] = opt
        print("=====================================")
        self.saver = tf.train.Saver(tf.global_variables())

    def step(self, session, task, x, input_lengths, y=None, mode="train"):
        """
        Performs a model step. Performs different sequence of operations
        (computing predictions/loss, optimizing) depending on mode.
        :param session:
        :param task:
        :param x:
        :param input_lengths:
        :param y:
        :param mode: any of 'train', 'eval', 'decode'.
        :return:
        """
        pred = self.preds[task.task_id]
        optimizer = self.optimizers[task.task_id]
        loss = self.losses[task.task_id]
        gold = self.ys[task.task_id]

        l = None
        if mode == "decode":
            feed = {self.X_in: x, self.sequence_lengths: input_lengths,
                    self.dropout_rate_op: 0}
            p = session.run(pred, feed_dict=feed)
        elif mode == "eval":
            assert (y is not None), "Need to provide gold sequence y " \
                                    "when evaluating"
            feed = {self.X_in: x, gold: y, self.sequence_lengths: input_lengths,
                    self.dropout_rate_op: 0}
            p, l = session.run([pred, loss], feed_dict=feed)
        elif mode == "train":
            assert (y is not None), "Need to provide gold sequence y " \
                                    "when training"
            step_incr = tf.assign_add(self.global_step, 1)
            feed = {self.X_in: x, gold: y, self.sequence_lengths: input_lengths,
                    self.dropout_rate_op: self.dropout_rate_val}
            p, l, _, _ = session.run([pred, loss, optimizer, step_incr],
                                     feed_dict=feed)
        else:
            raise ValueError("Illegal step mode '{}'. Valid modes are 'train', "
                             "'decode' and 'eval'.")
        return p, l