# 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', (['(0.5 * logvar)'], {}), True, 'import tensorflow as tf\n'), (18, 'numpy.max', 'np.max', (['scores'], {}), True, 'import numpy as np\n'), (18, 'numpy.mean', 'np.mean', (['scores'], {}), True, 'import numpy as np\n'), (33, 'tensorflow.shape', 'tf.shape', (['logvar'], {}), True, 'import tensorflow as tf\n'), (35, 'tensorflow.multiply', 'tf.multiply', (['std', 'epsilon'], {}), True, 'import tensorflow as tf\n'), (57, 'tensorflow.variable_scope', 'tf.variable_scope', (['scope', '"""RnnEncoding"""'], {'reuse': 'reuse'}), True, 'import tensorflow as tf\n'), (61, 'tensorflow.nn.dynamic_rnn', 'tf.nn.dynamic_rnn', (['cell', 'embedding'], {'sequence_length': 'length_mask', 'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (71, 'tensorflow.variable_scope', 'tf.variable_scope', (['scope', '"""RnnEncoding"""'], {'reuse': 'reuse'}), True, 'import tensorflow as tf\n'), (75, 'tensorflow.nn.bidirectional_dynamic_rnn', 'tf.nn.bidirectional_dynamic_rnn', (['f_cell', 'b_cell', 'embedding'], {'sequence_length': 'length_mask', 'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (76, 'tensorflow.concat', 'tf.concat', (['encoded_input', '(1)'], {}), True, 'import tensorflow as tf\n'), (87, 'tensorflow.reshape', 'tf.reshape', (['usr_input_sent', '[-1, 50]'], {}), True, 'import tensorflow as tf\n'), (87, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['length_mask', '(50)'], {}), True, 'import tensorflow as tf\n'), (96, 'tensorflow.variable_scope', 'tf.variable_scope', (['"""repeat"""'], {}), True, 'import tensorflow as tf\n'), (97, 'tensorflow.expand_dims', 'tf.expand_dims', (['tensor', '(-1)'], {}), True, 'import tensorflow as tf\n'), (99, 'tensorflow.tile', 'tf.tile', (['expanded_tensor'], {'multiples': 'multiples'}), True, 'import tensorflow as tf\n'), (46, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['embedding'], {'reduction_indices': '[1]'}), True, 'import tensorflow as tf\n'), (48, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['embedding'], {'reduction_indices': '[1]'}), True, 'import tensorflow as tf\n'), (60, 'tensorflow.to_int32', 'tf.to_int32', (['length_mask'], {}), True, 'import tensorflow as tf\n'), (74, 'tensorflow.to_int32', 'tf.to_int32', (['length_mask'], {}), True, 'import tensorflow as tf\n'), (100, 'tensorflow.shape', 'tf.shape', (['tensor'], {}), True, 'import tensorflow as tf\n'), (24, 'tensorflow.exp', 'tf.exp', (['recog_logvar'], {}), True, 'import tensorflow as tf\n'), (24, 'tensorflow.exp', 'tf.exp', (['prior_logvar'], {}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.log', 'tf.log', (['(2 * np.pi)'], {}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.pow', 'tf.pow', (['(x - mu)', '(2)'], {}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.exp', 'tf.exp', (['logvar'], {}), True, 'import tensorflow as tf\n'), (23, 'tensorflow.pow', 'tf.pow', (['(prior_mu - recog_mu)', '(2)'], {}), True, 'import tensorflow as tf\n'), (23, 'tensorflow.exp', 'tf.exp', (['prior_logvar'], {}), True, 'import tensorflow as tf\n'), (59, 'tensorflow.abs', 'tf.abs', (['embedding'], {}), True, 'import tensorflow as tf\n'), (73, 'tensorflow.abs', 'tf.abs', (['embedding'], {}), True, 'import tensorflow as tf\n'), (14, 'nltk.translate.bleu_score.SmoothingFunction', 'SmoothingFunction', ([], {}), False, 'from nltk.translate.bleu_score import SmoothingFunction\n')]