"""
Modifications copyright (C) 2020 Michael Strobl
"""

import time
import tensorflow as tf
import numpy as np

from entity_linker.models.base import Model

class LabelingModel(Model):
    """Unsupervised Clustering using Discrete-State VAE"""

    def __init__(self, batch_size, num_labels, context_encoded_dim,
                 true_entity_embeddings,
                 word_embed_dim, context_encoded, mention_embed, scope_name, device):

        self.batch_size = batch_size
        self.num_labels = num_labels
        self.word_embed_dim = word_embed_dim

        with tf.variable_scope(scope_name) as s, tf.device(device) as d:
            if mention_embed == None:
                self.label_weights = tf.get_variable(
                  name="label_weights",
                  shape=[context_encoded_dim, num_labels],
                  initializer=tf.random_normal_initializer(mean=0.0,
                                                           stddev=1.0/(100.0)))
            else:
                context_encoded = tf.concat(
                  1, [context_encoded, mention_embed], name='con_ment_repr')
                self.label_weights = tf.get_variable(
                  name="label_weights",
                  shape=[context_encoded_dim+word_embed_dim, num_labels],
                  initializer=tf.random_normal_initializer(mean=0.0,
                                                           stddev=1.0/(100.0)))

            # [B, L]
            self.label_scores = tf.matmul(context_encoded, self.label_weights)
            self.label_probs = tf.sigmoid(self.label_scores)

            ### PREDICT TYPES FROM ENTITIES
            #true_entity_embeddings = tf.nn.dropout(true_entity_embeddings, keep_prob=0.5)
            self.entity_label_scores = tf.matmul(true_entity_embeddings, self.label_weights)
            self.entity_label_probs = tf.sigmoid(self.label_scores)


    def loss_graph(self, true_label_ids, scope_name, device_gpu):
        with tf.variable_scope(scope_name) as s, tf.device(device_gpu) as d:
            # [B, L]
            self.cross_entropy_losses = tf.nn.sigmoid_cross_entropy_with_logits(
              logits=self.label_scores,
              targets=true_label_ids,
              name="labeling_loss")

            self.labeling_loss = tf.reduce_sum(
              self.cross_entropy_losses) / tf.to_float(self.batch_size)


            self.enlabel_cross_entropy_losses = tf.nn.sigmoid_cross_entropy_with_logits(
              logits=self.entity_label_scores,
              targets=true_label_ids,
              name="entity_labeling_loss")

            self.entity_labeling_loss = tf.reduce_sum(
              self.enlabel_cross_entropy_losses) / tf.to_float(self.batch_size)

[((22, 13, 22, 42), 'tensorflow.variable_scope', 'tf.variable_scope', ({(22, 31, 22, 41): 'scope_name'}, {}), '(scope_name)', True, 'import tensorflow as tf\n'), ((22, 49, 22, 66), 'tensorflow.device', 'tf.device', ({(22, 59, 22, 65): 'device'}, {}), '(device)', True, 'import tensorflow as tf\n'), ((39, 32, 39, 78), 'tensorflow.matmul', 'tf.matmul', ({(39, 42, 39, 57): 'context_encoded', (39, 59, 39, 77): 'self.label_weights'}, {}), '(context_encoded, self.label_weights)', True, 'import tensorflow as tf\n'), ((40, 31, 40, 60), 'tensorflow.sigmoid', 'tf.sigmoid', ({(40, 42, 40, 59): 'self.label_scores'}, {}), '(self.label_scores)', True, 'import tensorflow as tf\n'), ((44, 39, 44, 92), 'tensorflow.matmul', 'tf.matmul', ({(44, 49, 44, 71): 'true_entity_embeddings', (44, 73, 44, 91): 'self.label_weights'}, {}), '(true_entity_embeddings, self.label_weights)', True, 'import tensorflow as tf\n'), ((45, 38, 45, 67), 'tensorflow.sigmoid', 'tf.sigmoid', ({(45, 49, 45, 66): 'self.label_scores'}, {}), '(self.label_scores)', True, 'import tensorflow as tf\n'), ((49, 13, 49, 42), 'tensorflow.variable_scope', 'tf.variable_scope', ({(49, 31, 49, 41): 'scope_name'}, {}), '(scope_name)', True, 'import tensorflow as tf\n'), ((49, 49, 49, 70), 'tensorflow.device', 'tf.device', ({(49, 59, 49, 69): 'device_gpu'}, {}), '(device_gpu)', True, 'import tensorflow as tf\n'), ((51, 40, 54, 35), 'tensorflow.nn.sigmoid_cross_entropy_with_logits', 'tf.nn.sigmoid_cross_entropy_with_logits', (), '', True, 'import tensorflow as tf\n'), ((60, 48, 63, 42), 'tensorflow.nn.sigmoid_cross_entropy_with_logits', 'tf.nn.sigmoid_cross_entropy_with_logits', (), '', True, 'import tensorflow as tf\n'), ((30, 34, 31, 76), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((56, 33, 57, 40), 'tensorflow.reduce_sum', 'tf.reduce_sum', ({(57, 14, 57, 39): 'self.cross_entropy_losses'}, {}), '(self.cross_entropy_losses)', True, 'import tensorflow as tf\n'), ((57, 43, 57, 71), 'tensorflow.to_float', 'tf.to_float', ({(57, 55, 57, 70): 'self.batch_size'}, {}), '(self.batch_size)', True, 'import tensorflow as tf\n'), ((65, 40, 66, 48), 'tensorflow.reduce_sum', 'tf.reduce_sum', ({(66, 14, 66, 47): 'self.enlabel_cross_entropy_losses'}, {}), '(self.enlabel_cross_entropy_losses)', True, 'import tensorflow as tf\n'), ((66, 51, 66, 79), 'tensorflow.to_float', 'tf.to_float', ({(66, 63, 66, 78): 'self.batch_size'}, {}), '(self.batch_size)', True, 'import tensorflow as tf\n'), ((27, 30, 28, 78), 'tensorflow.random_normal_initializer', 'tf.random_normal_initializer', (), '', True, 'import tensorflow as tf\n'), ((35, 30, 36, 78), 'tensorflow.random_normal_initializer', 'tf.random_normal_initializer', (), '', True, 'import tensorflow as tf\n')]