# coding=utf-8
# Copyright 2019 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Utilities for building a LaserTagger TF model."""

from __future__ import absolute_import
from __future__ import division

from __future__ import print_function
from bert import modeling
from bert import optimization
from src import transformer_decoder
import tensorflow as tf
from src.official_transformer import model_params


class LaserTaggerConfig(modeling.BertConfig):
    """Model configuration for LaserTagger."""

    def __init__(self,
                 use_t2t_decoder=True,
                 decoder_num_hidden_layers=1,
                 decoder_hidden_size=768,
                 decoder_num_attention_heads=4,
                 decoder_filter_size=3072,
                 use_full_attention=False,
                 **kwargs):
        """Initializes an instance of LaserTagger configuration.

        This initializer expects both the BERT specific arguments and the
        Transformer decoder arguments listed below.

        Args:
          use_t2t_decoder: Whether to use the Transformer decoder (i.e.
            LaserTagger_AR). If False, the remaining args do not affect anything and
            can be set to default values.
          decoder_num_hidden_layers: Number of hidden decoder layers.
          decoder_hidden_size: Decoder hidden size.
          decoder_num_attention_heads: Number of decoder attention heads.
          decoder_filter_size: Decoder filter size.
          use_full_attention: Whether to use full encoder-decoder attention.
          **kwargs: The arguments that the modeling.BertConfig initializer expects.
        """
        super(LaserTaggerConfig, self).__init__(**kwargs)
        self.use_t2t_decoder = use_t2t_decoder
        self.decoder_num_hidden_layers = decoder_num_hidden_layers
        self.decoder_hidden_size = decoder_hidden_size
        self.decoder_num_attention_heads = decoder_num_attention_heads
        self.decoder_filter_size = decoder_filter_size
        self.use_full_attention = use_full_attention


class ModelFnBuilder(object):
    """Class for building `model_fn` closure for TPUEstimator."""

    def __init__(self, config, num_tags,
                 init_checkpoint,
                 learning_rate, num_train_steps,
                 num_warmup_steps, use_tpu,
                 use_one_hot_embeddings, max_seq_length):
        """Initializes an instance of a LaserTagger model.

        Args:
          config: LaserTagger model configuration.
          num_tags: Number of different tags to be predicted.
          init_checkpoint: Path to a pretrained BERT checkpoint (optional).
          learning_rate: Learning rate.
          num_train_steps: Number of training steps.
          num_warmup_steps: Number of warmup steps.
          use_tpu: Whether to use TPU.
          use_one_hot_embeddings: Whether to use one-hot embeddings for word
            embeddings.
          max_seq_length: Maximum sequence length.
        """
        self._config = config
        self._num_tags = num_tags
        self._init_checkpoint = init_checkpoint
        self._learning_rate = learning_rate
        self._num_train_steps = num_train_steps
        self._num_warmup_steps = num_warmup_steps
        self._use_tpu = use_tpu
        self._use_one_hot_embeddings = use_one_hot_embeddings
        self._max_seq_length = max_seq_length

    def _create_model(self, mode, input_ids, input_mask, segment_ids, labels,
                      labels_mask):
        """Creates a LaserTagger model."""
        is_training = (mode == tf.estimator.ModeKeys.TRAIN)
        model = modeling.BertModel(
            config=self._config,
            is_training=is_training,
            input_ids=input_ids,
            input_mask=input_mask,
            token_type_ids=segment_ids,
            use_one_hot_embeddings=self._use_one_hot_embeddings)

        final_hidden = model.get_sequence_output()

        # 使用更好的transformer_decoder解码层
        if self._config.use_t2t_decoder:
            # Size of the output vocabulary which contains the tags + begin and end
            # tokens used by the Transformer decoder.
            output_vocab_size = self._num_tags + 2  # TODO  为什么加２，begin and end
            params = _get_decoder_params(self._config, self._use_tpu,
                                         self._max_seq_length, output_vocab_size)
            decoder = transformer_decoder.TransformerDecoder(params, is_training)
            logits = decoder(input_mask, final_hidden, labels)  # labels is the id of operation
        else:
            # 使用普通的dense出概率
            if is_training:
                # I.e., 0.1 dropout
                final_hidden = tf.nn.dropout(final_hidden, keep_prob=0.9)

            logits = tf.layers.dense(
                final_hidden,
                self._num_tags,
                kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),
                name="output_projection")

        with tf.variable_scope("loss"):
            loss = None
            per_example_loss = None
            # 训练阶段
            if mode != tf.estimator.ModeKeys.PREDICT:
                # softmax_cross_entropy loss
                loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
                    labels=labels, logits=logits)
                per_example_loss = tf.truediv(
                    tf.reduce_sum(loss, axis=1),
                    tf.dtypes.cast(tf.reduce_sum(labels_mask, axis=1), tf.float32))
                loss = tf.reduce_mean(per_example_loss)
                pred = tf.cast(tf.argmax(logits, axis=-1), tf.int32)
            else:
                #   预测阶段
                if self._config.use_t2t_decoder:
                    pred = logits["outputs"]
                    # Transformer decoder reserves the first two IDs to the begin and the
                    # end token so we shift the IDs back.
                    pred -= 2
                else:
                    pred = tf.cast(tf.argmax(logits, axis=-1), tf.int32)

            return (loss, per_example_loss, pred)

    def build(self):
        """Returns `model_fn` closure for TPUEstimator."""

        def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
            """The `model_fn` for TPUEstimator."""

            tf.logging.info("*** Features ***")
            for name in sorted(features.keys()):
                tf.logging.info("  name = %s, shape = %s", name, features[name].shape)

            input_ids = features["input_ids"]
            input_mask = features["input_mask"]
            segment_ids = features["segment_ids"]

            labels = None
            labels_mask = None
            if mode != tf.estimator.ModeKeys.PREDICT:
                if self._config.use_t2t_decoder:
                    # Account for the begin and end tokens used by Transformer.
                    labels = features["labels"] + 2
                else:
                    labels = features["labels"]
                labels_mask = tf.cast(features["labels_mask"], tf.float32)

            (total_loss, per_example_loss, predictions) = self._create_model(
                mode, input_ids, input_mask, segment_ids, labels, labels_mask)

            tvars = tf.trainable_variables()
            initialized_variable_names = {}
            scaffold_fn = None
            if self._init_checkpoint:
                (assignment_map, initialized_variable_names
                 ) = modeling.get_assignment_map_from_checkpoint(tvars,
                                                                 self._init_checkpoint)
                if self._use_tpu:
                    def tpu_scaffold():
                        tf.train.init_from_checkpoint(self._init_checkpoint, assignment_map)
                        return tf.train.Scaffold()

                    scaffold_fn = tpu_scaffold
                else:
                    tf.train.init_from_checkpoint(self._init_checkpoint, assignment_map)

            tf.logging.info("**** Trainable Variables ****")
            # for var in tvars:
            #   tf.logging.info("Initializing the model from: %s",
            #                   self._init_checkpoint)
            #   init_string = ""
            #   if var.name in initialized_variable_names:
            #     init_string = ", *INIT_FROM_CKPT*"
            #   tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
            #                   init_string)

            output_spec = None
            if mode == tf.estimator.ModeKeys.TRAIN:
                train_op = optimization.create_optimizer(
                    total_loss, self._learning_rate, self._num_train_steps,
                    self._num_warmup_steps, self._use_tpu)

                output_spec = tf.contrib.tpu.TPUEstimatorSpec(
                    mode=mode,
                    loss=total_loss,
                    train_op=train_op,
                    scaffold_fn=scaffold_fn)

            elif mode == tf.estimator.ModeKeys.EVAL:
                def metric_fn(per_example_loss, labels, labels_mask, predictions):
                    """Compute eval metrics."""
                    accuracy = tf.cast(
                        tf.reduce_all(  # tf.reduce_all  相当于＂逻辑ＡＮＤ＂操作，找到输出完全正确的才算正确
                            tf.logical_or(
                                tf.equal(labels, predictions),
                                ~tf.cast(labels_mask, tf.bool)),
                            axis=1), tf.float32)
                    return {
                        # This is equal to the Exact score if the final realization step
                        # doesn't introduce errors.
                        "sentence_level_acc": tf.metrics.mean(accuracy),
                        "eval_loss": tf.metrics.mean(per_example_loss),
                    }

                eval_metrics = (metric_fn,
                                [per_example_loss, labels, labels_mask, predictions])
                output_spec = tf.contrib.tpu.TPUEstimatorSpec(
                    mode=mode,
                    loss=total_loss,
                    eval_metrics=eval_metrics,
                    scaffold_fn=scaffold_fn)
            else:
                output_spec = tf.contrib.tpu.TPUEstimatorSpec(
                    mode=mode, predictions={"pred": predictions},
                    scaffold_fn=scaffold_fn)
            return output_spec

        return model_fn


def _get_decoder_params(config, use_tpu,
                        max_seq_length,
                        output_vocab_size):
    """Returns hyperparameters for TransformerDecoder.

    Args:
      config: LaserTagger model configuration.
      use_tpu: Whether to train on TPUs.
      max_seq_length: Maximum sequence length.
      output_vocab_size: Size of the output vocabulary.

    Returns:
      Hyperparameter dictionary.
    """
    params = model_params.BASE_PARAMS
    params.update(
        num_hidden_layers=config.decoder_num_hidden_layers,
        hidden_size=config.decoder_hidden_size,
        num_heads=config.decoder_num_attention_heads,
        filter_size=config.decoder_filter_size,
        vocab_size=output_vocab_size,
        use_tpu=use_tpu,
        max_length=max_seq_length,
        # This parameter should not be changed since we want the number of decoded
        # tags to equal the number of source tokens.
        extra_decode_length=0)
    return params
