# Copyright 2018 Google LLC.
#
# 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.
"""VariantCall utilities."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from nucleus.protos import struct_pb2
from nucleus.protos import variants_pb2
from nucleus.util import struct_utils
from nucleus.util import vcf_constants

# Special-cased FORMAT fields that are first-class fields in the VariantCall.
_GL = 'GL'
_GT = 'GT'


def set_format(variant_call, field_name, value, vcf_object=None):
  """Sets a field of the info map of the `VariantCall` to the given value(s).

  `variant_call.info` is analogous to the FORMAT field of a VCF call.

  Example usage:
  with vcf.VcfReader('/path/to/my.vcf') as vcf_reader:
    for variant in vcf_reader:
      first_call = variant.calls[0]
      # Type can be inferred for reserved VCF fields.
      set_format(first_call, 'AD', 25)
      # Specify the reader explicitly for unknown fields.
      set_format(first_call, 'MYFIELD', 30, vcf_reader)

  Args:
    variant_call: VariantCall proto. The VariantCall to modify.
    field_name: str. The name of the field to set.
    value: A single value or list of values to update the VariantCall with.
      The type of the value is determined by the `vcf_object` if one is given,
      otherwise is looked up based on the reserved FORMAT fields in the VCF
      specification.
    vcf_object: (Optional) A VcfReader or VcfWriter object. If not None, the
      type of the field is inferred from the associated VcfReader or VcfWriter
      based on its name. Otherwise, the type is inferred if it is a reserved
      field.
  """
  if field_name == _GL:
    set_gl(variant_call, value)
    return
  if field_name == _GT:
    set_gt(variant_call, value)
    return

  if vcf_object is None:
    set_field_fn = vcf_constants.reserved_format_field_set_fn(field_name)
  else:
    set_field_fn = vcf_object.field_access_cache.format_field_set_fn(field_name)
  set_field_fn(variant_call.info, field_name, value)


def get_format(variant_call, field_name, vcf_object=None):
  """Returns the value of the `field_name` FORMAT field.

  The `vcf_object` is used to determine the type of the resulting value. If it
  is a single value or a Flag, that single value will be returned. Otherwise,
  the list of values is returned.

  Args:
    variant_call: VariantCall proto. The VariantCall of interest.
    field_name: str. The name of the field to retrieve values from.
    vcf_object: (Optional) A VcfReader or VcfWriter object. If not None, the
      type of the field is inferred from the associated VcfReader or VcfWriter
      based on its name. Otherwise, the type is inferred if it is a reserved
      field.
  """
  if field_name == _GL:
    return get_gl(variant_call)
  if field_name == _GT:
    return get_gt(variant_call)

  if vcf_object is None:
    get_field_fn = vcf_constants.reserved_format_field_get_fn(field_name)
  else:
    get_field_fn = vcf_object.field_access_cache.format_field_get_fn(field_name)
  return get_field_fn(variant_call.info, field_name)


# The following functions are convenience methods for getting/setting some
# reserved FORMAT fields of a VariantCall as well as some non-reserved FORMAT
# fields used by DeepVariant. Note that these functions will use the types of
# each field as defined by the VCF 4.3 specification, mirrored in
# vcf_constants.py, so if you have redefined any of these fields to have
# different types these functions will not do what you want.
def set_ad(variant_call, ad):
  """Sets the allele depth of the VariantCall."""
  set_format(variant_call, 'AD', ad)


def get_ad(variant_call):
  """Gets the allele depth of the VariantCall."""
  return get_format(variant_call, 'AD')


def set_gl(variant_call, gl):
  """Sets the genotype likelihoods of the VariantCall.

  Args:
    variant_call: VariantCall proto. The VariantCall to modify.
    gl: list(float). The list of genotype likelihoods for the VariantCall.
  """
  # Note: genotype_likelihood is extracted to a first-class field within
  # VariantCall. Consequently, we just set its value directly here.
  variant_call.genotype_likelihood[:] = gl


def get_gl(variant_call):
  """Returns the genotype likelihoods of the VariantCall.

  Args:
    variant_call: VariantCall proto. The VariantCall for which to return GLs.

  Returns:
    A list of floats representing the genotype likelihoods of this call.
  """
  return variant_call.genotype_likelihood


def set_gt(variant_call, gt):
  """Sets the genotypes of the VariantCall.

  Args:
    variant_call: VariantCall proto. The VariantCall to modify.
    gt: list(int). The list of genotypes for the VariantCall.
  """
  # Note: genotype is extracted to a first-class field within
  # VariantCall. Consequently, we just set its value directly here.
  variant_call.genotype[:] = gt


def get_gt(variant_call):
  """Returns the genotypes of the VariantCall.

  Args:
    variant_call: VariantCall proto. The VariantCall for which to return GTs.

  Returns:
    A list of ints representing the genotype indices of this call.
  """
  return variant_call.genotype


def set_gq(variant_call, gq):
  """Sets the genotype quality of the VariantCall."""
  set_format(variant_call, 'GQ', gq)


def get_gq(variant_call):
  """Gets the genotype quality of the VariantCall."""
  return get_format(variant_call, 'GQ')


def set_min_dp(variant_call, min_dp):
  """Sets the 'MIN_DP' field of the VariantCall."""
  struct_utils.set_int_field(variant_call.info, 'MIN_DP', min_dp)


def get_min_dp(variant_call):
  """Gets the 'MIN_DP' field of the VariantCall."""
  return struct_utils.get_int_field(
      variant_call.info, 'MIN_DP', is_single_field=True)


def has_genotypes(variant_call):
  """Returns True iff the VariantCall has one or more called genotypes.

  Args:
    variant_call: VariantCall proto. The VariantCall to evaluate.

  Returns:
    True if the VariantCall has one or more called genotypes, False otherwise.
  """
  return any(gt >= 0 for gt in variant_call.genotype)


def has_full_genotypes(variant_call):
  """Returns True iff the VariantCall has only known genotypes.

  Args:
    variant_call: VariantCall proto. The VariantCall to evaluate.

  Returns:
    True if all `genotype` fields are known genotypes.
  """
  return all(gt >= 0 for gt in variant_call.genotype)


def ploidy(variant_call):
  """Returns the ploidy of the VariantCall.

  Args:
    variant_call: VariantCall proto. The VariantCall to evaluate.

  Returns:
    The ploidy of the call (a non-negative integer).
  """
  # Unknown genotypes are represented as -1 in VariantCall protos. When
  # a VCF is parsed that contains multiple ploidies in different samples,
  # a separate padding value of -2**30 - 1 is inserted into the calls.
  return sum(gt >= -1 for gt in variant_call.genotype)


def has_variation(variant_call):
  """Returns True if and only if the call has a non-reference genotype.

  Args:
    variant_call: VariantCall proto. The VariantCall to evaluate.

  Returns:
    True if and only if the call has a non-reference genotype.
  """
  return any(gt > 0 for gt in variant_call.genotype)


def is_heterozygous(variant_call):
  """Returns True if and only if the call is heterozygous.

  Args:
    variant_call: VariantCall proto. The VariantCall to evaluate.

  Returns:
    True if and only if the call is heterozygous.
  """
  return len({gt for gt in variant_call.genotype if gt >= 0}) >= 2
