# 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.

"""Utility functions for DeepVariant.

Started with a collection of utilities for working with the TF models. Now this

file includes broader utilities we use in DeepVariant.

"""

import json

from typing import Optional, Tuple

from absl import logging

import numpy as np

import tensorflow as tf

from deepvariant.protos import deepvariant_pb2

from third_party.nucleus.io import sharded_file_utils

# TODO: this dep still uses CLIF.

from third_party.nucleus.io import tfrecord

from third_party.nucleus.protos import variants_pb2

from tensorflow.core.example import example_pb2

# Convert strings up to this length, then clip.  We picked a number that

# was less than 1K, with a bit of extra space for the length element,

# to give enough space without overflowing to a larger multiple of 128.

STRING_TO_INT_MAX_CONTENTS_LEN = 1020

# This is the length of the resulting buffer (including the length entry).

STRING_TO_INT_BUFFER_LENGTH = STRING_TO_INT_MAX_CONTENTS_LEN + 1

def example_variant_type(example):

  """Gets the locus field from example as a string."""

  return example.features.feature['variant_type'].int64_list.value[0]

def example_locus(example):

  """Gets the locus field from example as a string."""

  return example.features.feature['locus'].bytes_list.value[0]

def example_alt_alleles_indices(example):

  """Gets an iterable of the alt allele indices in example."""

  return deepvariant_pb2.CallVariantsOutput.AltAlleleIndices.FromString(

      example.features.feature['alt_allele_indices/encoded'].bytes_list.value[0]

  ).indices

def example_alt_alleles(example, variant=None):

  """Gets a list of the alt alleles in example."""

  variant = variant if variant else example_variant(example)

  return [

      variant.alternate_bases[i] for i in example_alt_alleles_indices(example)

  ]

def example_encoded_image(example):

  """Gets image field from example as a string."""

  return example.features.feature['image/encoded'].bytes_list.value[0]

def example_variant(example):

  """Gets and decodes the variant field from example as a Variant."""

  encoded = example.features.feature['variant/encoded'].bytes_list.value[0]

  return variants_pb2.Variant.FromString(encoded)

def example_label(example: example_pb2.Example) -> Optional[int]:

  """Gets the label field from example as a string."""

  if 'label' not in example.features.feature:

    return None

  return int(example.features.feature['label'].int64_list.value[0])

def example_denovo_label(example: example_pb2.Example) -> Optional[int]:

  """Gets the label field from example as a string.

  Args:

    example: A tf.Example containing DeepVariant example.

  Returns:

    De novo label for the example.

  """

  if 'denovo_label' not in example.features.feature:

    return None

  return int(example.features.feature['denovo_label'].int64_list.value[0])

def example_image_shape(example):

  """Gets the image shape field from example as a list of int64."""

  if len(example.features.feature['image/shape'].int64_list.value) != 3:

    raise ValueError(

        'Invalid image/shape: we expect to find an image/shape '

        'field with length 3.'

    )

  return example.features.feature['image/shape'].int64_list.value[0:3]

def example_key(example):

  """Constructs a key for example based on its position and alleles."""

  variant = example_variant(example)

  alts = example_alt_alleles(example)

  return '{}:{}:{}->{}'.format(

      variant.reference_name,

      variant.start + 1,

      variant.reference_bases,

      '/'.join(alts),

  )

def example_set_label(example, numeric_label):

  """Sets the label features of example.

  Sets the label feature of example to numeric_label.

  Args:

    example: A tf.Example proto.

    numeric_label: A numeric (int64 compatible) label for example.

  """

  example.features.feature['label'].int64_list.value[:] = [numeric_label]

def example_set_denovo_label(

    example: example_pb2.Example, numeric_label: int

) -> None:

  """Sets the denovo label features of example.

  Sets the label feature of example to numeric_label.

  Args:

    example: a tf.Example proto.

    numeric_label: A numeric (int64 compatible) label for example.

  """

  example.features.feature['denovo_label'].int64_list.value[:] = [numeric_label]

def example_set_variant(example, variant, deterministic=False):

  """Sets the variant/encoded feature of example to variant.SerializeToString().

  Args:

    example: a tf.Example proto.

    variant: third_party.nucleus.protos.Variant protobuf containing information

      about a candidate variant call.

    deterministic: Used to set SerializeToString.

  """

  example.features.feature['variant/encoded'].bytes_list.value[:] = [

      variant.SerializeToString(deterministic=deterministic)

  ]

def example_sequencing_type(example):

  return example.features.feature['sequencing_type'].int64_list.value[0]

def get_one_example_from_examples_path(source, proto=None):

  """Get the first record from `source`.

  Args:

    source: str. A pattern or a comma-separated list of patterns that represent

      file names.

    proto: A proto class. proto.FromString() will be called on each serialized

      record in path to parse it.

  Returns:

    The first record, or None.

  """

  files = sharded_file_utils.glob_list_sharded_file_patterns(source)

  if not files:

    raise ValueError(

        'Cannot find matching files with the pattern "{}"'.format(source)

    )

  for f in files:

    try:

      return next(

          tfrecord.read_tfrecords(f, proto=proto, compression_type='GZIP')

      )

    except StopIteration:

      # Getting a StopIteration from one next() means source_path is empty.

      # Move on to the next one to try to get one example.

      pass

  return None

def get_shape_from_examples_path(source):

  """Reads one record from source to determine the tensor shape for all."""

  one_example = get_one_example_from_examples_path(source)

  if one_example:

    return example_image_shape(one_example)

  return None

def _simplify_variant(variant):

  """Returns a new Variant with only the basic fields of variant."""

  def _simplify_variant_call(call):

    """Returns a new VariantCall with the basic fields of call."""

    return variants_pb2.VariantCall(

        call_set_name=call.call_set_name,

        genotype=call.genotype,

        info=dict(call.info),

    )  # dict() is necessary to actually set info.

  return variants_pb2.Variant(

      reference_name=variant.reference_name,

      start=variant.start,

      end=variant.end,

      reference_bases=variant.reference_bases,

      alternate_bases=variant.alternate_bases,

      filter=variant.filter,

      quality=variant.quality,

      calls=[_simplify_variant_call(call) for call in variant.calls],

  )

def string_to_int_tensor(x):

  """Graph operations decode a string into a fixed-size tensor of ints."""

  decoded = tf.compat.v1.decode_raw(x, tf.uint8)

  clipped = decoded[:STRING_TO_INT_MAX_CONTENTS_LEN]  # clip to allowed max_len

  shape = tf.shape(input=clipped)

  slen = shape[0]

  # pad to desired max_len

  padded = tf.pad(

      tensor=clipped, paddings=[[0, STRING_TO_INT_MAX_CONTENTS_LEN - slen]]

  )

  casted = tf.cast(padded, tf.int32)

  casted.set_shape([STRING_TO_INT_MAX_CONTENTS_LEN])

  return tf.concat([[slen], casted], 0)

def int_tensor_to_string(x):

  """Python operations to encode a tensor of ints into string of bytes."""

  slen = x[0]

  v = x[1 : slen + 1]

  return np.array(v, dtype=np.uint8).tostring()

def tpu_available(sess=None):

  """Return true if a TPU device is available to the default session."""

  if sess is None:

    init_op = tf.group(

        tf.compat.v1.global_variables_initializer(),

        tf.compat.v1.local_variables_initializer(),

    )

    with tf.compat.v1.Session() as sess:

      sess.run(init_op)

      devices = sess.list_devices()

  else:

    devices = sess.list_devices()

  return any(dev.device_type == 'TPU' for dev in devices)

def resolve_master(master, tpu_name, tpu_zone, gcp_project):

  """Resolve the master's URL given standard flags."""

  if master is not None:

    return master

  elif tpu_name is not None:

    return tf.distribute.cluster_resolver.TPUClusterResolver(

        tpu=[tpu_name], zone=tpu_zone, project=gcp_project

    ).get_master()

  else:

    # For k8s TPU we do not have/need tpu_name. See

    # https://cloud.google.com/tpu/docs/kubernetes-engine-setup#tensorflow-code

    return tf.distribute.cluster_resolver.TPUClusterResolver().get_master()

def get_example_info_json_filename(

    examples_filename: str, task_id: Optional[int]

) -> str:

  """Returns corresponding example_info.json filename for examples_filename."""

  if sharded_file_utils.is_sharded_file_spec(examples_filename):

    assert task_id is not None

    # If examples_filename has the @shards representation, resolve it into

    # the first shard. We only write .example_info.json to the first shard.

    example_info_prefix = sharded_file_utils.sharded_filename(

        examples_filename, task_id

    )

  else:

    # In all other cases, including non-sharded files,

    # or sharded filenames with -ddddd-of-ddddd, just append.

    example_info_prefix = examples_filename

  return example_info_prefix + '.example_info.json'

def get_shape_and_channels_from_json(example_info_json):

  """Returns the shape and channels list from the input json."""

  if not tf.io.gfile.exists(example_info_json):

    logging.warning(

        (

            'Starting from v1.4.0, we expect %s to '

            'include information for shape and channels.'

        ),

        example_info_json,

    )

    return None, None

  with tf.io.gfile.GFile(example_info_json) as f:

    example_info = json.load(f)

  example_shape = example_info['shape']

  example_channels_enum = example_info['channels']

  logging.info(

      'From %s: Shape of input examples: %s, Channels of input examples: %s.',

      example_info_json,

      str(example_shape),

      str(example_channels_enum),

  )

  return example_shape, example_channels_enum

def get_tf_record_writer(output_filename: str) -> tf.io.TFRecordWriter:

  tf_options = tf.io.TFRecordOptions(compression_type='GZIP')

  return tf.io.TFRecordWriter(output_filename, options=tf_options)

def preprocess_images(images):

  """Applies preprocessing operations for Inception images.

  Because this will run in model_fn, on the accelerator, we use operations

  that efficiently execute there.

  Args:

    images: A Tensor of with uint8 values.

  Returns:

    A tensor of images the same shape, containing floating point values, with

    all points rescaled between -1 and 1 and possibly resized.

  """

  images = tf.cast(images, dtype=tf.float32)

  images = tf.subtract(images, 128.0)

  images = tf.math.divide(images, 128.0)

  return images

def unpreprocess_images(images: np.ndarray) -> np.ndarray:

  """Reverses preprocess_images in numpy format.

  Args:

    images: A numpy array with floating point values.

  Returns:

    A numpy array of images the same shape.

  """

  images *= 128.0

  images += 128.0

  # We can optionally convert it to uint8 by .astype(np.uint8).

  # But for now we'll just return it as floating points.

  return images

def call_variant_to_tfexample(

    cvo: deepvariant_pb2.CallVariantsOutput,

    image_shape: Tuple[int, int, int] = (100, 221, 7),

) -> tf.train.Example:

  """Converts CallVariantsOutput to tf.train.Example if possible.

  This function is for specific debugging purpose, so it will only

  work on CallVariantsOutput with debug_info.image_encoded.

  Note: Not all values are transferred as there isn't a 1:1 mapping. No mapping

  exists for 'variant/encoded' or 'sequencing_type' for example.

  Args:

    cvo: A CallVariantsOutput to convert to a TF.Example.

    image_shape: The shape of the image contained within cvo.

  Returns:

    A Tf.Example created from the given CallVariantsOutput.

  Raises:

    ValueError if the input data lacks the needed fields.

  """

  tfexample = tf.train.Example()

  features = tfexample.features.feature

  features['image/shape'].int64_list.value[:] = list(image_shape)

  if cvo.debug_info and cvo.debug_info.image_encoded:

    features['image/encoded'].bytes_list.value[:] = [

        cvo.debug_info.image_encoded

    ]

  else:

    raise ValueError('CallVariantsOutput does not contain an image.')

  features['label'].int64_list.value[:] = [cvo.debug_info.true_label]

  if cvo.alt_allele_indices:

    features['alt_allele_indices'].int64_list.value[

        :

    ] = cvo.alt_allele_indices.indices

  # Create and assign locus

  features['locus'].bytes_list.value[:] = [

      bytes(

          f'{cvo.variant.reference_name}:{cvo.variant.start}-{cvo.variant.end}',

          'utf-8',

      )

  ]

  return tfexample