# 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.
"""Classes for reading and writing VCF files.

The VCF format is described at
https://samtools.github.io/hts-specs/VCFv4.3.pdf

API for reading:

```python
from nucleus.io import vcf

with vcf.VcfReader(input_path) as reader:
  for variant in reader:
    print(variant)
```

API for writing:

```python
from nucleus.io import vcf

# variants is an iterable of nucleus.genomics.v1.Variant protocol buffers.
variants = ...

with vcf.VcfWriter(output_path, header=header) as writer:
  for variant in variants:
    writer.write(variant)
```

The class attempts to infer the file format (`TFRecord` vs VCF) from the file
path provided to the constructor.

1. For files that end with '.tfrecord' and '.tfrecord.gz' (a gzipped version),
  a `TFRecord` file is assumed and is attempted to be read or written.

2. For all other files, the VCF format will be used.

  VCF format used in writing is inferred from file paths:
    - ending in '.bcf.gz': BGZF compressed BCF format will be written;
    - ending in '.bcf': uncompressed BCF format will be written;
    - ending in '.gz' and not in '.bcf.gz': BGZP compressed VCF format will be
        written;
    - all other suffixes: uncompressed VCF format will be written.

  VCF format used in reading is inferred from the contents of the file.
"""

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

from nucleus.io import genomics_reader
from nucleus.io import genomics_writer
from nucleus.io.python import vcf_reader
from nucleus.io.python import vcf_writer
from nucleus.protos import variants_pb2
from nucleus.util import ranges
from nucleus.util import variant_utils
from nucleus.util import vcf_constants


def _create_get_fn_cache(fields):
  """Returns a dictionary from field to a callable that extracts its value."""
  return {
      field.id: vcf_constants.create_get_fn(field.type, field.number)
      for field in fields
  }


def _create_set_fn_cache(fields):
  """Returns a dictionary from field to a callable that sets its value."""
  return {field.id: vcf_constants.SET_FN_LOOKUP[field.type] for field in fields}


class VcfHeaderCache(object):
  """This class creates a cache of accessors to structured fields in Variants.

  The INFO and FORMAT fields within Variant protos are structured and typed,
  with types defined by the corresponding VCF header. This cache object provides
  provides {info,format}_field_{get,set}_fn functions that can be used to
  extract information from the structured Variant protos based on the types
  defined therein.

  NOTE: Users should not need to interact with this class at all. It is used
  by the variant_utils.{get,set}_info and variantcall_utils.{get,set}_format
  functions for interacting with the INFO and FORMAT fields in a Variant proto.
  """

  def __init__(self, header):
    """Initializer.

    Args:
      header: nucleus.genomics.v1.VcfHeader proto. Used to define the accessor
        functions needed.
    """
    if header is None:
      header = variants_pb2.VcfHeader()
    self._info_get_cache = _create_get_fn_cache(header.infos)
    self._info_set_cache = _create_set_fn_cache(header.infos)
    self._format_get_cache = _create_get_fn_cache(header.formats)
    self._format_set_cache = _create_set_fn_cache(header.formats)

  def info_field_get_fn(self, field_name):
    """Returns a callable that extracts the given INFO field based on its type.

    Args:
      field_name: str. The INFO field name of interest, e.g. 'AA', 'DB', 'AF'.

    Returns:
      A callable used to extract the given INFO field from a Variant proto.
    """
    return self._info_get_cache[field_name]

  def info_field_set_fn(self, field_name):
    """Returns a callable that sets the given INFO field based on its type."""
    return self._info_set_cache[field_name]

  def format_field_get_fn(self, field_name):
    """Returns a callable that gets the given FORMAT field based on its type."""
    return self._format_get_cache[field_name]

  def format_field_set_fn(self, field_name):
    """Returns a callable that sets the given FORMAT field based on its type."""
    return self._format_set_cache[field_name]


class NativeVcfReader(genomics_reader.GenomicsReader):
  """Class for reading from native VCF files.

  Most users will want to use VcfReader instead, because it dynamically
  dispatches between reading native VCF files and TFRecord files based
  on the filename's extensions.
  """

  def __init__(self,
               input_path,
               excluded_info_fields=None,
               excluded_format_fields=None,
               store_gl_and_pl_in_info_map=False,
               header=None):
    """Initializer for NativeVcfReader.

    Args:
      input_path: str. The path to the VCF file to read.
      excluded_info_fields: list(str). A list of INFO field IDs that should not
        be parsed into the Variants. If None, all INFO fields are included.
      excluded_format_fields: list(str). A list of FORMAT field IDs that should
        not be parsed into the Variants. If None, all FORMAT fields are
        included.
      store_gl_and_pl_in_info_map: bool. If True, the "GL" and "PL" FORMAT
        fields are stored in the VariantCall.info map rather than as top-level
        values in the VariantCall.genotype_likelihood field.
      header: If not None, specifies the variants_pb2.VcfHeader. The file at
        input_path must not contain any header information.
    """
    super(NativeVcfReader, self).__init__()

    options = variants_pb2.VcfReaderOptions(
        excluded_info_fields=excluded_info_fields,
        excluded_format_fields=excluded_format_fields,
        store_gl_and_pl_in_info_map=store_gl_and_pl_in_info_map)
    if header is not None:
      self._reader = vcf_reader.VcfReader.from_file_with_header(
          input_path.encode('utf8'), options, header)
    else:
      self._reader = vcf_reader.VcfReader.from_file(
          input_path.encode('utf8'), options)

    self.header = self._reader.header
    self.field_access_cache = VcfHeaderCache(self.header)

  def iterate(self):
    """Returns an iterable of Variant protos in the file."""
    return self._reader.iterate()

  def query(self, region):
    """Returns an iterator for going through variants in the region."""
    return self._reader.query(region)

  def __exit__(self, exit_type, exit_value, exit_traceback):
    self._reader.__exit__(exit_type, exit_value, exit_traceback)

  @property
  def c_reader(self):
    """Returns the underlying C++ reader."""
    return self._reader


class VcfReader(genomics_reader.DispatchingGenomicsReader):
  """Class for reading Variant protos from VCF or TFRecord files."""

  def _native_reader(self, input_path, **kwargs):
    return NativeVcfReader(input_path, **kwargs)

  def _record_proto(self):
    return variants_pb2.Variant

  def _post_init_hook(self):
    # Initialize field_access_cache.  If we are dispatching to a
    # NativeVcfReader, we use its field_access_cache. Otherwise, we need to
    # create a new one.
    self.field_access_cache = getattr(
        self._reader, 'field_access_cache', VcfHeaderCache(self.header))

  @property
  def c_reader(self):
    """Returns the underlying C++ reader.

    Note that the C++ reader might be a VcfReader or it might be a
    TFRecordReader, depending on the input_path's extension.
    """
    return self._reader.c_reader


class NativeVcfWriter(genomics_writer.GenomicsWriter):
  """Class for writing to native VCF files.

  Most users will want VcfWriter, which will write to either native VCF
  files or TFRecords files, based on the output filename's extensions.
  """

  def __init__(self,
               output_path,
               header=None,
               round_qualities=False,
               excluded_info_fields=None,
               excluded_format_fields=None,
               retrieve_gl_and_pl_from_info_map=False,
               exclude_header=False):
    """Initializer for NativeVcfWriter.

    Args:
      output_path: str. The path to which to write the VCF file.
      header: nucleus.genomics.v1.VcfHeader. The header that defines all
        information germane to the constituent variants. This includes contigs,
        FILTER fields, INFO fields, FORMAT fields, samples, and all other
        structured and unstructured header lines.
      round_qualities: bool. If True, the QUAL field is rounded to one point
        past the decimal.
      excluded_info_fields: list(str). A list of INFO field IDs that should not
        be written to the output. If None, all INFO fields are included.
      excluded_format_fields: list(str). A list of FORMAT field IDs that should
        not be written to the output. If None, all FORMAT fields are included.
      retrieve_gl_and_pl_from_info_map: bool. If True, the "GL" and "PL" FORMAT
        fields are retrieved from the VariantCall.info map rather than from the
        top-level value in the VariantCall.genotype_likelihood field.
      exclude_header: bool. If True, write a headerless VCF.
    """
    super(NativeVcfWriter, self).__init__()

    if header is None:
      header = variants_pb2.VcfHeader()
    writer_options = variants_pb2.VcfWriterOptions(
        round_qual_values=round_qualities,
        excluded_info_fields=excluded_info_fields,
        excluded_format_fields=excluded_format_fields,
        retrieve_gl_and_pl_from_info_map=retrieve_gl_and_pl_from_info_map,
        exclude_header=exclude_header,
    )
    self._writer = vcf_writer.VcfWriter.to_file(output_path, header,
                                                writer_options)
    self.field_access_cache = VcfHeaderCache(header)

  def write(self, proto):
    self._writer.write(proto)

  def __exit__(self, exit_type, exit_value, exit_traceback):
    self._writer.__exit__(exit_type, exit_value, exit_traceback)


class VcfWriter(genomics_writer.DispatchingGenomicsWriter):
  """Class for writing Variant protos to VCF or TFRecord files."""

  def _native_writer(self,
                     output_path,
                     header,
                     round_qualities=False,
                     excluded_info_fields=None,
                     excluded_format_fields=None,
                     retrieve_gl_and_pl_from_info_map=False,
                     exclude_header=False):
    return NativeVcfWriter(
        output_path,
        header=header,
        round_qualities=round_qualities,
        excluded_info_fields=excluded_info_fields,
        excluded_format_fields=excluded_format_fields,
        retrieve_gl_and_pl_from_info_map=retrieve_gl_and_pl_from_info_map,
        exclude_header=exclude_header)

  def _post_init_hook(self):
    # Initialize field_access_cache.  If we are dispatching to a
    # NativeVcfWriter, we use its field_access_cache.  Otherwise, we
    # need to create a new one.
    self.field_access_cache = getattr(
        self._writer, 'field_access_cache', VcfHeaderCache(self.header))


class InMemoryVcfReader(genomics_reader.GenomicsReader):
  """Class for "reading" Variant protos from an in-memory cache of variants.

  ```python
  from nucleus.io import vcf
  from nucleus.protos import variants_pb2

  variants = [... Variant protos ...]
  header = variants_pb2.VcfHeader()

  with vcf.InMemoryVcfReader(variants, header) as reader:
    for variant in reader:
      print(variant)
  ```

  This class accepts a collection of variants and optionally a header and
  provides all of the standard API functions of VcfReader but instead of
  fetching variants from a file the variants are queried from an in-memory cache
  of variant protos.

  Note that the input variants provided to this class aren't checked in any way,
  and their ordering determines the order of variants emitted by this class for
  the iterate() and query() operations. This is intentional, to make this class
  easy to use for testing where you often want to use less-than-perfectly formed
  inputs. In order to fully meet the contract of a standard VcfReader, variants
  should be sorted by their contig ordering and then by their start and finally
  by their ends.

  Implementation note:
    The current implementation will be very slow for query() if the provided
    cache of variants is large, as we do a O(n) search to collect all of the
    overlapping variants for each query. There are several straightforward
    optimizations to do if we need/want to scale this up. (a) sort the variants
    and use a binary search to find overlapping variants (b) partition the
    variants by contig, so we have dict[contig] => [variants on contig], which
    allows us to completely avoid considering any variants on any other contigs.
    Neither of these optimizations are worth it if len(variants) is small, but
    it may be worth considering if we want to use this functionality with a
    large number of variants.
  """

  def __init__(self, variants, header=None):
    """Creates a VCFReader backed by a collection of variants.

    Args:
      variants: list of nucleus.genomics.v1.Variant protos we will "read"
        from.
      header: a VCFHeader object to provide as a result to calls to self.header,
        or None, indicating that we don't have a header associated with this
        reader.
    """
    super(InMemoryVcfReader, self).__init__()
    self.variants = list(variants)
    self.header = header

  def iterate(self):
    return iter(self.variants)

  def query(self, region):
    return iter(
        variant for variant in self.variants
        if ranges.ranges_overlap(variant_utils.variant_range(variant), region)
    )
