# Copyright 2019 Google LLC
#
# 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
#
#     https://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.

"""Utility functions for building models."""


import collections
import logging
import re
import time
import tensorflow.compat.v1 as tf
from rnn_decoder.multi_rnn import MultiRNNCell
from utils import dialogue_utils
from utils import iterator_utils
from utils import misc_utils as utils
from utils import vocab_utils


def get_initializer(init_op, seed=None, init_weight=None):
  """Create an initializer. init_weight is only for uniform."""
  if init_op == "uniform":
    assert init_weight
    return tf.random_uniform_initializer(-init_weight, init_weight, seed=seed)
  elif init_op == "glorot_normal":
    return tf.keras.initializers.glorot_normal(seed=seed)
  elif init_op == "glorot_uniform":
    return tf.keras.initializers.glorot_uniform(seed=seed)
  else:
    raise ValueError("Unknown init_op %s" % init_op)


def get_device_str(device_id, num_gpus):
  """Return a device string for multi-GPU setup."""
  if num_gpus == 0:
    return "/cpu:0"
  device_str_output = "/gpu:%d" % (device_id % num_gpus)
  return device_str_output


class ExtraArgs(
    collections.namedtuple(
        "ExtraArgs",
        ("single_cell_fn", "model_device_fn", "attention_mechanism_fn"))):
  pass


class TrainModel(
    collections.namedtuple(
        "TrainModel",
        ("graph", "model", "placeholder_iterator", "placeholder_handle",
         "train_iterator", "skip_count_placeholder"))):
  pass


class EvalModel(
    collections.namedtuple(
        "EvalModel",
        ("graph", "model", "placeholder_iterator", "placeholder_handle",
         "eval_iterator", "data_file_placeholder", "kb_file_placeholder"))):
  pass


class InferModel(
    collections.namedtuple(
        "InferModel",
        ("graph", "model", "placeholder_iterator", "placeholder_handle",
         "infer_iterator", "data_src_placeholder", "kb_placeholder",
         "batch_size_placeholder"))):
  pass


class SelfplayModel(
    collections.namedtuple(
        "SelfplayModel",
        ("graph", "model", "placeholder_iterator", "placeholder_handle",
         "train_iterator", "self_play_ft_iterator", "self_play_st_iterator",
         "data_placeholder", "kb_placeholder", "skip_count_placeholder",
         "batch_size_placeholder"))):
  pass


def create_train_model(model_creator,
                       hparams,
                       scope=None,
                       num_workers=1,
                       jobid=0,
                       extra_args=None):
  """Create graph, model and iterator for training."""
  graph = tf.Graph()

  with graph.as_default(), tf.container(scope or "train"):
    vocab_table, reverse_vocab_table = vocab_utils.create_vocab_tables(hparams.vocab_file)
    data_dataset = tf.data.TextLineDataset(hparams.train_data)
    kb_dataset = tf.data.TextLineDataset(hparams.train_kb)
    skip_count_placeholder = tf.placeholder(shape=(), dtype=tf.int64)
    # this is the actual train_iterator
    train_iterator = iterator_utils.get_iterator(
        data_dataset,
        kb_dataset,
        vocab_table,
        batch_size=hparams.batch_size,
        t1=hparams.t1.encode(),
        t2=hparams.t2.encode(),
        eod=hparams.eod,
        len_action=hparams.len_action,
        random_seed=hparams.random_seed,
        num_buckets=hparams.num_buckets,
        max_dialogue_len=hparams.max_dialogue_len,
        skip_count=skip_count_placeholder,
        num_shards=num_workers,
        shard_index=jobid)

    # this is the placeholder iterator. One can use this placeholder iterator
    # to switch between training and evauation.
    handle = tf.placeholder(tf.string, shape=[])
    iterator = tf.data.Iterator.from_string_handle(
        handle, train_iterator.output_types, train_iterator.output_shapes)
    batched_iterator = iterator_utils.get_batched_iterator(iterator)
    model_device_fn = None
    if extra_args:
      model_device_fn = extra_args.model_device_fn
    with tf.device(model_device_fn):
      model = model_creator(
          hparams,
          iterator=batched_iterator,
          handle=handle,
          mode=tf.estimator.ModeKeys.TRAIN,
          vocab_table=vocab_table,
          scope=scope,
          extra_args=extra_args,
          reverse_vocab_table=reverse_vocab_table)
  return TrainModel(
      graph=graph,
      model=model,
      placeholder_iterator=iterator,
      train_iterator=train_iterator,
      placeholder_handle=handle,
      skip_count_placeholder=skip_count_placeholder)


def create_eval_model(model_creator, hparams, scope=None, extra_args=None):
  """Create train graph, model, src/tgt file holders, and iterator."""
  vocab_file = hparams.vocab_file
  graph = tf.Graph()

  with graph.as_default(), tf.container(scope or "eval"):
    vocab_table = vocab_utils.create_vocab_tables(vocab_file)[0]
    data_file_placeholder = tf.placeholder(shape=(), dtype=tf.string)
    kb_file_placeholder = tf.placeholder(shape=(), dtype=tf.string)
    data_dataset = tf.data.TextLineDataset(data_file_placeholder)
    kb_dataset = tf.data.TextLineDataset(kb_file_placeholder)
    # this is the eval_actual iterator
    eval_iterator = iterator_utils.get_iterator(
        data_dataset,
        kb_dataset,
        vocab_table,
        batch_size=hparams.batch_size,
        t1=hparams.t1.encode(),
        t2=hparams.t2.encode(),
        eod=hparams.eod,