# coding=utf-8
# Copyright 2019 The Google AI Language Team 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
"""Table retriever experiment."""

import csv
import functools
import os
import traceback
from typing import Text, Optional

from absl import app
from absl import flags
from tapas.models import table_retriever_model
from tapas.scripts import eval_table_retriever_utils
from tapas.utils import experiment_utils  # pylint: disable=unused-import
import tensorflow.compat.v1 as tf

tf.disable_v2_behavior()

FLAGS = flags.FLAGS

flags.DEFINE_string("data_format", "tfrecord", "The input data format.")

flags.DEFINE_multi_string(
    "input_file_train", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_multi_string(
    "input_file_eval", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_multi_string(
    "input_file_predict", None,
    "Input TF example files (can be a glob or comma separated).")

flags.DEFINE_string(
    "prediction_output_dir", None,
    "If not none or empty writes predictions to this directory. Otherwise "
    "writes predictions to model_dir.")

flags.DEFINE_string(
    "init_checkpoint", None,
    "Initial checkpoint (usually from a pre-trained BERT model).")

flags.DEFINE_integer(
    "max_seq_length", 128,
    "The maximum total input sequence length after WordPiece tokenization. "
    "Sequences longer than this will be truncated, and sequences shorter "
    "than this will be padded. Must match data generation.")

flags.DEFINE_integer("minutes_to_sleep_before_predictions", 5,
                     "Time in minutes to sleep before starting to predict.")

flags.DEFINE_bool("do_train", False, "Whether to run training.")

flags.DEFINE_bool("do_eval", False, "Whether to run eval on the dev set.")

flags.DEFINE_string(
    "eval_name", "default",
    "Name of the current evaluation set. Will be used in Tensorboard.")

flags.DEFINE_bool(
    "do_predict", False,
    "Whether to run the model in inference mode on the test set.")

flags.DEFINE_float(
    "grad_clipping", None, "If not None, gradients greater in absolute value"
    "than this number are clipped.")

flags.DEFINE_integer(
    "down_projection_dim", 0, "Representation dimension of the query/table"
    "after down projection. If smaller than 1, no projection occurs.")

flags.DEFINE_bool(
    "init_from_single_encoder", True, "If true, expects to load"
    "a checkpoint of a single encoder, that would be used to"
    "initialize both encoders.")

flags.DEFINE_integer("max_query_length", 128,
                     "The query is capped to this length.")

flags.DEFINE_string(
    "compression_type",
    "",
    "Compression to use when reading tfrecords. '' for no compression.",
)

flags.DEFINE_integer(
    "evaluated_checkpoint_step", None,
    "The step for a specific model checkpoint to be evaluated. If None, then"
    "all checkpoints are used.")

flags.DEFINE_string(
    "evaluated_checkpoint_metric", None,
    "The metric used to chose a model checkpoint to be evaluated. If None, then"
    "all checkpoints are used.")

flags.DEFINE_bool("use_out_of_core_negatives", False,
                  "If true, use all the negatives when"
                  "using many TPU cores.")

flags.DEFINE_bool("mask_repeated_tables", False,
                  "If true, mask tables that are repeated within a batch.")

flags.DEFINE_bool("mask_repeated_questions", False,
                  "If true, mask questions that are repeated within a batch.")

flags.DEFINE_bool(
    "ignore_table_content", False,
    "If true, use only the table headers to represent the table.")

flags.DEFINE_bool(
    "use_mined_negatives", False,
    "If true, use mined negatives that should be given as"
    "additional table features.")

flags.DEFINE_list("disabled_features", [],
                  "Features that should be disabled (for ablation studies).")


def _get_test_input_fn(name, input_file):
  """Gets input_fn for eval/predict modes."""
  if input_file is None:
    return None
  input_fn = functools.partial(
      table_retriever_model.input_fn,
      name=name,
      file_patterns=input_file,
      data_format=FLAGS.data_format,
      is_training=False,
      max_seq_length=FLAGS.max_seq_length,
      compression_type=FLAGS.compression_type,
      use_mined_negatives=FLAGS.use_mined_negatives,
      include_id=True)
  return input_fn


def _predict_and_export_metrics(
    mode,
    input_fn,
    checkpoint_path,
    step,
    estimator,
    output_dir,
):
  """Exports model predictions and calculates precision@k."""
  tf.logging.info("Running predictor for step %d.", step)
  result = estimator.predict(input_fn=input_fn, checkpoint_path=checkpoint_path)
  output_predict_file = os.path.join(output_dir, f"{mode}_results_{step}.tsv")
  write_predictions(result, output_predict_file)

  # Compute precision@k.
  if (not FLAGS.evaluated_checkpoint_step or
      not FLAGS.evaluated_checkpoint_metric):
    p_at_k = eval_table_retriever_utils.eval_precision_at_k(
        query_prediction_files=output_predict_file,
        table_prediction_files=output_predict_file,
        make_tables_unique=True)
    experiment_utils.save_metrics(output_dir, mode, step, p_at_k)


def write_predictions(predictions,
                      output_predict_file):
  """Writes predictions to an output TSV file.

  Predictions header: [query_id, query_rep, table_id, table_rep]
  Args:
    predictions: model predictions
    output_predict_file: Path for wrinting the predicitons.
  """
  with tf.io.gfile.GFile(output_predict_file, "w") as write_file:
    header = [
        "query_id",
        "query_rep",
        "table_id",
        "table_rep",
    ]
    writer = csv.DictWriter(write_file, fieldnames=header, delimiter="\t")
    writer.writeheader()

    for prediction in predictions:
      query_id = prediction["query_id"]
      table_id = prediction["table_id"]
      query_rep = prediction["query_rep"]
      table_rep = prediction["table_rep"]

      prediction_to_write = {
          "query_id": query_id[0].decode("utf-8"),
          "query_rep": query_rep.tolist(),
          "table_id": table_id[0].decode("utf-8"),
          "table_rep": table_rep.tolist(),
      }
      writer.writerow(prediction_to_write)


def main(_):
  bert_config = experiment_utils.bert_config_from_flags()
  total_steps = experiment_utils.num_train_steps()
  retriever_config = table_retriever_model.RetrieverConfig(
      bert_config=bert_config,
      init_checkpoint=FLAGS.init_checkpoint,
      learning_rate=FLAGS.learning_rate,
      num_train_steps=total_steps,
      num_warmup_steps=experiment_utils.num_warmup_steps(),
      use_tpu=FLAGS.use_tpu,
      grad_clipping=FLAGS.grad_clipping,
      down_projection_dim=FLAGS.down_projection_dim,
      init_from_single_encoder=FLAGS.init_from_single_encoder,
      max_query_length=FLAGS.max_query_length,
      mask_repeated_tables=FLAGS.mask_repeated_tables,
      mask_repeated_questions=FLAGS.mask_repeated_questions,
      use_out_of_core_negatives=FLAGS.use_out_of_core_negatives,
      ignore_table_content=FLAGS.ignore_table_content,
      disabled_features=FLAGS.disabled_features,
      use_mined_negatives=FLAGS.use_mined_negatives,
  )

  model_fn = table_retriever_model.model_fn_builder(retriever_config)
  estimator = experiment_utils.build_estimator(model_fn)

  if FLAGS.do_train:
    tf.io.gfile.makedirs(FLAGS.model_dir)
    bert_config.to_json_file(os.path.join(FLAGS.model_dir, "bert_config.json"))
    retriever_config.to_json_file(
        os.path.join(FLAGS.model_dir, "tapas_config.json"))
    train_input_fn = functools.partial(
        table_retriever_model.input_fn,
        name="train",
        file_patterns=FLAGS.input_file_train,
        data_format=FLAGS.data_format,
        is_training=True,
        max_seq_length=FLAGS.max_seq_length,
        compression_type=FLAGS.compression_type,
        use_mined_negatives=FLAGS.use_mined_negatives,
        include_id=False)
    estimator.train(input_fn=train_input_fn, max_steps=total_steps)

  eval_input_fn = _get_test_input_fn("eval", FLAGS.input_file_eval)
  if FLAGS.do_eval:
    if eval_input_fn is None:
      raise ValueError("No input_file_eval specified!")
    for _, checkpoint in experiment_utils.iterate_checkpoints(
        model_dir=estimator.model_dir,
        total_steps=total_steps,
        marker_file_prefix=os.path.join(estimator.model_dir,
                                        f"eval_{FLAGS.eval_name}"),
        minutes_to_sleep=FLAGS.minutes_to_sleep_before_predictions):
      tf.logging.info("Running eval: %s", FLAGS.eval_name)
      try:
        result = estimator.evaluate(
            input_fn=eval_input_fn,
            steps=FLAGS.num_eval_steps,
            name=FLAGS.eval_name,
            checkpoint_path=checkpoint)
        tf.logging.info("Eval result:\n%s", result)
      except (ValueError, tf.errors.NotFoundError):
        tf.logging.error("Error getting predictions for checkpoint %s: %s",
                         checkpoint, traceback.format_exc())

  if FLAGS.do_predict:
    predict_input_fn = _get_test_input_fn("predict", FLAGS.input_file_predict)
    if FLAGS.prediction_output_dir:
      prediction_output_dir = FLAGS.prediction_output_dir
      tf.io.gfile.makedirs(prediction_output_dir)
    else:
      prediction_output_dir = estimator.model_dir

    marker_file_prefix = os.path.join(prediction_output_dir, "predict")
    # When two separate jobs are launched we don't want conflicting markers.
    if predict_input_fn is not None:
      marker_file_prefix += "_test"
    if eval_input_fn is not None:
      marker_file_prefix += "_dev"

    single_step = FLAGS.evaluated_checkpoint_step
    if FLAGS.evaluated_checkpoint_metric:
      single_step = experiment_utils.get_best_step_for_metric(
          estimator.model_dir, FLAGS.evaluated_checkpoint_metric)
    for current_step, checkpoint in experiment_utils.iterate_checkpoints(
        model_dir=estimator.model_dir,
        total_steps=total_steps,
        marker_file_prefix=marker_file_prefix,
        single_step=single_step):
      try:
        if predict_input_fn is not None:
          _predict_and_export_metrics(
              mode="predict",
              input_fn=predict_input_fn,
              checkpoint_path=checkpoint,
              step=current_step,
              estimator=estimator,
              output_dir=prediction_output_dir)

        if eval_input_fn is not None:
          _predict_and_export_metrics(
              mode="eval",
              input_fn=eval_input_fn,
              checkpoint_path=checkpoint,
              step=current_step,
              estimator=estimator,
              output_dir=prediction_output_dir)
      except (ValueError, tf.errors.NotFoundError):
        tf.logging.error("Error getting predictions for checkpoint %s: %s",
                         checkpoint, traceback.format_exc())


if __name__ == "__main__":
  app.run(main)

tensorflow.compat.v1.disable_v2_behavior tensorflow.compat.v1.io.gfile.makedirs tensorflow.compat.v1.logging.info tensorflow.compat.v1.io.gfile.GFile