# Copyright 2017 The TensorFlow Authors All Rights Reserved.
#
# 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.
# ==============================================================================

import numpy as np
import tensorflow as tf
import tensorflow_fold as td
from model_n2nmn import netgen_att
from model_n2nmn import assembler
from model_n2nmn.modules import Modules


class Model:

  def __init__(self,
               config,
               kb,
               text_seq_batch,
               seq_length_batch,
               num_vocab_txt,
               num_vocab_nmn,
               EOS_idx,
               num_choices,
               decoder_sampling,
               use_gt_layout=None,
               gt_layout_batch=None,
               scope='neural_module_network',
               reuse=None):
    with tf.variable_scope(scope, reuse=reuse):
      # Part 1: Seq2seq RNN to generate module layout tokens

      embedding_mat = tf.get_variable(
        'embedding_mat', [num_vocab_txt, config.embed_dim_txt],
        initializer=tf.contrib.layers.xavier_initializer())

      with tf.variable_scope('layout_generation'):
        att_seq2seq = netgen_att.AttentionSeq2Seq(
            config, text_seq_batch, seq_length_batch, num_vocab_txt,
            num_vocab_nmn, EOS_idx, decoder_sampling, embedding_mat,
            use_gt_layout, gt_layout_batch)
        self.att_seq2seq = att_seq2seq
        predicted_tokens = att_seq2seq.predicted_tokens
        token_probs = att_seq2seq.token_probs
        word_vecs = att_seq2seq.word_vecs
        neg_entropy = att_seq2seq.neg_entropy
        self.atts = att_seq2seq.atts

        self.predicted_tokens = predicted_tokens
        self.token_probs = token_probs
        self.word_vecs = word_vecs
        self.neg_entropy = neg_entropy

        # log probability of each generated sequence
        self.log_seq_prob = tf.reduce_sum(tf.log(token_probs), axis=0)

      # Part 2: Neural Module Network
      with tf.variable_scope('layout_execution'):
        modules = Modules(config, kb, word_vecs, num_choices, embedding_mat)
        self.modules = modules
        # Recursion of modules
        att_shape = [len(kb)]
        # Forward declaration of module recursion
        att_expr_decl = td.ForwardDeclaration(td.PyObjectType(),
                                              td.TensorType(att_shape))
        # _key_find
        case_key_find = td.Record([('time_idx', td.Scalar(dtype='int32')),
                                   ('batch_idx', td.Scalar(dtype='int32'))])
        case_key_find = case_key_find >> td.ScopedLayer(
            modules.KeyFindModule, name_or_scope='KeyFindModule')
        # _key_filter
        case_key_filter = td.Record([('input_0', att_expr_decl()),
                                     ('time_idx', td.Scalar('int32')),
                                     ('batch_idx', td.Scalar('int32'))])
        case_key_filter = case_key_filter >> td.ScopedLayer(
            modules.KeyFilterModule, name_or_scope='KeyFilterModule')
        recursion_cases = td.OneOf(
            td.GetItem('module'),
            {'_key_find': case_key_find,
             '_key_filter': case_key_filter})
        att_expr_decl.resolve_to(recursion_cases)
        # _val_desc: output scores for choice (for valid expressions)
        predicted_scores = td.Record([('input_0', recursion_cases),
                                      ('time_idx', td.Scalar('int32')),
                                      ('batch_idx', td.Scalar('int32'))])
        predicted_scores = predicted_scores >> td.ScopedLayer(
            modules.ValDescribeModule, name_or_scope='ValDescribeModule')

        # For invalid expressions, define a dummy answer
        # so that all answers have the same form
        INVALID = assembler.INVALID_EXPR
        dummy_scores = td.Void() >> td.FromTensor(
            np.zeros(num_choices, np.float32))
        output_scores = td.OneOf(
            td.GetItem('module'),
            {'_val_desc': predicted_scores,
             INVALID: dummy_scores})

        # compile and get the output scores
        self.compiler = td.Compiler.create(output_scores)
        self.scores = self.compiler.output_tensors[0]

      # Regularization: Entropy + L2
      self.entropy_reg = tf.reduce_mean(neg_entropy)
      module_weights = [
          v for v in tf.trainable_variables()
          if (scope in v.op.name and v.op.name.endswith('weights'))
      ]
      self.l2_reg = tf.add_n([tf.nn.l2_loss(v) for v in module_weights])