# 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
"""TAPAS BERT model for classification."""

import dataclasses
import enum
import json
from typing import Iterable, List, Mapping, Optional, Set, Text

from tapas.datasets import dataset
from tapas.datasets import table_dataset
from tapas.models import segmented_tensor
from tapas.models import tapas_classifier_model_utils as utils
from tapas.models.bert import modeling
from tapas.models.bert import optimization
from tapas.models.bert import table_bert
from tapas.utils import attention_utils
from tapas.utils import span_prediction_utils
from tapas.utils import table_pruning
import tensorflow.compat.v1 as tf
import tensorflow_probability as tfp

SpanPredictionMode = span_prediction_utils.SpanPredictionMode

_EPSILON_ZERO_DIVISION = utils.EPSILON_ZERO_DIVISION
_CLOSE_ENOUGH_TO_LOG_ZERO = utils.CLOSE_ENOUGH_TO_LOG_ZERO
_classification_initializer = utils.classification_initializer


class AverageApproximationFunction(str, enum.Enum):
  RATIO = "ratio"
  FIRST_ORDER = "first_order"
  SECOND_ORDER = "second_order"


@dataclasses.dataclass
class TapasClassifierConfig:
  """Helper class for configuration of Tapas model.

  bert_config: Config object for general bert hyper params.
  init_checkpoint: Location of the model checkpoint.
  learning_rate: Optimizer learning rate.
  num_train_steps: Total number of training steps for optimizer schedule.
  num_warmup_steps: Number of training steps to warm up optimizer.
  use_tpu: Use TPU for training.
  positive_weight: Weight for positive labels.
  num_aggregation_labels: The number of aggregation classes to predict.
  num_classification_labels: The number of classes to predict.
  aggregation_loss_importance: Importance weight for the aggregation loss.
  use_answer_as_supervision: Whether to use the answer as the only supervision
    for aggregation examples.
  answer_loss_importance: Importance weight for the regression loss.
  use_normalized_answer_loss: Normalize loss by max of predicted and expected
    value.
  huber_loss_delta: Delta parameter used to calculate the regression loss.
  temperature: Scales cell logits to control the skewness of probabilities.
  agg_temperature: Scales aggregation logits to control the skewness of
    probabilities.
  use_gumbel_for_cells: Applies Gumbel-Softmax to cell selection.
  use_gumbel_for_agg: Applies Gumbel-Softmax to aggregation selection.
  average_approximation_function: Method to calculate expected average of
    cells in the relaxed case.
  cell_select_pref: Preference for cell selection in ambiguous cases.
  answer_loss_cutoff: Ignore examples with answer loss larger than cutoff.
  grad_clipping: If not None, clip the gradient norm to this value
  max_num_rows: Maximum number of rows.
  max_num_columns: Maximum number of columns.
  average_logits_per_cell: Wheher to average logits per cell.
  select_one_column: Whether to constrain the model to only select cells from
    a single column.
  allow_empty_column_selection: Allow not to select any column.
  disabled_features: Set of embeddings ids to disable in the input layer.
    Posible values are "segment_ids", "column_ids", "row_ids", "prev_label_ids",
    "column_ranks", "inv_column_ranks", "numeric_relations"
  init_cell_selection_weights_to_zero: Whether to initialize cell selection.
    weights to 0 so that the initial probabilities are 50%.
  disable_position_embeddings: Disable positional embeddings in the input layer.
  reset_position_index_per_cell: Restart position indexes at every cell.
  disable_per_token_loss: Disable any (strong or weak) supervision on cells.
  span_prediction: Span selection mode to use.
  proj_value_length: Optional. down-project keys and values when computing
    self-attention. Following https://arxiv.org/pdf/2006.04768.pdf.
  reset_output_cls: If true, reset classification output layer on init from cp.
  classification_label_weight: Per label loss multiplier for classification.
  mask_examples_without_labels: This is useful for e2e retrieval where we do
  want to train on negative tables but we don't want the model to learn to
  predict no answer for them.
  table_pruning_config_file: The config file of table pruning.
  restrict_attention_mode: Restrict attention between cells depending on row or
    column position.
  restrict_attention_bucket_size: For sparse attention modes, further restrict
    attention to consecutive buckets of uniform size.
  restrict_attention_header_size: For sparse attention modes, size of the first
    section that will attend to/from everything else. If None use bucket_size.
  restrict_attention_row_heads_ratio: For sparse attention modes, proportion of
    heads that should focus on rows vs columns.
  cell_cross_entropy: Whether to use cross entropy for cell selection loss.
  cell_cross_entropy_hard_em: When using cell cross entropy, whether to use hard
    or soft EM for cases with more than one ground truth cell.
  """

  bert_config: modeling.BertConfig
  init_checkpoint: Text
  learning_rate: float
  num_train_steps: Optional[int]
  num_warmup_steps: Optional[int]
  use_tpu: bool
  positive_weight: float
  num_aggregation_labels: int
  num_classification_labels: int
  aggregation_loss_importance: float
  use_answer_as_supervision: bool
  answer_loss_importance: Optional[float]
  use_normalized_answer_loss: bool
  huber_loss_delta: Optional[float]
  temperature: float
  agg_temperature: float
  use_gumbel_for_cells: bool
  use_gumbel_for_agg: bool
  average_approximation_function: AverageApproximationFunction
  cell_select_pref: Optional[float]
  answer_loss_cutoff: Optional[float]
  grad_clipping: Optional[float]
  max_num_rows: int
  max_num_columns: int
  average_logits_per_cell: bool
  select_one_column: bool
  allow_empty_column_selection: bool = True
  disabled_features: Optional[List[Text]] = None
  init_cell_selection_weights_to_zero: bool = False
  disable_position_embeddings: bool = False
  reset_position_index_per_cell: bool = False
  disable_per_token_loss: bool = False
  span_prediction: SpanPredictionMode = SpanPredictionMode.NONE
  proj_value_length: Optional[int] = None
  reset_output_cls: bool = False
  classification_label_weight: Optional[Mapping[int, float]] = None
  mask_examples_without_labels: bool = False
  table_pruning_config_file: Optional[Text] = None
  restrict_attention_mode: attention_utils.RestrictAttentionMode = (
      attention_utils.RestrictAttentionMode.FULL)
  restrict_attention_bucket_size: int = 0
  restrict_attention_header_size: Optional[int] = None
  restrict_attention_row_heads_ratio: float = 0.5
  cell_cross_entropy: bool = False
  cell_cross_entropy_hard_em: bool = False

  def to_json_string(self):
    """Serializes this instance to a JSON string."""

    class EnhancedJSONEncoder(json.JSONEncoder):

      def default(self, o):
        if dataclasses.is_dataclass(o):
          return dataclasses.asdict(o)
        if isinstance(o, modeling.BertConfig):
          return o.to_dict()
        return super().default(o)

    return json.dumps(self, indent=2, sort_keys=True, cls=EnhancedJSONEncoder)

  def to_json_file(self, json_file):
    """Serializes this instance to a JSON file."""
    with tf.io.gfile.GFile(json_file, "w") as writer:
      writer.write(self.to_json_string() + "\n")

  @classmethod
  def from_dict(cls, json_object, for_prediction=False):
    """Constructs a config from a Python dictionary of parameters."""
    json_object = dict(json_object)
    # Overwrite json bert config with config object.
    json_object["bert_config"] = modeling.BertConfig.from_dict(
        json_object["bert_config"])
    # Delete deprecated option, if present.
    # TODO See of we can filter everything that's not an argument.
    if "restrict_attention" in json_object:
      del json_object["restrict_attention"]
    if for_prediction:
      # Disable training-only option to reduce input requirements.
      json_object["use_answer_as_supervision"] = False
      json_object["init_checkpoint"] = None
    return TapasClassifierConfig(**json_object)

  @classmethod
  def from_json_file(cls, json_file, for_prediction=False):
    """Constructs a `BertConfig` from a json file of parameters."""
    with tf.io.gfile.GFile(json_file, "r") as reader:
      text = reader.read()
    return cls.from_dict(json.loads(text), for_prediction=for_prediction)


def _get_probs(dist):
  """Helper to extract probability from a distribution."""
  # In tensorflow_probabiliy 0.7 this attribute was filled on __init__ method
  if dist.probs is not None:
    return dist.probs
  # From 0.8 onwards the probs is not filled and a function should be used
  return dist.probs_parameter()


def _calculate_aggregation_logits(output_layer_aggregation, output_weights_agg,
                                  output_bias_agg):
  """Calculates the aggregation logits.

  Args:
    output_layer_aggregation: <float32>[batch_size, hidden_size]
    output_weights_agg: <float32>[num_aggregation_labels, hidden_size_agg]
    output_bias_agg: <float32>[num_aggregation_labels]

  Returns:
    logits_aggregation: <float32>[batch_size, num_aggregation_labels]
  """
  logits_aggregation = tf.matmul(
      output_layer_aggregation, output_weights_agg, transpose_b=True)
  logits_aggregation = tf.nn.bias_add(logits_aggregation, output_bias_agg)
  return logits_aggregation


def _calculate_aggregation_loss_known(logits_aggregation, aggregate_mask,
                                      aggregation_function_id,
                                      config):
  """Calculates aggregation loss when its type is known during training.

  In the weakly supervised setting, the only known information is that for
  cell selection examples, "no aggregation" should be predicted. For other
  examples (those that require aggregation), no loss is accumulated.
  In the setting where aggregation type is always known, standard cross entropy
  loss is accumulated for all examples.

  Args:
    logits_aggregation: <float32>[batch_size, num_aggregation_labels]
    aggregate_mask: <float32>[batch_size]
    aggregation_function_id: <int32>[batch_size]
    config: Configuration for Tapas model.

  Returns:
    aggregation_loss_known: <float32>[batch_size, num_aggregation_labels]
  """
  if config.use_answer_as_supervision:
    # Prepare "no aggregation" targets for cell selection examples.
    target_aggregation = tf.zeros_like(aggregate_mask, dtype=tf.int32)
  else:
    # Use aggregation supervision as the target.
    target_aggregation = aggregation_function_id

  one_hot_labels = tf.one_hot(
      target_aggregation, depth=config.num_aggregation_labels, dtype=tf.float32)
  log_probs = tf.nn.log_softmax(logits_aggregation, axis=-1)
  # <float32>[batch_size]
  per_example_aggregation_intermediate = -tf.reduce_sum(
      one_hot_labels * log_probs, axis=-1)
  if config.use_answer_as_supervision:
    # Accumulate loss only for examples requiring cell selection
    # (no aggregation).
    return per_example_aggregation_intermediate * (1 - aggregate_mask)
  else:
    return per_example_aggregation_intermediate


def _calculate_aggregation_loss_unknown(logits_aggregation, aggregate_mask):
  """Calculates aggregation loss in the case of answer supervision."""
  dist_aggregation = tfp.distributions.Categorical(logits=logits_aggregation)
  # Index 0 correponds to "no aggregation".
  aggregation_ops_total_mass = tf.reduce_sum(
      _get_probs(dist_aggregation)[:, 1:], axis=1)
  # Predict some aggregation in case of an answer that needs aggregation.
  # This increases the probability of all aggregation functions, in a way
  # similar to MML, but without considering whether the function gives the
  # correct answer.
  return -tf.log(aggregation_ops_total_mass) * aggregate_mask


def _calculate_aggregation_loss(logits_aggregation, aggregate_mask,
                                aggregation_function_id,
                                config):
  """Calculates the aggregation loss per example."""
  per_example_aggregation_loss = _calculate_aggregation_loss_known(
      logits_aggregation, aggregate_mask, aggregation_function_id, config)

  if config.use_answer_as_supervision:
    # Add aggregation loss for numeric answers that need aggregation.
    per_example_aggregation_loss += _calculate_aggregation_loss_unknown(
        logits_aggregation, aggregate_mask)
  return config.aggregation_loss_importance * per_example_aggregation_loss


def _calculate_expected_result(dist_per_cell, numeric_values,
                               numeric_values_scale, input_mask_float,
                               logits_aggregation,
                               config):
  """Calculate the expected result given cell and aggregation probabilities."""
  if config.use_gumbel_for_cells:
    gumbel_dist = tfp.distributions.RelaxedBernoulli(
        # The token logits where already divided by the temperature and used for
        # computing cell selection errors so we need to multiply it again here
        config.temperature,
        logits=dist_per_cell.logits_parameter() * config.temperature)
    scaled_probability_per_cell = gumbel_dist.sample()
  else:
    scaled_probability_per_cell = _get_probs(dist_per_cell)

  # <float32>[batch_size, seq_length]
  scaled_probability_per_cell = (scaled_probability_per_cell /
                                 numeric_values_scale) * input_mask_float
  count_result = tf.reduce_sum(scaled_probability_per_cell, axis=1)
  numeric_values_masked = tf.where(
      tf.is_nan(numeric_values), tf.zeros_like(numeric_values),
      numeric_values)  # Mask non-numeric table values to zero.
  sum_result = tf.reduce_sum(
      scaled_probability_per_cell * numeric_values_masked, axis=1)
  avg_approximation = config.average_approximation_function
  if avg_approximation == AverageApproximationFunction.RATIO:
    average_result = sum_result / (count_result + _EPSILON_ZERO_DIVISION)
  elif avg_approximation == AverageApproximationFunction.FIRST_ORDER:
    # The sum of all probabilities exept that correspond to other cells
    ex = (
        tf.reduce_sum(scaled_probability_per_cell, axis=1, keepdims=True) -
        scaled_probability_per_cell + 1)
    average_result = tf.reduce_sum(
        numeric_values_masked * scaled_probability_per_cell / ex, axis=1)
  elif avg_approximation == AverageApproximationFunction.SECOND_ORDER:
    # The sum of all probabilities exept that correspond to other cells
    ex = (
        tf.reduce_sum(scaled_probability_per_cell, axis=1, keepdims=True) -
        scaled_probability_per_cell + 1)
    pointwise_var = (
        scaled_probability_per_cell * (1 - scaled_probability_per_cell))
    var = tf.reduce_sum(pointwise_var, axis=1, keepdims=True) - pointwise_var
    multiplier = (var / tf.math.square(ex) + 1) / ex
    average_result = tf.reduce_sum(
        numeric_values_masked * scaled_probability_per_cell * multiplier,
        axis=1)
  else:
    tf.logging.error("Invalid average_approximation_function: %s",
                     config.average_approximation_function)

  if config.use_gumbel_for_agg:
    gumbel_dist = tfp.distributions.RelaxedOneHotCategorical(
        config.agg_temperature, logits=logits_aggregation[:, 1:])
    # <float32>[batch_size, num_aggregation_labels - 1]
    aggregation_op_only_probs = gumbel_dist.sample()
  else:
    # <float32>[batch_size, num_aggregation_labels - 1]
    aggregation_op_only_probs = tf.nn.softmax(
        logits_aggregation[:, 1:] / config.agg_temperature, axis=-1)
  all_results = tf.concat([
      tf.expand_dims(sum_result, axis=1),
      tf.expand_dims(average_result, axis=1),
      tf.expand_dims(count_result, axis=1)
  ],
                          axis=1)
  expected_result = tf.reduce_sum(
      all_results * aggregation_op_only_probs, axis=1)
  return expected_result


def _calculate_regression_loss(answer, aggregate_mask, dist_per_cell,
                               numeric_values, numeric_values_scale,
                               input_mask_float, logits_aggregation,
                               config):
  """Calculates the regression loss per example.

  Args:
    answer: <float32>[batch_size]
    aggregate_mask: <float32>[batch_size]
    dist_per_cell: Cell selection distribution for each cell.
    numeric_values: <float32>[batch_size, seq_length]
    numeric_values_scale: <float32>[batch_size, seq_length]
    input_mask_float: <float32>[batch_size, seq_length]
    logits_aggregation: <float32>[batch_size, num_aggregation_labels]
      probabilities.
    config: Configuration for Tapas model.

  Returns:
    per_example_answer_loss_scaled: <float32>[batch_size]. Scales answer loss
    for each example in the batch.
    large_answer_loss_mask: <float32>[batch_size]. A mask which is 1 for
    examples for which their answer loss is larger than the answer_loss_cutoff.
  """
  # <float32>[batch_size]
  expected_result = _calculate_expected_result(dist_per_cell, numeric_values,
                                               numeric_values_scale,
                                               input_mask_float,
                                               logits_aggregation, config)
  # <float32>[batch_size]
  answer_masked = tf.where(tf.is_nan(answer), tf.zeros_like(answer), answer)

  if config.use_normalized_answer_loss:
    normalizer = tf.stop_gradient(
        tf.math.maximum(
            tf.math.abs(expected_result), tf.math.abs(answer_masked)) +
        _EPSILON_ZERO_DIVISION)
    normalized_answer_masked = answer_masked / normalizer
    normalized_expected_result = expected_result / normalizer
    per_example_answer_loss = tf.losses.huber_loss(
        normalized_answer_masked * aggregate_mask,
        normalized_expected_result * aggregate_mask,
        delta=tf.cast(config.huber_loss_delta, tf.float32),
        reduction=tf.losses.Reduction.NONE)
  else:
    per_example_answer_loss = tf.losses.huber_loss(
        answer_masked * aggregate_mask,
        expected_result * aggregate_mask,
        delta=tf.cast(config.huber_loss_delta, tf.float32),
        reduction=tf.losses.Reduction.NONE)
  if config.answer_loss_cutoff is None:
    large_answer_loss_mask = tf.ones_like(
        per_example_answer_loss, dtype=tf.float32)
  else:
    large_answer_loss_mask = tf.where(
        per_example_answer_loss > config.answer_loss_cutoff,
        tf.zeros_like(per_example_answer_loss, dtype=tf.float32),
        tf.ones_like(per_example_answer_loss, dtype=tf.float32))
  per_example_answer_loss_scaled = config.answer_loss_importance * (
      per_example_answer_loss * aggregate_mask)
  return per_example_answer_loss_scaled, large_answer_loss_mask


def _calculate_aggregate_mask(answer, output_layer_aggregation, output_bias_agg,
                              output_weights_agg, cell_select_pref, label_ids):
  """Finds examples where the model should select cells with no aggregation.

  Returns a mask that determines for which examples should the model select
  answers directly from the table, without any aggregation function. If the
  answer is a piece of text the case is unambiguous as aggregation functions
  only apply to numbers. If the answer is a number but does not appear in the
  table then we must use some aggregation case. The ambiguous case is when the
  answer is a number that also appears in the table. In this case we use the
  aggregation function probabilities predicted by the model to decide whether
  to select or aggregate. The threshold for this is a hyperparameter
  `cell_select_pref`.

  Args:
    answer: <float32>[batch_size]
    output_layer_aggregation: <float32>[batch_size, hidden_size]
    output_bias_agg: <float32>[num_aggregation_labels]
    output_weights_agg: <float32>[num_aggregation_labels, hidden_size_agg]
    cell_select_pref: Preference for cell selection in ambiguous cases.
    label_ids: <int32>[batch_size, seq_length]

  Returns:
    aggregate_mask: <float32>[batch_size] A mask set to 1 for examples that
      should use aggregation functions.
  """
  # <float32>[batch_size]
  aggregate_mask_init = tf.cast(tf.logical_not(tf.is_nan(answer)), tf.float32)
  logits_aggregation = _calculate_aggregation_logits(output_layer_aggregation,
                                                     output_weights_agg,
                                                     output_bias_agg)
  dist_aggregation = tfp.distributions.Categorical(logits=logits_aggregation)
  aggregation_ops_total_mass = tf.reduce_sum(
      _get_probs(dist_aggregation)[:, 1:], axis=1)
  # Cell selection examples according to current model.
  is_pred_cell_selection = aggregation_ops_total_mass <= cell_select_pref
  # Examples with non-empty cell selection supervision.
  is_cell_supervision_available = tf.reduce_sum(label_ids, axis=1) > 0
  aggregate_mask = tf.where(
      tf.logical_and(is_pred_cell_selection, is_cell_supervision_available),
      tf.zeros_like(aggregate_mask_init, dtype=tf.float32), aggregate_mask_init)
  aggregate_mask = tf.stop_gradient(aggregate_mask)
  return aggregate_mask


def compute_classification_logits(num_classification_labels, output_layer):
  """Computes logits for each classification of the sequence.

  Args:
    num_classification_labels: int Number of class to predict
    output_layer: <float>[batch_size, hidden_dim] Output of the encoder layer.

  Returns:
    <float>[batch_size, num_classification_labels] Logits per class.
  """
  hidden_size_agg = output_layer.shape[-1].value
  output_weights_cls = tf.get_variable(
      "output_weights_cls",
      shape=[num_classification_labels, hidden_size_agg],
      initializer=_classification_initializer())
  output_bias_cls = tf.get_variable(
      "output_bias_cls",
      shape=[num_classification_labels],
      initializer=tf.zeros_initializer())
  logits_cls = tf.matmul(output_layer, output_weights_cls, transpose_b=True)
  logits_cls = tf.nn.bias_add(logits_cls, output_bias_cls)
  return logits_cls


def _single_column_cell_selection_loss(token_logits, column_logits, label_ids,
                                       cell_index, col_index, cell_mask):
  """Computes the loss for cell selection constrained to a single column.

  The loss is a hierarchical log-likelihood. The model first predicts a column
  and then selects cells within that column (conditioned on the column). Cells
  outside the selected column are never selected.

  Args: