# coding=utf-8
# Copyright 2022 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.

"""Build FelixInsert BERT Examples from text (source, target) pairs."""

import collections
import itertools
from typing import List, Mapping, Optional, Sequence, Tuple

import tensorflow as tf

from felix import converter_for_felix_insert as converter
from felix import felix_constants as constants
from felix import tokenization
from felix import utils


class TaggingBertExample:
  """Class for training and inference examples for BERT tagging model.

  Attributes:
    features: A dictionary of features with numeral lists as values.
    features_float: A dictionary of features with float lists as values.
    scalar_features: A dictionary of features with scalar values.
  """

  def __init__(self,
               input_ids,
               input_mask,
               segment_ids,
               labels = None,
               labels_mask = None,
               input_tokens = None,
               label_tokens = None,
               source_text = None,
               target_text = None,
               target_tokens = None,
              ):
    """Initializes an instance of TaggingBertExample.

    Args:
      input_ids: List of ids of source tokens.
      input_mask: List of 1s and 0s. 0 indicates a PAD token.
      segment_ids: List of segment ids for BERT.
      labels: List of label ids. Optional if we are at test time.
      labels_mask: List of 1s and 0s. Optional if we are at test time.
      input_tokens: List of tokens (as text). Optional. For building an
        insertion example and for debugging.
      label_tokens: List of labels (as text). Optional. For building an
        insertion example and for debugging.
      source_text: Raw string of source input. Optional. For debug purposes.
      target_text: Raw string of target output. Optional. For debug purposes.
      target_tokens: List of target tokens (as text). Optional. For debug
        purposes.
    """
    # Check if the labels(_mask) is None or an empty list.
    if not labels or not labels_mask:
      self.features = collections.OrderedDict([
          ('input_ids', list(input_ids)),
          ('input_mask', list(input_mask)),
          ('segment_ids', list(segment_ids)),
      ])
      self.features_float = {}
    else:
      self.features = collections.OrderedDict([
          ('input_ids', list(input_ids)),
          ('input_mask', list(input_mask)),
          ('segment_ids', list(segment_ids)),
          ('labels', list(labels)),
      ])

    self.features_float = collections.OrderedDict([
        ('labels_mask', labels_mask),
    ])

    self.scalar_features = collections.OrderedDict()

    self._debug_features = collections.OrderedDict()
    if input_tokens is not None:
      self._debug_features['input_tokens'] = input_tokens
    if label_tokens is not None:
      self._debug_features['label_tokens'] = label_tokens
    if source_text is not None:
      self._debug_features['text_source'] = [source_text]
    if target_text is not None:
      self._debug_features['text_target'] = [target_text]
    if target_tokens is not None:
      self._debug_features['target_tokens'] = target_tokens

  def pad_to_max_length(self, max_seq_length, pad_token_id):
    """Pad the feature vectors so that they all have max_seq_length.

    Args:
      max_seq_length: The length that features will have after padding.
      pad_token_id: input_ids feature is padded with this ID, other features
        with ID 0.
    """
    pad_len = max_seq_length - len(self.features['input_ids'])
    for key, feature in itertools.chain(self.features.items(),
                                        self.features_float.items()):
      if feature is None:
        continue

      pad_id = pad_token_id if key == 'input_ids' else 0

      if feature is not None:
        feature.extend([pad_id] * pad_len)  # pytype: disable=attribute-error
      if len(feature) != max_seq_length:
        raise ValueError(f'{key} has length {len(feature)} (should be '
                         f'{max_seq_length}).')

  def _int_feature(self, values):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))

  def _float_feature(self, values):
    return tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))

  def to_tf_example(self):
    """Returns this object as a tf.Example."""

    tf_features = collections.OrderedDict([
        (key, self._int_feature(val)) for key, val in self.features.items()
    ])

    # Add scalar integer features.
    for key, value in self.scalar_features.items():
      tf_features[key] = self._int_feature([value])

    # Add label mask feature.
    for key, value in self.features_float.items():
      tf_features[key] = self._float_feature(value)

    # Add debug fields.
    for key, value in self._debug_features.items():
      tf_features[key] = tf.train.Feature(
          bytes_list=tf.train.BytesList(
              value=[element.encode('utf8') for element in value]))

    return tf.train.Example(features=tf.train.Features(feature=tf_features))


class FelixInsertExampleBuilder:
  """Builder class for FelixInsert tagging and insertion examples.

  Attributes:
    label_map: Mapping from tags to tag IDs.
    tokenizer: A tokenization.FullTokenizer, which converts between strings and
      lists of tokens.
  """

  def __init__(self,
               label_map,
               vocab_file,
               do_lower_case,
               max_seq_length,
               max_predictions_per_seq,
               max_insertions_per_token,
               insert_after_token = True,
               special_glue_string_for_sources = None):
    """Initializes an instance of BertExampleBuilder.

    Args:
      label_map: Mapping from tags represented as (base_tag, num_insertions)
        tuples to tag IDs.
      vocab_file: Path to vocab file.
      do_lower_case: should text be lowercased.
      max_seq_length: Maximum sequence length.
      max_predictions_per_seq: Maximum number of tokens to insert per input.
      max_insertions_per_token: Maximum number of tokens/masks to insert per
        token.
      insert_after_token: Whether to insert tokens after the current token.
      special_glue_string_for_sources: If there are multiple sources, this
        string is used to combine them into one string. The empty string is a
        valid value. Optional.
    """
    self.label_map = label_map
    self.tokenizer = tokenization.FullTokenizer(
        vocab_file, do_lower_case=do_lower_case)
    self._max_seq_length = max_seq_length
    self._max_predictions_per_seq = max_predictions_per_seq
    self._max_insertions_per_token = max_insertions_per_token
    self._insert_after_token = insert_after_token
    try:
      self._pad_id = self.tokenizer.convert_tokens_to_ids([constants.PAD])[0]
    except KeyError:
      self._pad_id = 0
    if special_glue_string_for_sources is not None:
      self._special_glue_string_for_sources = special_glue_string_for_sources
    else:
      self._special_glue_string_for_sources = ' '

  def build_bert_example(
      self,
      sources,
      target = None,
      is_test_time = False
  ):
    """Constructs a tagging and an insertion BERT Example.

    Args:
      sources: List of source texts.
      target: Target text or None when building an example during inference. If
        the target is None then we don't calculate gold labels or tags, this is
        equivaltn to setting is_test_time to True.
      is_test_time: Controls whether the dataset is to be used at test time.
        Unlike setting target = None to indicate test time, this flags allows
        for saving the target in the tfrecord.  For compatibility with old
        scripts, setting target to None has the same behavior as setting
        is_test_time to True.

    Returns:
      A tuple with:
      1. TaggingBertExample (or None if more than
      `self._max_insertions_per_token` insertions are required).
      2. A feed_dict object for creating the insertion BERT example (or None if
      `target` is None, `is_test_time` is True, or the above TaggingBertExample
      is None.

    Raises:
      KeyError: If a label not in `self.label_map` is produced.
    """
    merged_sources = self._special_glue_string_for_sources.join(sources).strip()
    input_tokens = self._tokenize_text(merged_sources)

    input_ids = self.tokenizer.convert_tokens_to_ids(input_tokens)
    input_mask = [1] * len(input_ids)
    segment_ids = [0] * len(input_ids)
    if target is None or is_test_time:
      example = TaggingBertExample(
          input_ids=input_ids,
          input_mask=input_mask,
          segment_ids=segment_ids,
          labels=None,
          labels_mask=None,
          input_tokens=input_tokens,
          source_text=merged_sources,
          target_text=target)
      example.pad_to_max_length(self._max_seq_length, self._pad_id)
      return example, None

    output_tokens = self._tokenize_text(target)
    edits_and_insertions = converter.compute_edits_and_insertions(
        input_tokens, output_tokens, self._max_insertions_per_token,
        self._insert_after_token)
    if edits_and_insertions is None:
      return None, None
    else:
      edits, insertions = edits_and_insertions

    label_tokens = []  # Labels as strings.
    label_tuples = []  # Labels as (base_tag, num_insertions) tuples.
    labels = []  # Labels as IDs.
    for edit, insertion in zip(edits, insertions):
      label_token = edit
      if insertion:
        label_token += f'|{len(insertion)}'
      label_tokens.append(label_token)
      label_tuple = (edit, len(insertion))
      label_tuples.append(label_tuple)
      if label_tuple in self.label_map:
        labels.append(self.label_map[label_tuple])
      else:
        raise KeyError(
            f"Label map doesn't contain a computed label: {label_tuple}")

    label_counter = collections.Counter(labels)
    label_weight = {
        label: len(labels) / count / len(label_counter)
        for label, count in label_counter.items()
    }
    # Weight the labels inversely proportional to their frequency.
    labels_mask = [label_weight[label] for label in labels]
    if self._insert_after_token:
      # When inserting after the current token, we never need to insert after
      # the final [SEP] token and thus the edit label for that token is constant
      # ('KEEP') and could be excluded from loss computations.
      labels_mask[-1] = 0
    else:
      # When inserting before the current token, the first edit is constant.
      labels_mask[0] = 0

    example = TaggingBertExample(
        input_ids=input_ids,
        input_mask=input_mask,
        segment_ids=segment_ids,
        labels=labels,
        labels_mask=labels_mask,
        input_tokens=input_tokens,
        label_tokens=label_tokens,
        source_text=merged_sources,
        target_text=target)
    example.pad_to_max_length(self._max_seq_length, self._pad_id)

    insertion_example = self.build_insertion_example(
        source_tokens=input_tokens,
        labels=label_tuples,
        target_insertions=insertions)
    return example, insertion_example

  def build_insertion_tokens(
      self,
      source_tokens,
      labels,
      target_insertions = None,
  ):
    """Constructs the masked input TF Example for the insertion model.

    Args:
      source_tokens: List of source tokens.
      labels: List of edit label tuples (base_tag, num_insertions).
      target_insertions: Inserted target tokens per source token. Only provided
        when constructing training examples.

    Returns:
      A tuple with a masked insertion token sequence and the unmasked target
      sequence.

    Raises:
      ValuError: If the number of insertions in a label doesn't match the
        number of elements in the corresponding target_insertions[i] list.
    """
    masked_tokens = []
    target_tokens = []
    deletion_started = False
    # Tokens to add after finishing the deleted span.
    mask_buffer = []
    target_buffer = []
    for i, (source_token,
            (base_tag,
             num_insertions)) in enumerate(zip(source_tokens, labels)):
      masks = [constants.MASK] * num_insertions
      if target_insertions is not None:
        insertions = target_insertions[i]
        if len(insertions) != num_insertions:
          raise ValueError(f'Inserted token counts do not match at {i}: '
                           f'{len(insertions)} vs. {num_insertions}')
      else:
        insertions = masks

      if base_tag == constants.KEEP:
        if deletion_started:
          deletion_started = False
          masked_tokens.extend([constants.DELETE_SPAN_END] + mask_buffer)
          target_tokens.extend([constants.DELETE_SPAN_END] + target_buffer)
          mask_buffer = []
          target_buffer = []
        if self._insert_after_token:
          masked_tokens.extend([source_token] + masks)
          target_tokens.extend([source_token] + insertions)
        else:
          masked_tokens.extend(masks + [source_token])
          target_tokens.extend(insertions + [source_token])
      else:  # base_tag == constants.DELETE
        if not deletion_started:
          masked_tokens.append(constants.DELETE_SPAN_START)
          target_tokens.append(constants.DELETE_SPAN_START)
        deletion_started = True
        mask_buffer.extend(masks)
        target_buffer.extend(insertions)
        masked_tokens.append(source_token)
        target_tokens.append(source_token)
    if deletion_started:
      # Source ended with deletion so we need to add the buffered tokens.
      masked_tokens.extend([constants.DELETE_SPAN_END] + mask_buffer)
      target_tokens.extend([constants.DELETE_SPAN_END] + target_buffer)

    assert len(masked_tokens) == len(target_tokens), (
        f"Masked token count ({len(masked_tokens)}) doesn't match the "
        f"target token count ({len(target_tokens)}).")

    # Truncate lists. Note that although `source_tokens` is already guaranteed
    # to be truncated, `masked_tokens` can still be too long since it contains
    # added MASK tokens.
    masked_tokens = masked_tokens[:self._max_seq_length]
    target_tokens = target_tokens[:self._max_seq_length]
    # Don't truncate [SEP] to match BERT input format.
    if masked_tokens[-1] != constants.SEP:
      masked_tokens[-1] = constants.SEP
      target_tokens[-1] = constants.SEP

    return masked_tokens, target_tokens

  def build_insertion_example(
      self,
      source_tokens,
      labels,
      target_insertions = None
  ):
    """Constructs the masked input TF Example for the insertion model.

    Args:
      source_tokens: List of source tokens.
      labels: List of edit label tuples (base_tag, num_insertions).
      target_insertions: Inserted target tokens per source token. Only provided
        when constructing training examples.

    Returns:
      A feed_dict containing input features to be fed to a predictor (for
      inference) or to be converted to a tf.Example (for model training). If the
      labels don't contain any insertions, returns None.
    """
    masked_tokens, target_tokens = self.build_insertion_tokens(
        source_tokens, labels, target_insertions)
    if constants.MASK not in masked_tokens:
      # No need to create an insertion example so return None.
      return None
    return utils.build_feed_dict(
        masked_tokens,
        self.tokenizer,
        target_tokens=target_tokens,
        max_seq_length=self._max_seq_length,
        max_predictions_per_seq=self._max_predictions_per_seq)

  def _tokenize_text(self, text):
    """Returns the tokenized text with special start and end tokens."""
    tokens = list(self.tokenizer.tokenize(text))
    # Truncate list and add [CLS] and [SEP] tokens to make the input resemble
    # the data BERT has seen during pretraining and to allow inserting tokens to
    # the beginning and to the end of the sequence regardless of
    # `self._insert_after_token`.
    tokens = tokens[:self._max_seq_length - 2]
    return [constants.CLS] + tokens + [constants.SEP]
