# Copyright 2017 Google LLC.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from this
#    software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
"""Data providers for deepvariant images.

tf.data.Dataset and data providers for standard DeepVariant datasets for
training and evaluating germline calling accuracy.
"""


from absl import logging
import tensorflow as tf
from tensorflow import estimator as tf_estimator

from deepvariant import dv_constants
from deepvariant import dv_utils
from deepvariant.protos import deepvariant_pb2
from google.protobuf import text_format
from third_party.nucleus.io import sharded_file_utils


# These are empirically determined to work well on TPU with our data sets,
# where lots of buffering and concurrency is necessary to keep the device
# busy.
# These are settable in the constructor.
_DEFAULT_INPUT_READ_THREADS = 32
_DEFAULT_SHUFFLE_BUFFER_ELEMENTS = 100
_DEFAULT_INITIAL_SHUFFLE_BUFFER_ELEMENTS = 1024
_DEFAULT_PREFETCH_BUFFER_BYTES = 16 * 1000 * 1000


class DeepVariantInput(object):
  """This class serves as an `input_fn` for the `tf.estimator` framework."""

  # Calling this object like a function returns a stream of variadic tuples.
  # Essentially it is a buffered io library, that handles concurrently
  # reading and possibly shuffling input records from a set of files. It
  # knows how to parse features we care about from tf.examples. It records
  # some extra information about the source of the input, such as the name
  # and number of classes.

  def __init__(
      self,
      mode,
      input_file_spec,
      num_examples=None,
      num_classes=dv_constants.NUM_CLASSES,
      max_examples=None,
      tensor_shape=None,
      name=None,
      use_tpu=False,
      input_read_threads=_DEFAULT_INPUT_READ_THREADS,
      shuffle_buffer_size=_DEFAULT_SHUFFLE_BUFFER_ELEMENTS,
      initial_shuffle_buffer_size=_DEFAULT_INITIAL_SHUFFLE_BUFFER_ELEMENTS,
      prefetch_dataset_buffer_size=_DEFAULT_PREFETCH_BUFFER_BYTES,
      sloppy=True,
      list_files_shuffle=True,
      debugging_true_label_mode=False):
    """Create an DeepVariantInput object, usable as an `input_fn`.

    Args:
      mode: the mode string (from `tf.estimator.ModeKeys`).
      input_file_spec: the input filename for a tfrecord[.gz] file containing
        examples.  Can contain sharding designators.
      num_examples: the number of examples contained in the input file. Required
        for setting learning rate schedule in train/eval only.
      num_classes: The number of classes in the labels of this dataset.
        Currently defaults to DEFAULT_NUM_CLASSES.
      max_examples: The maximum number of examples to use. If None, all examples
        will be used. If not None, the first n = min(max_examples, num_examples)
        will be used. This works with training, and the n examples will repeat
        over and over.
      tensor_shape: None (which means we get the shape from the first example in
        source), or list of int [height, width, channel] for testing.
      name: string, name of the dataset.
      use_tpu: use code paths tuned for TPU, in particular protobuf encoding.
        Default False.
      input_read_threads: number of threads for reading data.  Default 32.
      shuffle_buffer_size: size of the final shuffle buffer, in elements.
        Default 100.
      initial_shuffle_buffer_size: int; the size of the dataset.shuffle buffer
        in elements.  Default is 1024.
      prefetch_dataset_buffer_size: int; the size of the TFRecordDataset buffer
        in bytes.  Default is 16 * 1000 * 1000.
      sloppy: boolean, allow parallel_interleave to be sloppy.  Default True.
      list_files_shuffle: boolean, allow list_files to shuffle.  Default True.
      debugging_true_label_mode: boolean. If true, the input examples are
        created with "training" mode. We'll parse the 'label' field even if the
        `mode` is PREDICT.

    Raises:
      ValueError: if `num_examples` not provided, in a context requiring it.
    """
    self.mode = mode
    self.input_file_spec = input_file_spec
    self.name = name
    self.num_examples = num_examples
    self.num_classes = num_classes
    self.max_examples = max_examples

    self.use_tpu = use_tpu
    self.sloppy = sloppy
    self.list_files_shuffle = list_files_shuffle
    self.input_read_threads = input_read_threads
    self.shuffle_buffer_size = shuffle_buffer_size
    self.initial_shuffle_buffer_size = initial_shuffle_buffer_size
    self.prefetch_dataset_buffer_size = prefetch_dataset_buffer_size
    self.debugging_true_label_mode = debugging_true_label_mode
    self.feature_extraction_spec = self.features_extraction_spec_for_mode(
        mode in (tf_estimator.ModeKeys.TRAIN, tf_estimator.ModeKeys.EVAL) or
        debugging_true_label_mode)

    if num_examples is None and mode != tf_estimator.ModeKeys.PREDICT:
      raise ValueError('num_examples argument required for DeepVariantInput'
                       'in TRAIN/EVAL modes.')

    if max_examples is not None:
      if max_examples <= 0:
        raise ValueError(
            'max_examples must be > 0 if not None. Got {}'.format(max_examples))
      # We update our num_examples in the situation where num_examples is set
      # (i.e., is not None) to the smaller of max_examples and num_examples.
      if self.num_examples is not None:
        self.num_examples = min(max_examples, self.num_examples)

    if tensor_shape:
      self.tensor_shape = tensor_shape
    else:
      self.tensor_shape = dv_utils.get_shape_from_examples_path(input_file_spec)
    self.input_files = sharded_file_utils.glob_list_sharded_file_patterns(
        self.input_file_spec)

  def features_extraction_spec_for_mode(self, include_label_and_locus):
    """Returns a dict describing features from a TF.example."""
    spec = {
        'image/encoded': tf.io.FixedLenFeature((), tf.string),
        'variant/encoded': tf.io.FixedLenFeature((), tf.string),
        'alt_allele_indices/encoded': tf.io.FixedLenFeature((), tf.string),
        'variant_type': tf.io.FixedLenFeature((), tf.int64),
        'sequencing_type': tf.io.FixedLenFeature([], tf.int64),
    }
    if include_label_and_locus:
      # N.B. int32 fails here on TPU.
      spec['label'] = tf.io.FixedLenFeature((), tf.int64)
      spec['locus'] = tf.io.FixedLenFeature((), tf.string)
    return spec

  def parse_tfexample(self, tf_example):
    """Parse a DeepVariant pileup tf.Example to features and labels.

    This potentially stores parsed strings as fixed length tensors of integers,
    as required by TPU.  They have to be handled properly by consumers.

    Args:
      tf_example: a serialized tf.Example for a DeepVariant "pileup".

    Returns:
      If (mode is EVAL or TRAIN) or debugging_true_label_mode:
        (features, label) ...
      If mode is PREDICT,
        features ...
    """
    # TODO: Add tests for parse_tf_example function
    with tf.compat.v1.name_scope('input'):
      parsed = tf.io.parse_single_example(
          serialized=tf_example, features=self.feature_extraction_spec)
      image = parsed['image/encoded']
      if self.tensor_shape:
        # If the input is empty there won't be a tensor_shape.
        image = tf.reshape(tf.io.decode_raw(image, tf.uint8), self.tensor_shape)
        if self.use_tpu:
          # Cast to int32 for loading onto the TPU
          image = tf.cast(image, tf.int32)

      variant = parsed['variant/encoded']
      alt_allele_indices = parsed['alt_allele_indices/encoded']
      if self.use_tpu:
        # Passing a string to a TPU draws this error: TypeError: <dtype:
        # 'string'> is not a supported TPU infeed type. Supported types are:
        # [tf.float32, tf.int32, tf.complex64, tf.int64, tf.bool, tf.bfloat16]
        # Thus, we must encode the string as a tensor of int.
        variant = dv_utils.string_to_int_tensor(variant)
        alt_allele_indices = dv_utils.string_to_int_tensor(alt_allele_indices)

      features = {
          'image': image,
          'variant': variant,
          'alt_allele_indices': alt_allele_indices,
          'sequencing_type': parsed['sequencing_type'],
      }

      if (self.mode in (tf_estimator.ModeKeys.TRAIN, tf_estimator.ModeKeys.EVAL)
          or self.debugging_true_label_mode):
        if self.use_tpu:
          features['locus'] = dv_utils.string_to_int_tensor(parsed['locus'])
        else:
          features['locus'] = parsed['locus']

        # Add variant_type to our features if are in TRAIN or EVAL mode.
        features['variant_type'] = parsed['variant_type']

        if self.mode in (tf_estimator.ModeKeys.TRAIN,
                         tf_estimator.ModeKeys.EVAL):
          label = parsed['label']
          return features, label
        features['label'] = parsed['label']

      # For predict model, label is not present. So, returns features only.
      return features

  def __call__(self, params):
    """Interface to get a data batch, fulfilling `input_fn` contract.

    Args:
      params: a dict containing an integer value for key 'batch_size'.

    Returns:
      the tuple (features, labels), where:
        - features is a dict of Tensor-valued input features; keys populated
          are:
            'image'
            'variant'
            'alt_allele_indices'
          and, if not PREDICT mode, also:
            'locus'

          Aside from 'image', these may be encoded specially for TPU.

        - label is the Tensor-valued prediction label; in train/eval
          mode the label value is is populated from the data source; in
          inference mode, the value is a constant empty Tensor value "()".
    """
    # See https://cloud.google.com/tpu/docs/tutorials/inception-v3-advanced
    # for some background on tuning this on TPU.

    # TPU optimized implementation for prediction mode
    if self.mode == tf_estimator.ModeKeys.PREDICT:
      return self.prediction_input_fn(params)

    # Optimized following:
    #   https://www.tensorflow.org/guide/performance/datasets
    # using the information available from xprof.
    def load_dataset(filename):
      dataset = tf.data.TFRecordDataset(
          filename,
          buffer_size=self.prefetch_dataset_buffer_size,
          compression_type=compression_type)
      return dataset

    batch_size = params['batch_size']
    compression_type = dv_utils.compression_type_of_files(self.input_files)

    # NOTE: The order of the file names returned can be non-deterministic,
    # even if shuffle is false.  See internal and the note in internal.
    # We need the shuffle flag to be able to disable reordering in EVAL mode.
    dataset = None
    for pattern in self.input_file_spec.split(','):
      one_dataset = tf.data.Dataset.list_files(
          sharded_file_utils.normalize_to_sharded_file_pattern(pattern),
          shuffle=self.mode == tf_estimator.ModeKeys.TRAIN)
      dataset = dataset.concatenate(one_dataset) if dataset else one_dataset

    # This shuffle applies to the set of files.
    # TODO: why would we shuffle the files?
    if (self.mode == tf_estimator.ModeKeys.TRAIN and
        self.initial_shuffle_buffer_size > 0):
      dataset = dataset.shuffle(self.initial_shuffle_buffer_size)

    # For both TRAIN and EVAL, use the following to speed up.
    if self.sloppy:
      options = tf.data.Options()
      options.experimental_deterministic = False
      dataset = dataset.with_options(options)
    dataset = dataset.interleave(
        load_dataset,
        cycle_length=self.input_read_threads,
        num_parallel_calls=tf.data.AUTOTUNE)

    if self.max_examples is not None:
      dataset = dataset.take(self.max_examples)

    if self.mode == tf_estimator.ModeKeys.TRAIN:
      dataset = dataset.repeat()

    # This shuffle applies to the set of records.
    if self.mode == tf_estimator.ModeKeys.TRAIN:
      if self.shuffle_buffer_size > 0:
        dataset = dataset.shuffle(self.shuffle_buffer_size)

    dataset = dataset.map(
        map_func=self.parse_tfexample, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size=batch_size, drop_remainder=True)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)

    return dataset

  def prediction_input_fn(self, params):
    """Implementation of `input_fn` contract for prediction mode.

    Args:
      params: a dict containing an integer value for key 'batch_size'.

    Returns:
      the tuple (features, labels), where:
        - features is a dict of Tensor-valued input features; keys populated
          are:
            'image'
            'variant'
            'alt_allele_indices'

          Aside from 'image', these may be encoded specially for TPU.
    """

    def load_dataset(filename):
      dataset = tf.data.TFRecordDataset(
          filename,
          buffer_size=self.prefetch_dataset_buffer_size,
          compression_type=compression_type)
      return dataset

    batch_size = params['batch_size']
    compression_type = dv_utils.compression_type_of_files(self.input_files)
    dataset = tf.data.Dataset.list_files(
        sharded_file_utils.normalize_to_sharded_file_pattern(
            self.input_file_spec),
        shuffle=False,
    )
    logging.vlog(3,
                 'self.input_read_threads={}'.format(self.input_read_threads))
    if self.sloppy:
      options = tf.data.Options()
      options.experimental_deterministic = False
      dataset = dataset.with_options(options)
    dataset = dataset.interleave(
        load_dataset,
        cycle_length=self.input_read_threads,
        num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.map(
        map_func=self.parse_tfexample, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size=batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

  def __str__(self):
    return ('DeepVariantInput(name={}, input_file_spec={}, num_examples={}, '
            'mode={})').format(self.name, self.input_file_spec,
                               self.num_examples, self.mode)


# This is the entry point to get a DeepVariantInput when you start with
# a dataset configuration file name.
def get_input_fn_from_dataset(dataset_config_filename, mode, **kwargs):
  """Creates an input_fn from the dataset config file.

  Args:
    dataset_config_filename: str. Path to the dataset config pbtxt file.
    mode: one of tf.estimator.ModeKeys.{TRAIN,EVAL,PREDICT}
    **kwargs: Additional keyword arguments for DeepVariantInput.

  Returns:
    An input_fn from the specified split in the dataset_config file.

  Raises:
    ValueError: if the dataset config doesn't have the necessary information.
  """
  # Get the metadata.
  dataset_config = read_dataset_config(dataset_config_filename)
  # Return a reader for the data.
  return get_input_fn_from_filespec(
      input_file_spec=dataset_config.tfrecord_path,
      mode=mode,
      num_examples=dataset_config.num_examples,
      name=dataset_config.name,
      **kwargs)


# This is the entry point to get a DeepVariantInput when you start with
# a tf.example file specification, and associated metadata.
def get_input_fn_from_filespec(input_file_spec, mode, **kwargs):
  """Create a DeepVariantInput function object from a file spec.

  Args:
    input_file_spec: the tf.example input file specification, possibly sharded.
    mode: tf.estimator.ModeKeys.
    **kwargs: Additional keyword arguments for DeepVariantInput.

  Returns:
    A DeepVariantInput object usable as an input_fn.
  """
  return DeepVariantInput(mode=mode, input_file_spec=input_file_spec, **kwargs)


# Return the stream of batched images from a dataset.
def get_batches(tf_dataset, model, batch_size):
  """Provides batches of pileup images from this dataset.

  Creates a DeepVariantInput for tf_dataset. It instantiates an iterator
  on the dataset, and returns the images, labels, encoded_variant
  features in batches. This calls model.preprocess_images on the images
  (but note that we will be moving that step into model_fn for the
  Estimator api).

  Args:
    tf_dataset: a DeepVariantInput object
    model: a model object
    batch_size: int batch size

  Returns:
    (images, labels, encoded_variant)

  Raises:
    ValueError: if the dataset has the wrong mode.
  """
  if tf_dataset.mode not in (tf_estimator.ModeKeys.TRAIN,
                             tf_estimator.ModeKeys.EVAL):
    raise ValueError(
        'tf_dataset.mode is {} but must be one of TRAIN or EVAL.'.format(
            tf_dataset.mode))

  params = dict(batch_size=batch_size)
  features, labels = tf.compat.v1.data.make_one_shot_iterator(
      tf_dataset(params)).get_next()

  images = features['image']
  encoded_variant = features['variant']

  images = model.preprocess_images(images)
  return images, labels, encoded_variant


# This reads a pbtxt file and returns the config proto.
def read_dataset_config(dataset_config_filename):
  """Returns a DeepVariantDatasetConfig proto read from the dataset config file.

  Args:
    dataset_config_filename: String. Path to the dataset config pbtxt file.

  Returns:
    A DeepVariantDatasetConfig proto from the dataset_config file.

  Raises:
    ValueError: if the dataset config doesn't have the necessary information.
  """
  with tf.io.gfile.GFile(dataset_config_filename) as f:
    dataset_config = text_format.Parse(
        f.read(), deepvariant_pb2.DeepVariantDatasetConfig())

  if not dataset_config.name:
    raise ValueError('dataset_config needs to have a name')

  if not dataset_config.tfrecord_path:
    raise ValueError('The dataset in the config {} does not have a '
                     'tfrecord_path.'.format(dataset_config_filename))

  # TODO: remove this check once we're able to deal with absence
  # of num_examples.
  if not dataset_config.num_examples:
    raise ValueError('The dataset in the config {} does not have a '
                     'num_examples.'.format(dataset_config_filename))

  return dataset_config


def write_dataset_config_to_pbtxt(dataset_config, dataset_config_filename):
  """Writes the dataset_config to a human-readable text format.

  Args:
    dataset_config: DeepVariantDatasetConfig. The config to be written out.
    dataset_config_filename: String. Path to the output pbtxt file.
  """
  with tf.io.gfile.GFile(dataset_config_filename, mode='w') as writer:
    writer.write(text_format.MessageToString(dataset_config))
