# Original work Copyright (C) 2017 Tiancheng Zhao, Carnegie Mellon University
# Modified work Copyright 2018 Weiyan Shi.

import tensorflow as tf
import numpy as np
from nltk.translate.bleu_score import sentence_bleu
from nltk.translate.bleu_score import SmoothingFunction


def get_bleu_stats(ref, hyps):
    scores = []
    for hyp in hyps:
        try:
            scores.append(sentence_bleu([ref], hyp, smoothing_function=SmoothingFunction().method7,
                                        weights=[1./3, 1./3,1./3]))
        except:
            scores.append(0.0)
    return np.max(scores), np.mean(scores)


def gaussian_kld(recog_mu, recog_logvar, prior_mu, prior_logvar):
    kld = -0.5 * tf.reduce_sum(1 + (recog_logvar - prior_logvar)
                               - tf.div(tf.pow(prior_mu - recog_mu, 2), tf.exp(prior_logvar))
                               - tf.div(tf.exp(recog_logvar), tf.exp(prior_logvar)), reduction_indices=1)
    return kld


def norm_log_liklihood(x, mu, logvar):
    return -0.5*tf.reduce_sum(tf.log(2*np.pi) + logvar + tf.div(tf.pow((x-mu), 2), tf.exp(logvar)), reduction_indices=1)


def sample_gaussian(mu, logvar):
    epsilon = tf.random_normal(tf.shape(logvar), name="epsilon")
    std = tf.exp(0.5 * logvar)
    z= mu + tf.multiply(std, epsilon)
    return z


def get_bow(embedding, avg=False):
    """
    Assumption, the last dimension is the embedding
    The second last dimension is the sentence length. The rank must be 3
    """
    embedding_size = embedding.get_shape()[2].value
    if avg:
        return tf.reduce_mean(embedding, reduction_indices=[1]), embedding_size
    else:
        return tf.reduce_sum(embedding, reduction_indices=[1]), embedding_size


def get_rnn_encode(embedding, cell, length_mask=None, scope=None, reuse=None):
    """
    Assumption, the last dimension is the embedding
    The second last dimension is the sentence length. The rank must be 3
    The padding should have zero
    """
    with tf.variable_scope(scope, 'RnnEncoding', reuse=reuse):
        if length_mask is None:
            length_mask = tf.reduce_sum(tf.sign(tf.reduce_max(tf.abs(embedding), reduction_indices=2)),reduction_indices=1)
            length_mask = tf.to_int32(length_mask)
        _, encoded_input = tf.nn.dynamic_rnn(cell, embedding, sequence_length=length_mask, dtype=tf.float32)
        return encoded_input, cell.state_size


def get_bi_rnn_encode(embedding, f_cell, b_cell, length_mask=None, scope=None, reuse=None):
    """
    Assumption, the last dimension is the embedding
    The second last dimension is the sentence length. The rank must be 3
    The padding should have zero
    """
    with tf.variable_scope(scope, 'RnnEncoding', reuse=reuse):
        if length_mask is None:
            length_mask = tf.reduce_sum(tf.sign(tf.reduce_max(tf.abs(embedding), reduction_indices=2)),reduction_indices=1)
            length_mask = tf.to_int32(length_mask)
        _, encoded_input = tf.nn.bidirectional_dynamic_rnn(f_cell, b_cell, embedding, sequence_length=length_mask, dtype=tf.float32)
        encoded_input = tf.concat(encoded_input, 1)
        return encoded_input, f_cell.state_size+b_cell.state_size

def get_prob_for_one_sent(vocab_prob, sent, length_mask=None):
    """

    :param vocab_prob:
    :param sent:
    :param length_mask:
    :return:
    """
    tf.boolean_mask(tf.reshape(usr_input_sent, [-1, 50]), tf.sequence_mask(length_mask, 50))

def tf_repeat(tensor, repeats):
    """

    :param tensor:
    :param repeats:
    :return:
    """
    with tf.variable_scope("repeat"):
        expanded_tensor = tf.expand_dims(tensor, -1)
        multiples = [1] + repeats
        tiled_tensor = tf.tile(expanded_tensor, multiples=multiples)
        repeated_tensor = tf.reshape(tiled_tensor, tf.shape(tensor) * repeats)
    return repeated_tensor
[(34, 'tensorflow.exp', 'tf.exp', True, 'import tensorflow as tf\n'), (18, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (18, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (33, 'tensorflow.shape', 'tf.shape', True, 'import tensorflow as tf\n'), (35, 'tensorflow.multiply', 'tf.multiply', True, 'import tensorflow as tf\n'), (57, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (61, 'tensorflow.nn.dynamic_rnn', 'tf.nn.dynamic_rnn', True, 'import tensorflow as tf\n'), (71, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (75, 'tensorflow.nn.bidirectional_dynamic_rnn', 'tf.nn.bidirectional_dynamic_rnn', True, 'import tensorflow as tf\n'), (76, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (87, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (87, 'tensorflow.sequence_mask', 'tf.sequence_mask', True, 'import tensorflow as tf\n'), (96, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (97, 'tensorflow.expand_dims', 'tf.expand_dims', True, 'import tensorflow as tf\n'), (99, 'tensorflow.tile', 'tf.tile', True, 'import tensorflow as tf\n'), (46, 'tensorflow.reduce_mean', 'tf.reduce_mean', True, 'import tensorflow as tf\n'), (48, 'tensorflow.reduce_sum', 'tf.reduce_sum', True, 'import tensorflow as tf\n'), (60, 'tensorflow.to_int32', 'tf.to_int32', True, 'import tensorflow as tf\n'), (74, 'tensorflow.to_int32', 'tf.to_int32', True, 'import tensorflow as tf\n'), (100, 'tensorflow.shape', 'tf.shape', True, 'import tensorflow as tf\n'), (24, 'tensorflow.exp', 'tf.exp', True, 'import tensorflow as tf\n'), (24, 'tensorflow.exp', 'tf.exp', True, 'import tensorflow as tf\n'), (29, 'tensorflow.log', 'tf.log', True, 'import tensorflow as tf\n'), (29, 'tensorflow.pow', 'tf.pow', True, 'import tensorflow as tf\n'), (29, 'tensorflow.exp', 'tf.exp', True, 'import tensorflow as tf\n'), (23, 'tensorflow.pow', 'tf.pow', True, 'import tensorflow as tf\n'), (23, 'tensorflow.exp', 'tf.exp', True, 'import tensorflow as tf\n'), (59, 'tensorflow.abs', 'tf.abs', True, 'import tensorflow as tf\n'), (73, 'tensorflow.abs', 'tf.abs', True, 'import tensorflow as tf\n'), (14, 'nltk.translate.bleu_score.SmoothingFunction', 'SmoothingFunction', False, 'from nltk.translate.bleu_score import SmoothingFunction\n')]