"""
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)

[((575, 604), 'tensorflow.variable_scope', 'tf.variable_scope', (['scope_name'], {}), '(scope_name)\n', (592, 604), True, 'import tensorflow as tf\n'), ((611, 628), 'tensorflow.device', 'tf.device', (['device'], {}), '(device)\n', (620, 628), True, 'import tensorflow as tf\n'), ((1486, 1532), 'tensorflow.matmul', 'tf.matmul', (['context_encoded', 'self.label_weights'], {}), '(context_encoded, self.label_weights)\n', (1495, 1532), True, 'import tensorflow as tf\n'), ((1564, 1593), 'tensorflow.sigmoid', 'tf.sigmoid', (['self.label_scores'], {}), '(self.label_scores)\n', (1574, 1593), True, 'import tensorflow as tf\n'), ((1769, 1822), 'tensorflow.matmul', 'tf.matmul', (['true_entity_embeddings', 'self.label_weights'], {}), '(true_entity_embeddings, self.label_weights)\n', (1778, 1822), True, 'import tensorflow as tf\n'), ((1861, 1890), 'tensorflow.sigmoid', 'tf.sigmoid', (['self.label_scores'], {}), '(self.label_scores)\n', (1871, 1890), True, 'import tensorflow as tf\n'), ((1972, 2001), 'tensorflow.variable_scope', 'tf.variable_scope', (['scope_name'], {}), '(scope_name)\n', (1989, 2001), True, 'import tensorflow as tf\n'), ((2008, 2029), 'tensorflow.device', 'tf.device', (['device_gpu'], {}), '(device_gpu)\n', (2017, 2029), True, 'import tensorflow as tf\n'), ((2097, 2213), 'tensorflow.nn.sigmoid_cross_entropy_with_logits', 'tf.nn.sigmoid_cross_entropy_with_logits', ([], {'logits': 'self.label_scores', 'targets': 'true_label_ids', 'name': '"""labeling_loss"""'}), "(logits=self.label_scores, targets=\n    true_label_ids, name='labeling_loss')\n", (2136, 2213), True, 'import tensorflow as tf\n'), ((2423, 2552), 'tensorflow.nn.sigmoid_cross_entropy_with_logits', 'tf.nn.sigmoid_cross_entropy_with_logits', ([], {'logits': 'self.entity_label_scores', 'targets': 'true_label_ids', 'name': '"""entity_labeling_loss"""'}), "(logits=self.entity_label_scores,\n    targets=true_label_ids, name='entity_labeling_loss')\n", (2462, 2552), True, 'import tensorflow as tf\n'), ((1027, 1095), 'tensorflow.concat', 'tf.concat', (['(1)', '[context_encoded, mention_embed]'], {'name': '"""con_ment_repr"""'}), "(1, [context_encoded, mention_embed], name='con_ment_repr')\n", (1036, 1095), True, 'import tensorflow as tf\n'), ((2286, 2326), 'tensorflow.reduce_sum', 'tf.reduce_sum', (['self.cross_entropy_losses'], {}), '(self.cross_entropy_losses)\n', (2299, 2326), True, 'import tensorflow as tf\n'), ((2344, 2372), 'tensorflow.to_float', 'tf.to_float', (['self.batch_size'], {}), '(self.batch_size)\n', (2355, 2372), True, 'import tensorflow as tf\n'), ((2633, 2681), 'tensorflow.reduce_sum', 'tf.reduce_sum', (['self.enlabel_cross_entropy_losses'], {}), '(self.enlabel_cross_entropy_losses)\n', (2646, 2681), True, 'import tensorflow as tf\n'), ((2699, 2727), 'tensorflow.to_float', 'tf.to_float', (['self.batch_size'], {}), '(self.batch_size)\n', (2710, 2727), True, 'import tensorflow as tf\n'), ((856, 914), 'tensorflow.random_normal_initializer', 'tf.random_normal_initializer', ([], {'mean': '(0.0)', 'stddev': '(1.0 / 100.0)'}), '(mean=0.0, stddev=1.0 / 100.0)\n', (884, 914), True, 'import tensorflow as tf\n'), ((1313, 1371), 'tensorflow.random_normal_initializer', 'tf.random_normal_initializer', ([], {'mean': '(0.0)', 'stddev': '(1.0 / 100.0)'}), '(mean=0.0, stddev=1.0 / 100.0)\n', (1341, 1371), True, 'import tensorflow as tf\n')]