# Copyright 2018 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 working with alignment CIGAR operations.

The CIGAR format is defined within the SAM spec, available at

https://samtools.github.io/hts-specs/SAMv1.pdf

This module provides utility functions for interacting with the parsed

representations of CIGAR strings.

"""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import re

import six

from third_party.nucleus.protos import cigar_pb2

# A frozenset of all CigarUnit.Operation enum values that advance the alignment

# with respect to the reference genome and read, respectively.

REF_ADVANCING_OPS = frozenset([

    cigar_pb2.CigarUnit.ALIGNMENT_MATCH,

    cigar_pb2.CigarUnit.SEQUENCE_MATCH,

    cigar_pb2.CigarUnit.DELETE,

    cigar_pb2.CigarUnit.SKIP,

    cigar_pb2.CigarUnit.SEQUENCE_MISMATCH

])

READ_ADVANCING_OPS = frozenset([

    cigar_pb2.CigarUnit.ALIGNMENT_MATCH,

    cigar_pb2.CigarUnit.SEQUENCE_MATCH,

    cigar_pb2.CigarUnit.INSERT,

    cigar_pb2.CigarUnit.CLIP_SOFT,

    cigar_pb2.CigarUnit.SEQUENCE_MISMATCH

])

# A map from CigarUnit.Operation (e.g., CigarUnit.ALIGNMENT_MATCH) enum values

# to their corresponding single character cigar codes (e.g., 'M').

CIGAR_OPS_TO_CHAR = {

    cigar_pb2.CigarUnit.ALIGNMENT_MATCH: 'M',

    cigar_pb2.CigarUnit.INSERT: 'I',

    cigar_pb2.CigarUnit.DELETE: 'D',

    cigar_pb2.CigarUnit.SKIP: 'N',

    cigar_pb2.CigarUnit.CLIP_SOFT: 'S',

    cigar_pb2.CigarUnit.CLIP_HARD: 'H',

    cigar_pb2.CigarUnit.PAD: 'P',

    cigar_pb2.CigarUnit.SEQUENCE_MATCH: '=',

    cigar_pb2.CigarUnit.SEQUENCE_MISMATCH: 'X',

}

# A map from single character cigar codes (e.g., 'M') to their corresponding

# CigarUnit.Operation (e.g., CigarUnit.ALIGNMENT_MATCH) enum values.

CHAR_TO_CIGAR_OPS = {v: k for k, v in CIGAR_OPS_TO_CHAR.items()}

# All of the CigarUnit.Operation values in a frozen set.

ALL_CIGAR_OPS = frozenset(CIGAR_OPS_TO_CHAR.keys())

# Regular expression that matches only valid full cigar strings.

VALID_CIGAR_RE = re.compile(

    r'^(\d+[' + ''.join(CHAR_TO_CIGAR_OPS.keys()) + '])+$')

# Regular expression that matches a single len/op cigar element. The match is

# grouped, so CIGAR_STR_SPLITTER_RE.finditer(cigar_str) returns grouped units

# of the cigar string in order.

CIGAR_STR_SPLITTER_RE = re.compile(

    r'(\d+[' + ''.join(CHAR_TO_CIGAR_OPS.keys()) + '])')

def format_cigar_units(cigar_units):

  """Returns the string version of an iterable of CigarUnit protos.

  Args:

    cigar_units: iterable[CigarUnit] protos.

  Returns:

    A string representation of the CigarUnit protos that conforms to the

    CIGAR string specification.

  """

  return ''.join(

      str(unit.operation_length) + CIGAR_OPS_TO_CHAR[unit.operation]

      for unit in cigar_units)

def parse_cigar_string(cigar_str):

  """Parse a cigar string into a list of cigar units.

  For example, if cigar_str is 150M2S, this function will return:

  [

    CigarUnit(operation=ALIGNMENT_MATCH, operation_length=150),

    CigarUnit(operation=SOFT_CLIP, operation_length=2)

  ]

  Args:

    cigar_str: str containing a valid cigar.

  Returns:

    list[cigar_pb2.CigarUnit].

  Raises:

    ValueError: If cigar_str isn't a well-formed CIGAR.

  """

  if not cigar_str:

    raise ValueError('cigar_str cannot be empty')

  if not VALID_CIGAR_RE.match(cigar_str):

    raise ValueError('Malformed CIGAR string {}'.format(cigar_str))

  parts = CIGAR_STR_SPLITTER_RE.finditer(cigar_str)

  return [to_cigar_unit(part.group(1)) for part in parts]

def alignment_length(cigar_units):

  """Computes the span in basepairs of the cigar units.

  Args:

    cigar_units: iterable[CigarUnit] whose alignment length we want to compute.

  Returns:

    The number of basepairs spanned by the cigar_units.

  """

  return sum(unit.operation_length

             for unit in cigar_units

             if unit.operation in REF_ADVANCING_OPS)

def to_cigar_unit(source):

  """Creates a cigar_pb2 CigarUnit from source.

  This function attempts to convert source into a CigarUnit protobuf. If

  source is a string, it must be a single CIGAR string specification like

  '12M'. If source is a tuple or a list, must have exactly two elements

  (operation_length, opstr). operation_length can be a string or int, and must

  be >= 1. opstr should be a single character CIGAR specification (e.g., 'M').

  If source is already a CigarUnit, it is just passed through unmodified.

  Args:

    source: many types allowed. The object we want to convert to a CigarUnit

      proto.

  Returns:

    CigarUnit proto with operation_length and operation set to values from

      source.

  Raises:

    ValueError: if source cannot be converted or is malformed.

  """

  try:

    if isinstance(source, cigar_pb2.CigarUnit):

      return source

    elif isinstance(source, six.string_types):

      l, op = source[:-1], source[-1]

    elif isinstance(source, (tuple, list)):

      l, op = source

    else:

      raise ValueError('Unexpected source', source)

    if isinstance(op, six.string_types):

      op = CHAR_TO_CIGAR_OPS[op]

    l = int(l)

    if l < 1:

      raise ValueError('Length must be >= 1', l)

    return cigar_pb2.CigarUnit(operation=op, operation_length=int(l))

  except (KeyError, IndexError):

    raise ValueError('Failed to convert {} into a CigarUnit'.format(source))

def to_cigar_units(source):

  """Converts object to a list of CigarUnit.

  This function attempts to convert source into a list of CigarUnit protos.

  If source is a string, we assume it is a CIGAR string and call

  parse_cigar_string on it, returning the result. It not, we assume it's an

  iterable containing element to be converted with to_cigar_unit(). The

  resulting list of converted elements is returned.

  Args:

    source: str or iterable to convert to CigarUnit protos.

  Returns:

    list[CigarUnit].

  """

  if isinstance(source, six.string_types):

    return parse_cigar_string(source)

  else:

    return [to_cigar_unit(singleton) for singleton in source]