# 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.
"""Determine genomic ranges to perform local assembly."""

from third_party.nucleus.protos import reads_pb2
from third_party.nucleus.util import ranges
from deepvariant.protos import deepvariant_pb2
from deepvariant.protos import realigner_pb2
from deepvariant.python import allelecounter
from deepvariant.realigner.python import window_selector as cpp_window_selector


def _candidates_from_reads(config, ref_reader, reads, region):
  """Returns a list of candidate positions.

  Args:
    config: learning.genomics.deepvariant.realigner.WindowSelectorOptions
      options determining the behavior of this window selector.
    ref_reader: GenomeReference. Indexed reference genome to query bases.
    reads: list[nucleus.protos.Read]. The reads we are processing into candidate
      positions.
    region: nucleus.protos.Range. The region we are processing.

  Returns:
    A list. The elements are reference positions within region.

  Raises:
    ValueError: if config.window_selector_model.model_type isn't a valid enum
    name in realigner_pb2.WindowSelectorModel.ModelType.
  """
  allele_counter_options = deepvariant_pb2.AlleleCounterOptions(
      read_requirements=reads_pb2.ReadRequirements(
          min_mapping_quality=config.min_mapq,
          min_base_quality=config.min_base_quality),
      keep_legacy_behavior=config.keep_legacy_behavior)
  expanded_region = ranges.expand(
      region,
      config.region_expansion_in_bp,
      contig_map=ranges.contigs_dict(ref_reader.header.contigs))

  allele_counter = allelecounter.AlleleCounter(ref_reader.c_reader,
                                               expanded_region, [],
                                               allele_counter_options)

  for read in reads:
    allele_counter.add(read, 'placeholder_sample_id')

  model_type = config.window_selector_model.model_type
  if model_type == realigner_pb2.WindowSelectorModel.VARIANT_READS:
    return _variant_reads_threshold_selector(
        allele_counter, config.window_selector_model.variant_reads_model,
        expanded_region)
  elif model_type == realigner_pb2.WindowSelectorModel.ALLELE_COUNT_LINEAR:
    return _allele_count_linear_selector(
        allele_counter, config.window_selector_model.allele_count_linear_model,
        expanded_region)
  else:
    raise ValueError('Unknown enum option "{}" for '
                     'WindowSelectorModel.model_type'.format(
                         config.window_selector_model.model_type))


def _variant_reads_threshold_selector(allele_counter, model_conf,
                                      expanded_region):
  """Returns a list of candidate positions.

  Following cigar operations generate candidate position:
    - ALIGNMENT_MATCH, SEQUENCE_MISMATCH, SEQUENCE_MATCH: at mismatch positions
      in the read when compared to the reference sequence.
    - DELETE: at positions within [cigar_start, cigar_start + cigar_len)
    - INSERT, CLIP_SOFT: at positions within
        [cigar_start - cigar_len, cigar_start + cigar_len)

   Note. Function implementation has changed to return positions beyond input
   region in case we have variants there. See the change at internal and
   internal.

  Args:
    allele_counter: learning.genomics.deepvariant.realigner.AlleleCounter in the
      considered region.
    model_conf: learning.genomics.deepvariant.realigner
      .WindowSelectorOptions.VariantReadsThresholdModel options determining the
      behavior of this window selector.
    expanded_region: nucleus.protos.Range. The region we are processing.

  Returns:
    A list. The elements are reference positions within region.
  """

  counts_vec = cpp_window_selector.variant_reads_candidates_from_allele_counter(
      allele_counter)

  return [
      expanded_region.start + i
      for i, count in enumerate(counts_vec)
      if (count >= model_conf.min_num_supporting_reads and
          count <= model_conf.max_num_supporting_reads)
  ]


def _allele_count_linear_selector(allele_counter, model_conf, expanded_region):
  """Returns a list of candidate positions.

  Candidate positions for realignment are generated by scoring each location.
  The score at a location is a weighted sum of the number of reads with each
  CIGAR operation at the location, where the weights are determined by the model
  coefficients. Locations whose score exceed the model decision boundary value
  are used to create realignment windows.

  Note. Function implementation has changed to return positions beyond input
  region in case we have variants there. See the change at internal and
  internal.

  Args:
    allele_counter: learning.genomics.deepvariant.realigner.AlleleCounter in the
      considered region.
    model_conf: learning.genomics.deepvariant.realigner
      .WindowSelectorOptions.AlleleCountLinearModel options determining the
      behavior of this window selector.
    expanded_region: nucleus.protos.Range. The region we are processing.

  Returns:
    A list. The elements are reference positions within region.
  """

  scores_vec = (
      cpp_window_selector.allele_count_linear_candidates_from_allele_counter(
          allele_counter, model_conf))

  return [
      expanded_region.start + i
      for i, score in enumerate(scores_vec)
      if score > model_conf.decision_boundary
  ]


def _candidates_to_windows(config, candidate_pos, ref_name):
  """"Process candidate positions to determine windows for local assembly.

  Windows are within range of
    [min(pos) - config.min_windows_distance,
     max(pos) + config.min_windows_distance)

  Args:
    config: learning.genomics.deepvariant.realigner.WindowSelectorOptions
      options determining the behavior of this window selector.
    candidate_pos: A list of ref_pos.
    ref_name: Reference name, used in setting the output
      genomics.range.reference_name value.

  Returns:
    A sorted list of nucleus.protos.Range protos for all windows in this region.
  """
  windows = []

  def _add_window(start_pos, end_pos):
    windows.append(
        ranges.make_range(ref_name, start_pos - config.min_windows_distance,
                          end_pos + config.min_windows_distance))

  start_pos, end_pos = None, None
  for pos in sorted(candidate_pos):
    if start_pos is None:
      start_pos = pos
      end_pos = pos
    # We need to check if the previous end_pos is within 2*window_distance as we
    # generate a window of radius window_distance around each position.
    #
    #   <-------end_pos------->
    #                          <-------pos------->
    # where window_distance = ------->
    #
    # If this is the case, we need to merge the two windows.
    elif pos > end_pos + 2 * config.min_windows_distance:
      _add_window(start_pos, end_pos)
      start_pos = pos
      end_pos = pos
    else:
      end_pos = pos
  if start_pos is not None:
    _add_window(start_pos, end_pos)

  return sorted(windows, key=ranges.as_tuple)


def select_windows(config, ref_reader, reads, region):
  """"Process reads to determine candidate windows for local assembly.

  Windows are within range of
    [0 - config.min_windows_distance, ref_len + config.min_windows_distance)

  Args:
    config: learning.genomics.deepvariant.realigner.WindowSelectorOptions
      options determining the behavior of this window selector.
    ref_reader: GenomeReference. Indexed reference genome to query bases.
    reads: A list of genomics.Read records.
    region: nucleus.protos.Range. The region we are processing.

  Returns:
    A list of nucleus.protos.Range protos sorted by their genomic position.
  """
  # This is a fast path for the case where we have no reads, so we have no
  # windows to assemble.
  if not reads:
    return []

  candidates = _candidates_from_reads(config, ref_reader, reads, region)
  return _candidates_to_windows(config, candidates, region.reference_name)