# 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, 10, 34, 30), 'tensorflow.exp', 'tf.exp', ({(34, 17, 34, 29): '0.5 * logvar'}, {}), '(0.5 * logvar)', True, 'import tensorflow as tf\n'), ((18, 11, 18, 25), 'numpy.max', 'np.max', ({(18, 18, 18, 24): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((18, 27, 18, 42), 'numpy.mean', 'np.mean', ({(18, 35, 18, 41): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((33, 31, 33, 47), 'tensorflow.shape', 'tf.shape', ({(33, 40, 33, 46): 'logvar'}, {}), '(logvar)', True, 'import tensorflow as tf\n'), ((35, 12, 35, 37), 'tensorflow.multiply', 'tf.multiply', ({(35, 24, 35, 27): 'std', (35, 29, 35, 36): 'epsilon'}, {}), '(std, epsilon)', True, 'import tensorflow as tf\n'), ((57, 9, 57, 61), 'tensorflow.variable_scope', 'tf.variable_scope', (), '', True, 'import tensorflow as tf\n'), ((61, 27, 61, 108), 'tensorflow.nn.dynamic_rnn', 'tf.nn.dynamic_rnn', (), '', True, 'import tensorflow as tf\n'), ((71, 9, 71, 61), 'tensorflow.variable_scope', 'tf.variable_scope', (), '', True, 'import tensorflow as tf\n'), ((75, 27, 75, 132), 'tensorflow.nn.bidirectional_dynamic_rnn', 'tf.nn.bidirectional_dynamic_rnn', (), '', True, 'import tensorflow as tf\n'), ((76, 24, 76, 51), 'tensorflow.concat', 'tf.concat', ({(76, 34, 76, 47): 'encoded_input', (76, 49, 76, 50): '1'}, {}), '(encoded_input, 1)', True, 'import tensorflow as tf\n'), ((87, 20, 87, 56), 'tensorflow.reshape', 'tf.reshape', ({(87, 31, 87, 45): 'usr_input_sent', (87, 47, 87, 55): '[-1, 50]'}, {}), '(usr_input_sent, [-1, 50])', True, 'import tensorflow as tf\n'), ((87, 58, 87, 91), 'tensorflow.sequence_mask', 'tf.sequence_mask', ({(87, 75, 87, 86): 'length_mask', (87, 88, 87, 90): '(50)'}, {}), '(length_mask, 50)', True, 'import tensorflow as tf\n'), ((96, 9, 96, 36), 'tensorflow.variable_scope', 'tf.variable_scope', ({(96, 27, 96, 35): '"""repeat"""'}, {}), "('repeat')", True, 'import tensorflow as tf\n'), ((97, 26, 97, 52), 'tensorflow.expand_dims', 'tf.expand_dims', ({(97, 41, 97, 47): 'tensor', (97, 49, 97, 51): '-1'}, {}), '(tensor, -1)', True, 'import tensorflow as tf\n'), ((99, 23, 99, 68), 'tensorflow.tile', 'tf.tile', (), '', True, 'import tensorflow as tf\n'), ((46, 15, 46, 63), 'tensorflow.reduce_mean', 'tf.reduce_mean', (), '', True, 'import tensorflow as tf\n'), ((48, 15, 48, 62), 'tensorflow.reduce_sum', 'tf.reduce_sum', (), '', True, 'import tensorflow as tf\n'), ((60, 26, 60, 50), 'tensorflow.to_int32', 'tf.to_int32', ({(60, 38, 60, 49): 'length_mask'}, {}), '(length_mask)', True, 'import tensorflow as tf\n'), ((74, 26, 74, 50), 'tensorflow.to_int32', 'tf.to_int32', ({(74, 38, 74, 49): 'length_mask'}, {}), '(length_mask)', True, 'import tensorflow as tf\n'), ((100, 51, 100, 67), 'tensorflow.shape', 'tf.shape', ({(100, 60, 100, 66): 'tensor'}, {}), '(tensor)', True, 'import tensorflow as tf\n'), ((24, 40, 24, 60), 'tensorflow.exp', 'tf.exp', ({(24, 47, 24, 59): 'recog_logvar'}, {}), '(recog_logvar)', True, 'import tensorflow as tf\n'), ((24, 62, 24, 82), 'tensorflow.exp', 'tf.exp', ({(24, 69, 24, 81): 'prior_logvar'}, {}), '(prior_logvar)', True, 'import tensorflow as tf\n'), ((29, 30, 29, 45), 'tensorflow.log', 'tf.log', ({(29, 37, 29, 44): '(2 * np.pi)'}, {}), '(2 * np.pi)', True, 'import tensorflow as tf\n'), ((29, 64, 29, 81), 'tensorflow.pow', 'tf.pow', ({(29, 72, 29, 76): '(x - mu)', (29, 79, 29, 80): '(2)'}, {}), '(x - mu, 2)', True, 'import tensorflow as tf\n'), ((29, 83, 29, 97), 'tensorflow.exp', 'tf.exp', ({(29, 90, 29, 96): 'logvar'}, {}), '(logvar)', True, 'import tensorflow as tf\n'), ((23, 40, 23, 70), 'tensorflow.pow', 'tf.pow', ({(23, 47, 23, 66): '(prior_mu - recog_mu)', (23, 68, 23, 69): '(2)'}, {}), '(prior_mu - recog_mu, 2)', True, 'import tensorflow as tf\n'), ((23, 72, 23, 92), 'tensorflow.exp', 'tf.exp', ({(23, 79, 23, 91): 'prior_logvar'}, {}), '(prior_logvar)', True, 'import tensorflow as tf\n'), ((59, 62, 59, 79), 'tensorflow.abs', 'tf.abs', ({(59, 69, 59, 78): 'embedding'}, {}), '(embedding)', True, 'import tensorflow as tf\n'), ((73, 62, 73, 79), 'tensorflow.abs', 'tf.abs', ({(73, 69, 73, 78): 'embedding'}, {}), '(embedding)', True, 'import tensorflow as tf\n'), ((14, 71, 14, 90), 'nltk.translate.bleu_score.SmoothingFunction', 'SmoothingFunction', ({}, {}), '()', False, 'from nltk.translate.bleu_score import SmoothingFunction\n')]