<commit_before># cython: embedsignature=True
# cython: profile=True
###############################################################################
###############################################################################
# Cython wrapper for SAM/BAM/CRAM files based on htslib
###############################################################################
# The principal classes defined in this module are:
#
# class AlignedSegment  an aligned segment (read)
#
# class PileupColumn    a collection of segments (PileupRead) aligned to
#                       a particular genomic position.
#
# class PileupRead      an AlignedSegment aligned to a particular genomic
#                       position. Contains additional attributes with respect
#                       to this.
#
# Additionally this module defines numerous additional classes that are part
# of the internal API. These are:
#
# Various iterator classes to iterate over alignments in sequential (IteratorRow)
# or in a stacked fashion (IteratorColumn):
#
# class IteratorRow
# class IteratorRowRegion
# class IteratorRowHead
# class IteratorRowAll
# class IteratorRowAllRefs
# class IteratorRowSelection
#
###############################################################################
#
# The MIT License
#
# Copyright (c) 2015 <NAME>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
#
###############################################################################
import re
import array
import ctypes
import struct

cimport cython
from cpython cimport array as c_array
from cpython.version cimport PY_MAJOR_VERSION
from cpython cimport PyBytes_FromStringAndSize
from libc.string cimport strchr
from cpython cimport array as c_array
from libc.stdint cimport INT8_MIN, INT16_MIN, INT32_MIN, \
    INT8_MAX, INT16_MAX, INT32_MAX, \
    UINT8_MAX, UINT16_MAX, UINT32_MAX

from pysam.libcutils cimport force_bytes, force_str, \
    charptr_to_str, charptr_to_bytes
from pysam.libcutils cimport qualities_to_qualitystring, qualitystring_to_array, \
    array_to_qualitystring

# Constants for binary tag conversion
cdef char * htslib_types = 'cCsSiIf'
cdef char * parray_types = 'bBhHiIf'

cdef bint IS_PYTHON3 = PY_MAJOR_VERSION >= 3

# translation tables

# cigar code to character and vice versa
cdef char* CODE2CIGAR= "MIDNSHP=XB"
cdef int NCIGAR_CODES = 10

if IS_PYTHON3:
    CIGAR2CODE = dict([y, x] for x, y in enumerate(CODE2CIGAR))
else:
    CIGAR2CODE = dict([ord(y), x] for x, y in enumerate(CODE2CIGAR))

CIGAR_REGEX = re.compile("(\d+)([MIDNSHP=XB])")

#####################################################################
# C multiplication with wrapping around
cdef inline uint32_t c_mul(uint32_t a, uint32_t b):
    return (a * b) & 0xffffffff


#####################################################################
# typecode guessing
cdef inline char map_typecode_htslib_to_python(uint8_t s):
    """map an htslib typecode to the corresponding python typecode
    to be used in the struct or array modules."""

    # map type from htslib to python array
    cdef char * f = strchr(htslib_types, s)

    if f == NULL:
        return 0
    return parray_types[f - htslib_types]

cdef inline uint8_t map_typecode_python_to_htslib(char s):
    """determine value type from type code of array"""
    cdef char * f = strchr(parray_types, s)
    if f == NULL:
        return 0
    return htslib_types[f - parray_types]

# optional tag data manipulation
cdef convert_binary_tag(uint8_t * tag):
    """return bytesize, number of values and array of values
    in aux_data memory location pointed to by tag."""
    cdef uint8_t auxtype
    cdef uint8_t byte_size
    cdef int32_t nvalues
    # get byte size
    auxtype = tag[0]
    byte_size = aux_type2size(auxtype)
    tag += 1
    # get number of values in array
    nvalues = (<int32_t*>tag)[0]
    tag += 4

    # define python array
    cdef c_array.array c_values = array.array(
        chr(map_typecode_htslib_to_python(auxtype)))
    c_array.resize(c_values, nvalues)

    # copy data
    memcpy(c_values.data.as_voidptr, <uint8_t*>tag, nvalues * byte_size)

    # no need to check for endian-ness as bam1_core_t fields
    # and aux_data are in host endian-ness. See sam.c and calls
    # to swap_data
    return byte_size, nvalues, c_values


cdef inline uint8_t get_tag_typecode(value, value_type=None):
    """guess type code for a *value*. If *value_type* is None, the type
    code will be inferred based on the Python type of *value*

    """
    # 0 is unknown typecode
    cdef char typecode = 0
    
    if value_type is None:
        if isinstance(value, int):
            if value < 0:
                if value >= INT8_MIN:
                    typecode = 'c'
                elif value >= INT16_MIN:
                    typecode = 's'
                elif value >= INT32_MIN:
                    typecode = 'i'
            # unsigned ints
            else:
                if value <= UINT8_MAX:
                    typecode = 'C'
                elif value <= UINT16_MAX:
                    typecode = 'S'
                elif value <= UINT32_MAX:
                    typecode = 'I'
        elif isinstance(value, float):
            typecode = 'f'
        elif isinstance(value, str):
            typecode = 'Z'
        elif isinstance(value, bytes):
            typecode = 'Z'
        elif isinstance(value, array.array) or \
                isinstance(value, list) or \
                isinstance(value, tuple):
            typecode = 'B'
    else:
        if value_type in 'aAsSIcCZidfH':
            typecode = force_bytes(value_type)[0]

    return typecode


cdef inline uint8_t get_btag_typecode(value, min_value=None, max_value=None):
    '''returns the value typecode of a value.

    If max is specified, the appropriate type is returned for a range
    where value is the minimum.

    Note that this method returns types from the extended BAM alphabet
    of types that includes tags that are not part of the SAM
    specification.
    '''


    cdef uint8_t typecode

    t = type(value)

    if t is float:
        typecode = 'f'
    elif t is int:
        if max_value is None:
            max_value = value
        if min_value is None:
            min_value = value
        # signed ints
        if min_value < 0:
            if min_value >= INT8_MIN and max_value <= INT8_MAX:
                typecode = 'c'
            elif min_value >= INT16_MIN and max_value <= INT16_MAX:
                typecode = 's'
            elif min_value >= INT32_MIN or max_value <= INT32_MAX:
                typecode = 'i'
            else:
                raise ValueError(
                    "at least one signed integer out of range of "
                    "BAM/SAM specification")
        # unsigned ints
        else:
            if max_value <= UINT8_MAX:
                typecode = 'C'
            elif max_value <= UINT16_MAX:
                typecode = 'S'
            elif max_value <= UINT32_MAX:
                typecode = 'I'
            else:
                raise ValueError(
                    "at least one integer out of range of BAM/SAM specification")
    else:
        # Note: hex strings (H) are not supported yet
        if t is not bytes:
            value = value.encode('ascii')
        if len(value) == 1:
            typecode = 'A'
        else:
            typecode = 'Z'

    return typecode


# mapping python array.array and htslib typecodes to struct typecodes
DATATYPE2FORMAT = {
    ord('c'): ('b', 1),
    ord('C'): ('B', 1),
    ord('s'): ('h', 2),
    ord('S'): ('H', 2),
    ord('i'): ('i', 4),
    ord('I'): ('I', 4),
    ord('f'): ('f', 4),
    ord('d'): ('d', 8),
    ord('A'): ('c', 1),
    ord('a'): ('c', 1)}


cdef inline pack_tags(tags):
    """pack a list of tags. Each tag is a tuple of (tag, tuple).

    Values are packed into the most space efficient data structure
    possible unless the tag contains a third field with the typecode.

    Returns a format string and the associated list of arguments to be
    used in a call to struct.pack_into.
    """
    fmts, args = ["<"], []

    # htslib typecode 
    cdef uint8_t typecode
    for tag in tags:

        if len(tag) == 2:
            pytag, value = tag
            valuetype = None
        elif len(tag) == 3:
            pytag, value, valuetype = tag
        else:
            raise ValueError("malformatted tag: %s" % str(tag))

        if valuetype is None:
            typecode = 0
        else:
            # only first character in valuecode matters
            if IS_PYTHON3:
                typecode = force_bytes(valuetype)[0]
            else:
                typecode = ord(valuetype[0])

        pytag = force_bytes(pytag)
        pytype = type(value)

        if pytype is tuple or pytype is list:
            # binary tags from tuples or lists
            if not typecode:
                # automatically determine value type - first value
                # determines type. If there is a mix of types, the
                # result is undefined.
                typecode = get_btag_typecode(min(value),
                                             min_value=min(value),
                                             max_value=max(value))

            if typecode not in DATATYPE2FORMAT:
                raise ValueError("invalid value type '{}'".format(chr(typecode)))

            datafmt = "2sBBI%i%s" % (len(value), DATATYPE2FORMAT[typecode][0])
            args.extend([pytag[:2],
                         ord("B"),
                         typecode,
                         len(value)] + list(value))

        elif isinstance(value, array.array):
            # binary tags from arrays
            if typecode == 0:
                typecode = map_typecode_python_to_htslib(ord(value.typecode))

                if typecode == 0:
                    raise ValueError("unsupported type code '{}'".format(value.typecode))

            if typecode not in DATATYPE2FORMAT:
                raise ValueError("invalid value type '{}' ({})".format(chr(typecode), array.typecode))
            
            # use array.tostring() to retrieve byte representation and
            # save as bytes
            datafmt = "2sBBI%is" % (len(value) * DATATYPE2FORMAT[typecode][1])
            args.extend([pytag[:2],
                         ord("B"),
                         typecode,
                         len(value),
                         force_bytes(value.tostring())])

        else:
            if typecode == 0:
                typecode = get_tag_typecode(value)
                if typecode == 0:
                    raise ValueError("could not deduce typecode for value {}".format(value))
                
            if typecode == 'a' or typecode == 'A' or typecode == 'Z' or typecode == 'H':
                value = force_bytes(value)

            if typecode == "a":
                typecode = 'A'

            if typecode == 'Z' or typecode == 'H':
                datafmt = "2sB%is" % (len(value)+1)
            else:
                datafmt = "2sB%s" % DATATYPE2FORMAT[typecode][0]
                
            args.extend([pytag[:2],
                         typecode,
                         value])

        fmts.append(datafmt)

    return "".join(fmts), args


cdef inline int32_t calculateQueryLengthWithoutHardClipping(bam1_t * src):
    """return query length computed from CIGAR alignment.

    Length ignores hard-clipped bases.

    Return 0 if there is no CIGAR alignment.
    """

    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)

    if cigar_p == NULL:
        return 0

    cdef uint32_t k, qpos
    cdef int op
    qpos = 0

    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK

        if op == BAM_CMATCH or \
           op == BAM_CINS or \
           op == BAM_CSOFT_CLIP or \
           op == BAM_CEQUAL or \
           op == BAM_CDIFF:
            qpos += cigar_p[k] >> BAM_CIGAR_SHIFT

    return qpos


cdef inline int32_t calculateQueryLengthWithHardClipping(bam1_t * src):
    """return query length computed from CIGAR alignment.

    Length includes hard-clipped bases.

    Return 0 if there is no CIGAR alignment.
    """

    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)

    if cigar_p == NULL:
        return 0

    cdef uint32_t k, qpos
    cdef int op
    qpos = 0

    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK

        if op == BAM_CMATCH or \
           op == BAM_CINS or \
           op == BAM_CSOFT_CLIP or \
           op == BAM_CHARD_CLIP or \
           op == BAM_CEQUAL or \
           op == BAM_CDIFF:
            qpos += cigar_p[k] >> BAM_CIGAR_SHIFT

    return qpos


cdef inline int32_t getQueryStart(bam1_t *src) except -1:
    cdef uint32_t * cigar_p
    cdef uint32_t start_offset = 0
    cdef uint32_t k, op

    cigar_p = pysam_bam_get_cigar(src);
    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK
        if op == BAM_CHARD_CLIP:
            if start_offset != 0 and start_offset != src.core.l_qseq:
                raise ValueError('Invalid clipping in CIGAR string')
        elif op == BAM_CSOFT_CLIP:
            start_offset += cigar_p[k] >> BAM_CIGAR_SHIFT
        else:
            break

    return start_offset


cdef inline int32_t getQueryEnd(bam1_t *src) except -1:
    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)
    cdef uint32_t end_offset = src.core.l_qseq
    cdef uint32_t k, op

    # if there is no sequence, compute length from cigar string
    if end_offset == 0:
        for k from 0 <= k < pysam_get_n_cigar(src):
            op = cigar_p[k] & BAM_CIGAR_MASK
            if op == BAM_CMATCH or \
               op == BAM_CINS or \
               op == BAM_CEQUAL or \
               op == BAM_CDIFF or \
              (op == BAM_CSOFT_CLIP and end_offset == 0):
                end_offset += cigar_p[k] >> BAM_CIGAR_SHIFT
    else:
        # walk backwards in cigar string
        for k from pysam_get_n_cigar(src) > k >= 1:
            op = cigar_p[k] & BAM_CIGAR_MASK
            if op == BAM_CHARD_CLIP:
                if end_offset != src.core.l_qseq:
                    raise ValueError('Invalid clipping in CIGAR string')
            elif op == BAM_CSOFT_CLIP:
                end_offset -= cigar_p[k] >> BAM_CIGAR_SHIFT
            else:
                break

    return end_offset


cdef inline bytes getSequenceInRange(bam1_t *src,
                                     uint32_t start,
                                     uint32_t end):
    """return python string of the sequence in a bam1_t object.
    """

    cdef uint8_t * p
    cdef uint32_t k
    cdef char * s

    if not src.core.l_qseq:
        return None

    seq = PyBytes_FromStringAndSize(NULL, end - start)
    s   = <char*>seq
    p   = pysam_bam_get_seq(src)

    for k from start <= k < end:
        # equivalent to seq_nt16_str[bam1_seqi(s, i)] (see bam.c)
        # note: do not use string literal as it will be a python string
        s[k-start] = seq_nt16_str[p[k/2] >> 4 * (1 - k%2) & 0xf]

    return charptr_to_bytes(seq)


cdef inline object getQualitiesInRange(bam1_t *src,
                                       uint32_t start,
                                       uint32_t end):
    """return python array of quality values from a bam1_t object"""

    cdef uint8_t * p
    cdef uint32_t k

    p = pysam_bam_get_qual(src)
    if p[0] == 0xff:
        return None

    # 'B': unsigned char
    cdef c_array.array result = array.array('B', [0])
    c_array.resize(result, end - start)

    # copy data
    memcpy(result.data.as_voidptr, <void*>&p[start], end - start)

    return result


#####################################################################
## private factory methods
cdef class AlignedSegment
cdef makeAlignedSegment(bam1_t * src, AlignmentFile alignment_file):
    '''return an AlignedSegment object constructed from `src`'''
    # note that the following does not call __init__
    cdef AlignedSegment dest = AlignedSegment.__new__(AlignedSegment)
    dest._delegate = bam_dup1(src)
    dest._alignment_file = alignment_file
    return dest


cdef class PileupColumn
cdef makePileupColumn(bam_pileup1_t ** plp, int tid, int pos,
                      int n_pu, AlignmentFile alignment_file):
    '''return a PileupColumn object constructed from pileup in `plp` and
    setting additional attributes.

    '''
    # note that the following does not call __init__
    cdef PileupColumn dest = PileupColumn.__new__(PileupColumn)
    dest._alignment_file = alignment_file
    dest.plp = plp
    dest.tid = tid
    dest.pos = pos
    dest.n_pu = n_pu
    return dest

cdef class PileupRead
cdef inline makePileupRead(bam_pileup1_t * src, AlignmentFile alignment_file):
    '''return a PileupRead object construted from a bam_pileup1_t * object.'''
    cdef PileupRead dest = PileupRead.__new__(PileupRead)
    dest._alignment = makeAlignedSegment(src.b, alignment_file)
    dest._qpos = src.qpos
    dest._indel = src.indel
    dest._level = src.level
    dest._is_del = src.is_del
    dest._is_head = src.is_head
    dest._is_tail = src.is_tail
    dest._is_refskip = src.is_refskip
    return dest


cdef inline uint32_t get_alignment_length(bam1_t * src):
    cdef int k = 0
    cdef uint32_t l = 0
    if src == NULL:
        return 0
    cdef uint32_t * cigar_p = bam_get_cigar(src)
    if cigar_p == NULL:
        return 0
    cdef int op
    cdef int n = pysam_get_n_cigar(src)
    for k from 0 <= k < n:
        op = cigar_p[k] & BAM_CIGAR_MASK
        if op == BAM_CSOFT_CLIP or op == BAM_CHARD_CLIP:
            continue
        l += cigar_p[k] >> BAM_CIGAR_SHIFT
    return l

cdef inline uint32_t get_md_reference_length(char * md_tag):
    cdef int l = 0
    cdef int md_idx = 0
    cdef int nmatches = 0

    while md_tag[md_idx] != 0:
        if md_tag[md_idx] >= 48 and md_tag[md_idx] <= 57:
            nmatches *= 10
            nmatches += md_tag[md_idx] - 48
            md_idx += 1
            continue
        else:
            l += nmatches
            nmatches = 0
            if md_tag[md_idx] == '^':
                md_idx += 1
                while md_tag[md_idx] >= 65 and md_tag[md_idx] <= 90:
                    md_idx += 1
                    l += 1
            else:
                # save mismatch and change to lower case
                md_idx += 1
                l += 1

    l += nmatches
    return l

# TODO: avoid string copying for getSequenceInRange, reconstituneSequenceFromMD, ...
cdef inline bytes build_alignment_sequence(bam1_t * src):
    """return expanded sequence from MD tag.

    The sequence includes substitutions and both insertions in the
    reference as well as deletions to the reference sequence. Combine
    with the cigar string to reconstitute the query or the reference
    sequence.

    Positions corresponding to `N` (skipped region from the reference)
    in the CIGAR string will not appear in the returned sequence. The
    MD should correspondingly not contain these. Thus proper tags are::

       Deletion from the reference:   cigar=5M1D5M    MD=5^C5
       Skipped region from reference: cigar=5M1N5M    MD=10

    Returns
    -------

    None, if no MD tag is present.

    """
    if src == NULL:
        return None

    cdef uint32_t start = getQueryStart(src)
    cdef uint32_t end = getQueryEnd(src)
    # get read sequence, taking into account soft-clipping
    r = getSequenceInRange(src, start, end)
    cdef char * read_sequence = r
    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)
    if cigar_p == NULL:
        return None

    cdef uint32_t r_idx = 0
    cdef int op
    cdef uint32_t k, i, l, x
    cdef int nmatches = 0
    cdef int s_idx = 0

    cdef uint32_t max_len = get_alignment_length(src)
    if max_len == 0:
        raise ValueError("could not determine alignment length")

    cdef char * s = <char*>calloc(max_len + 1, sizeof(char))
    if s == NULL:
        raise ValueError(
            "could not allocate sequence of length %i" % max_len)

    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK
        l = cigar_p[k] >> BAM_CIGAR_SHIFT
        if op == BAM_CMATCH or op == BAM_CEQUAL or op == BAM_CDIFF:
            for i from 0 <= i < l:
                s[s_idx] = read_sequence[r_idx]
                r_idx += 1
                s_idx += 1
        elif op == BAM_CDEL:
            for i from 0 <= i < l:
                s[s_idx] = '-'
                s_idx += 1
        elif op == BAM_CREF_SKIP:
            pass
        elif op == BAM_CINS:
            for i from 0 <= i < l:
                # encode insertions into reference as lowercase
                s[s_idx] = read_sequence[r_idx] + 32
                r_idx += 1
                s_idx += 1
        elif op == BAM_CSOFT_CLIP:
            pass
        elif op == BAM_CHARD_CLIP:
            pass # advances neither
        elif op == BAM_CPAD:
            raise NotImplementedError(
                "Padding (BAM_CPAD, 6) is currently not supported. "
                "Please implement. Sorry about that.")

    cdef uint8_t * md_tag_ptr = bam_aux_get(src, "MD")
    if md_tag_ptr == NULL:
        seq = PyBytes_FromStringAndSize(s, s_idx)
        free(s)
        return seq

    cdef char * md_tag = <char*>bam_aux2Z(md_tag_ptr)
    cdef int md_idx = 0
    s_idx = 0

    cdef uint32_t md_len = get_md_reference_length(md_tag)
    if md_len > max_len:
        raise AssertionError("Invalid MD tag: MD length {} mismatch with CIGAR length {}".format(md_len, max_len))

    while md_tag[md_idx] != 0:
        # c is numerical
        if md_tag[md_idx] >= 48 and md_tag[md_idx] <= 57:
            nmatches *= 10
            nmatches += md_tag[md_idx] - 48
            md_idx += 1
            continue
        else:
            # save matches up to this point, skipping insertions
            for x from 0 <= x < nmatches:
                while s[s_idx] >= 'a':
                    s_idx += 1
                s_idx += 1
            while s[s_idx] >= 'a':
                s_idx += 1

            r_idx += nmatches
            nmatches = 0
            if md_tag[md_idx] == '^':
                md_idx += 1
                while md_tag[md_idx] >= 65 and md_tag[md_idx] <= 90:
                    assert s[s_idx] == '-'
                    s[s_idx] = md_tag[md_idx]
                    s_idx += 1
                    md_idx += 1<commit_msg>Account for insertion length in max MD ref calculation<commit_after># cython: embedsignature=True
# cython: profile=True
###############################################################################
###############################################################################
# Cython wrapper for SAM/BAM/CRAM files based on htslib
###############################################################################
# The principal classes defined in this module are:
#
# class AlignedSegment  an aligned segment (read)
#
# class PileupColumn    a collection of segments (PileupRead) aligned to
#                       a particular genomic position.
#
# class PileupRead      an AlignedSegment aligned to a particular genomic
#                       position. Contains additional attributes with respect
#                       to this.
#
# Additionally this module defines numerous additional classes that are part
# of the internal API. These are:
#
# Various iterator classes to iterate over alignments in sequential (IteratorRow)
# or in a stacked fashion (IteratorColumn):
#
# class IteratorRow
# class IteratorRowRegion
# class IteratorRowHead
# class IteratorRowAll
# class IteratorRowAllRefs
# class IteratorRowSelection
#
###############################################################################
#
# The MIT License
#
# Copyright (c) 2015 <NAME>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
#
###############################################################################
import re
import array
import ctypes
import struct

cimport cython
from cpython cimport array as c_array
from cpython.version cimport PY_MAJOR_VERSION
from cpython cimport PyBytes_FromStringAndSize
from libc.string cimport strchr
from cpython cimport array as c_array
from libc.stdint cimport INT8_MIN, INT16_MIN, INT32_MIN, \
    INT8_MAX, INT16_MAX, INT32_MAX, \
    UINT8_MAX, UINT16_MAX, UINT32_MAX

from pysam.libcutils cimport force_bytes, force_str, \
    charptr_to_str, charptr_to_bytes
from pysam.libcutils cimport qualities_to_qualitystring, qualitystring_to_array, \
    array_to_qualitystring

# Constants for binary tag conversion
cdef char * htslib_types = 'cCsSiIf'
cdef char * parray_types = 'bBhHiIf'

cdef bint IS_PYTHON3 = PY_MAJOR_VERSION >= 3

# translation tables

# cigar code to character and vice versa
cdef char* CODE2CIGAR= "MIDNSHP=XB"
cdef int NCIGAR_CODES = 10

if IS_PYTHON3:
    CIGAR2CODE = dict([y, x] for x, y in enumerate(CODE2CIGAR))
else:
    CIGAR2CODE = dict([ord(y), x] for x, y in enumerate(CODE2CIGAR))

CIGAR_REGEX = re.compile("(\d+)([MIDNSHP=XB])")

#####################################################################
# C multiplication with wrapping around
cdef inline uint32_t c_mul(uint32_t a, uint32_t b):
    return (a * b) & 0xffffffff


#####################################################################
# typecode guessing
cdef inline char map_typecode_htslib_to_python(uint8_t s):
    """map an htslib typecode to the corresponding python typecode
    to be used in the struct or array modules."""

    # map type from htslib to python array
    cdef char * f = strchr(htslib_types, s)

    if f == NULL:
        return 0
    return parray_types[f - htslib_types]

cdef inline uint8_t map_typecode_python_to_htslib(char s):
    """determine value type from type code of array"""
    cdef char * f = strchr(parray_types, s)
    if f == NULL:
        return 0
    return htslib_types[f - parray_types]

# optional tag data manipulation
cdef convert_binary_tag(uint8_t * tag):
    """return bytesize, number of values and array of values
    in aux_data memory location pointed to by tag."""
    cdef uint8_t auxtype
    cdef uint8_t byte_size
    cdef int32_t nvalues
    # get byte size
    auxtype = tag[0]
    byte_size = aux_type2size(auxtype)
    tag += 1
    # get number of values in array
    nvalues = (<int32_t*>tag)[0]
    tag += 4

    # define python array
    cdef c_array.array c_values = array.array(
        chr(map_typecode_htslib_to_python(auxtype)))
    c_array.resize(c_values, nvalues)

    # copy data
    memcpy(c_values.data.as_voidptr, <uint8_t*>tag, nvalues * byte_size)

    # no need to check for endian-ness as bam1_core_t fields
    # and aux_data are in host endian-ness. See sam.c and calls
    # to swap_data
    return byte_size, nvalues, c_values


cdef inline uint8_t get_tag_typecode(value, value_type=None):
    """guess type code for a *value*. If *value_type* is None, the type
    code will be inferred based on the Python type of *value*

    """
    # 0 is unknown typecode
    cdef char typecode = 0
    
    if value_type is None:
        if isinstance(value, int):
            if value < 0:
                if value >= INT8_MIN:
                    typecode = 'c'
                elif value >= INT16_MIN:
                    typecode = 's'
                elif value >= INT32_MIN:
                    typecode = 'i'
            # unsigned ints
            else:
                if value <= UINT8_MAX:
                    typecode = 'C'
                elif value <= UINT16_MAX:
                    typecode = 'S'
                elif value <= UINT32_MAX:
                    typecode = 'I'
        elif isinstance(value, float):
            typecode = 'f'
        elif isinstance(value, str):
            typecode = 'Z'
        elif isinstance(value, bytes):
            typecode = 'Z'
        elif isinstance(value, array.array) or \
                isinstance(value, list) or \
                isinstance(value, tuple):
            typecode = 'B'
    else:
        if value_type in 'aAsSIcCZidfH':
            typecode = force_bytes(value_type)[0]

    return typecode


cdef inline uint8_t get_btag_typecode(value, min_value=None, max_value=None):
    '''returns the value typecode of a value.

    If max is specified, the appropriate type is returned for a range
    where value is the minimum.

    Note that this method returns types from the extended BAM alphabet
    of types that includes tags that are not part of the SAM
    specification.
    '''


    cdef uint8_t typecode

    t = type(value)

    if t is float:
        typecode = 'f'
    elif t is int:
        if max_value is None:
            max_value = value
        if min_value is None:
            min_value = value
        # signed ints
        if min_value < 0:
            if min_value >= INT8_MIN and max_value <= INT8_MAX:
                typecode = 'c'
            elif min_value >= INT16_MIN and max_value <= INT16_MAX:
                typecode = 's'
            elif min_value >= INT32_MIN or max_value <= INT32_MAX:
                typecode = 'i'
            else:
                raise ValueError(
                    "at least one signed integer out of range of "
                    "BAM/SAM specification")
        # unsigned ints
        else:
            if max_value <= UINT8_MAX:
                typecode = 'C'
            elif max_value <= UINT16_MAX:
                typecode = 'S'
            elif max_value <= UINT32_MAX:
                typecode = 'I'
            else:
                raise ValueError(
                    "at least one integer out of range of BAM/SAM specification")
    else:
        # Note: hex strings (H) are not supported yet
        if t is not bytes:
            value = value.encode('ascii')
        if len(value) == 1:
            typecode = 'A'
        else:
            typecode = 'Z'

    return typecode


# mapping python array.array and htslib typecodes to struct typecodes
DATATYPE2FORMAT = {
    ord('c'): ('b', 1),
    ord('C'): ('B', 1),
    ord('s'): ('h', 2),
    ord('S'): ('H', 2),
    ord('i'): ('i', 4),
    ord('I'): ('I', 4),
    ord('f'): ('f', 4),
    ord('d'): ('d', 8),
    ord('A'): ('c', 1),
    ord('a'): ('c', 1)}


cdef inline pack_tags(tags):
    """pack a list of tags. Each tag is a tuple of (tag, tuple).

    Values are packed into the most space efficient data structure
    possible unless the tag contains a third field with the typecode.

    Returns a format string and the associated list of arguments to be
    used in a call to struct.pack_into.
    """
    fmts, args = ["<"], []

    # htslib typecode 
    cdef uint8_t typecode
    for tag in tags:

        if len(tag) == 2:
            pytag, value = tag
            valuetype = None
        elif len(tag) == 3:
            pytag, value, valuetype = tag
        else:
            raise ValueError("malformatted tag: %s" % str(tag))

        if valuetype is None:
            typecode = 0
        else:
            # only first character in valuecode matters
            if IS_PYTHON3:
                typecode = force_bytes(valuetype)[0]
            else:
                typecode = ord(valuetype[0])

        pytag = force_bytes(pytag)
        pytype = type(value)

        if pytype is tuple or pytype is list:
            # binary tags from tuples or lists
            if not typecode:
                # automatically determine value type - first value
                # determines type. If there is a mix of types, the
                # result is undefined.
                typecode = get_btag_typecode(min(value),
                                             min_value=min(value),
                                             max_value=max(value))

            if typecode not in DATATYPE2FORMAT:
                raise ValueError("invalid value type '{}'".format(chr(typecode)))

            datafmt = "2sBBI%i%s" % (len(value), DATATYPE2FORMAT[typecode][0])
            args.extend([pytag[:2],
                         ord("B"),
                         typecode,
                         len(value)] + list(value))

        elif isinstance(value, array.array):
            # binary tags from arrays
            if typecode == 0:
                typecode = map_typecode_python_to_htslib(ord(value.typecode))

                if typecode == 0:
                    raise ValueError("unsupported type code '{}'".format(value.typecode))

            if typecode not in DATATYPE2FORMAT:
                raise ValueError("invalid value type '{}' ({})".format(chr(typecode), array.typecode))
            
            # use array.tostring() to retrieve byte representation and
            # save as bytes
            datafmt = "2sBBI%is" % (len(value) * DATATYPE2FORMAT[typecode][1])
            args.extend([pytag[:2],
                         ord("B"),
                         typecode,
                         len(value),
                         force_bytes(value.tostring())])

        else:
            if typecode == 0:
                typecode = get_tag_typecode(value)
                if typecode == 0:
                    raise ValueError("could not deduce typecode for value {}".format(value))
                
            if typecode == 'a' or typecode == 'A' or typecode == 'Z' or typecode == 'H':
                value = force_bytes(value)

            if typecode == "a":
                typecode = 'A'

            if typecode == 'Z' or typecode == 'H':
                datafmt = "2sB%is" % (len(value)+1)
            else:
                datafmt = "2sB%s" % DATATYPE2FORMAT[typecode][0]
                
            args.extend([pytag[:2],
                         typecode,
                         value])

        fmts.append(datafmt)

    return "".join(fmts), args


cdef inline int32_t calculateQueryLengthWithoutHardClipping(bam1_t * src):
    """return query length computed from CIGAR alignment.

    Length ignores hard-clipped bases.

    Return 0 if there is no CIGAR alignment.
    """

    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)

    if cigar_p == NULL:
        return 0

    cdef uint32_t k, qpos
    cdef int op
    qpos = 0

    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK

        if op == BAM_CMATCH or \
           op == BAM_CINS or \
           op == BAM_CSOFT_CLIP or \
           op == BAM_CEQUAL or \
           op == BAM_CDIFF:
            qpos += cigar_p[k] >> BAM_CIGAR_SHIFT

    return qpos


cdef inline int32_t calculateQueryLengthWithHardClipping(bam1_t * src):
    """return query length computed from CIGAR alignment.

    Length includes hard-clipped bases.

    Return 0 if there is no CIGAR alignment.
    """

    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)

    if cigar_p == NULL:
        return 0

    cdef uint32_t k, qpos
    cdef int op
    qpos = 0

    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK

        if op == BAM_CMATCH or \
           op == BAM_CINS or \
           op == BAM_CSOFT_CLIP or \
           op == BAM_CHARD_CLIP or \
           op == BAM_CEQUAL or \
           op == BAM_CDIFF:
            qpos += cigar_p[k] >> BAM_CIGAR_SHIFT

    return qpos


cdef inline int32_t getQueryStart(bam1_t *src) except -1:
    cdef uint32_t * cigar_p
    cdef uint32_t start_offset = 0
    cdef uint32_t k, op

    cigar_p = pysam_bam_get_cigar(src);
    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK
        if op == BAM_CHARD_CLIP:
            if start_offset != 0 and start_offset != src.core.l_qseq:
                raise ValueError('Invalid clipping in CIGAR string')
        elif op == BAM_CSOFT_CLIP:
            start_offset += cigar_p[k] >> BAM_CIGAR_SHIFT
        else:
            break

    return start_offset


cdef inline int32_t getQueryEnd(bam1_t *src) except -1:
    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)
    cdef uint32_t end_offset = src.core.l_qseq
    cdef uint32_t k, op

    # if there is no sequence, compute length from cigar string
    if end_offset == 0:
        for k from 0 <= k < pysam_get_n_cigar(src):
            op = cigar_p[k] & BAM_CIGAR_MASK
            if op == BAM_CMATCH or \
               op == BAM_CINS or \
               op == BAM_CEQUAL or \
               op == BAM_CDIFF or \
              (op == BAM_CSOFT_CLIP and end_offset == 0):
                end_offset += cigar_p[k] >> BAM_CIGAR_SHIFT
    else:
        # walk backwards in cigar string
        for k from pysam_get_n_cigar(src) > k >= 1:
            op = cigar_p[k] & BAM_CIGAR_MASK
            if op == BAM_CHARD_CLIP:
                if end_offset != src.core.l_qseq:
                    raise ValueError('Invalid clipping in CIGAR string')
            elif op == BAM_CSOFT_CLIP:
                end_offset -= cigar_p[k] >> BAM_CIGAR_SHIFT
            else:
                break

    return end_offset


cdef inline bytes getSequenceInRange(bam1_t *src,
                                     uint32_t start,
                                     uint32_t end):
    """return python string of the sequence in a bam1_t object.
    """

    cdef uint8_t * p
    cdef uint32_t k
    cdef char * s

    if not src.core.l_qseq:
        return None

    seq = PyBytes_FromStringAndSize(NULL, end - start)
    s   = <char*>seq
    p   = pysam_bam_get_seq(src)

    for k from start <= k < end:
        # equivalent to seq_nt16_str[bam1_seqi(s, i)] (see bam.c)
        # note: do not use string literal as it will be a python string
        s[k-start] = seq_nt16_str[p[k/2] >> 4 * (1 - k%2) & 0xf]

    return charptr_to_bytes(seq)


cdef inline object getQualitiesInRange(bam1_t *src,
                                       uint32_t start,
                                       uint32_t end):
    """return python array of quality values from a bam1_t object"""

    cdef uint8_t * p
    cdef uint32_t k

    p = pysam_bam_get_qual(src)
    if p[0] == 0xff:
        return None

    # 'B': unsigned char
    cdef c_array.array result = array.array('B', [0])
    c_array.resize(result, end - start)

    # copy data
    memcpy(result.data.as_voidptr, <void*>&p[start], end - start)

    return result


#####################################################################
## private factory methods
cdef class AlignedSegment
cdef makeAlignedSegment(bam1_t * src, AlignmentFile alignment_file):
    '''return an AlignedSegment object constructed from `src`'''
    # note that the following does not call __init__
    cdef AlignedSegment dest = AlignedSegment.__new__(AlignedSegment)
    dest._delegate = bam_dup1(src)
    dest._alignment_file = alignment_file
    return dest


cdef class PileupColumn
cdef makePileupColumn(bam_pileup1_t ** plp, int tid, int pos,
                      int n_pu, AlignmentFile alignment_file):
    '''return a PileupColumn object constructed from pileup in `plp` and
    setting additional attributes.

    '''
    # note that the following does not call __init__
    cdef PileupColumn dest = PileupColumn.__new__(PileupColumn)
    dest._alignment_file = alignment_file
    dest.plp = plp
    dest.tid = tid
    dest.pos = pos
    dest.n_pu = n_pu
    return dest

cdef class PileupRead
cdef inline makePileupRead(bam_pileup1_t * src, AlignmentFile alignment_file):
    '''return a PileupRead object construted from a bam_pileup1_t * object.'''
    cdef PileupRead dest = PileupRead.__new__(PileupRead)
    dest._alignment = makeAlignedSegment(src.b, alignment_file)
    dest._qpos = src.qpos
    dest._indel = src.indel
    dest._level = src.level
    dest._is_del = src.is_del
    dest._is_head = src.is_head
    dest._is_tail = src.is_tail
    dest._is_refskip = src.is_refskip
    return dest


cdef inline uint32_t get_alignment_length(bam1_t * src):
    cdef int k = 0
    cdef uint32_t l = 0
    if src == NULL:
        return 0
    cdef uint32_t * cigar_p = bam_get_cigar(src)
    if cigar_p == NULL:
        return 0
    cdef int op
    cdef int n = pysam_get_n_cigar(src)
    for k from 0 <= k < n:
        op = cigar_p[k] & BAM_CIGAR_MASK
        if op == BAM_CSOFT_CLIP or op == BAM_CHARD_CLIP:
            continue
        l += cigar_p[k] >> BAM_CIGAR_SHIFT
    return l

cdef inline uint32_t get_md_reference_length(char * md_tag):
    cdef int l = 0
    cdef int md_idx = 0
    cdef int nmatches = 0

    while md_tag[md_idx] != 0:
        if md_tag[md_idx] >= 48 and md_tag[md_idx] <= 57:
            nmatches *= 10
            nmatches += md_tag[md_idx] - 48
            md_idx += 1
            continue
        else:
            l += nmatches
            nmatches = 0
            if md_tag[md_idx] == '^':
                md_idx += 1
                while md_tag[md_idx] >= 65 and md_tag[md_idx] <= 90:
                    md_idx += 1
                    l += 1
            else:
                md_idx += 1
                l += 1

    l += nmatches
    return l

# TODO: avoid string copying for getSequenceInRange, reconstituneSequenceFromMD, ...
cdef inline bytes build_alignment_sequence(bam1_t * src):
    """return expanded sequence from MD tag.

    The sequence includes substitutions and both insertions in the
    reference as well as deletions to the reference sequence. Combine
    with the cigar string to reconstitute the query or the reference
    sequence.

    Positions corresponding to `N` (skipped region from the reference)
    in the CIGAR string will not appear in the returned sequence. The
    MD should correspondingly not contain these. Thus proper tags are::

       Deletion from the reference:   cigar=5M1D5M    MD=5^C5
       Skipped region from reference: cigar=5M1N5M    MD=10

    Returns
    -------

    None, if no MD tag is present.

    """
    if src == NULL:
        return None

    cdef uint32_t start = getQueryStart(src)
    cdef uint32_t end = getQueryEnd(src)
    # get read sequence, taking into account soft-clipping
    r = getSequenceInRange(src, start, end)
    cdef char * read_sequence = r
    cdef uint32_t * cigar_p = pysam_bam_get_cigar(src)
    if cigar_p == NULL:
        return None

    cdef uint32_t r_idx = 0
    cdef int op
    cdef uint32_t k, i, l, x
    cdef int nmatches = 0
    cdef int s_idx = 0

    cdef uint32_t max_len = get_alignment_length(src)
    if max_len == 0:
        raise ValueError("could not determine alignment length")

    cdef char * s = <char*>calloc(max_len + 1, sizeof(char))
    if s == NULL:
        raise ValueError(
            "could not allocate sequence of length %i" % max_len)

    for k from 0 <= k < pysam_get_n_cigar(src):
        op = cigar_p[k] & BAM_CIGAR_MASK
        l = cigar_p[k] >> BAM_CIGAR_SHIFT
        if op == BAM_CMATCH or op == BAM_CEQUAL or op == BAM_CDIFF:
            for i from 0 <= i < l:
                s[s_idx] = read_sequence[r_idx]
                r_idx += 1
                s_idx += 1
        elif op == BAM_CDEL:
            for i from 0 <= i < l:
                s[s_idx] = '-'
                s_idx += 1
        elif op == BAM_CREF_SKIP:
            pass
        elif op == BAM_CINS:
            for i from 0 <= i < l:
                # encode insertions into reference as lowercase
                s[s_idx] = read_sequence[r_idx] + 32
                r_idx += 1
                s_idx += 1
        elif op == BAM_CSOFT_CLIP:
            pass
        elif op == BAM_CHARD_CLIP:
            pass # advances neither
        elif op == BAM_CPAD:
            raise NotImplementedError(
                "Padding (BAM_CPAD, 6) is currently not supported. "
                "Please implement. Sorry about that.")

    cdef uint8_t * md_tag_ptr = bam_aux_get(src, "MD")
    if md_tag_ptr == NULL:
        seq = PyBytes_FromStringAndSize(s, s_idx)
        free(s)
        return seq

    cdef char * md_tag = <char*>bam_aux2Z(md_tag_ptr)
    cdef int md_idx = 0
    s_idx = 0

    # Check if MD tag is valid by matching CIGAR length to MD tag defined length
    # Insertions would be in addition to what is described by MD, so we calculate
    # the number of insertions seperately.
    insertions = 0

    while s[s_idx] != 0:
        if s[s_idx] >= 'a':
            insertions += 1
        s_idx += 1
    s_idx = 0

    cdef uint32_t md_len = get_md_reference_length(md_tag)
    if md_len + insertions > max_len:
        raise AssertionError("Invalid MD tag: MD length {} mismatch with CIGAR length {}".format(md_len, max_len))

    while md_tag[md_idx] != 0:
        # c is numerical
        if md_tag[md_idx] >= 48 and md_tag[md_idx] <= 57:
            nmatches *= 10
            nmatches += md_tag[md_idx] - 48
            md_idx += 1
            continue
        else:
            # save matches up to this point, skipping insertions
            for x from 0 <= x < nmatches:
                while s[s_idx] >= 'a':
                    s_idx += 1
                s_idx += 1
            while s[s_idx] >= 'a':
                s_idx += 1

            r_idx += nmatches
            nmatches = 0
            if md_tag[md_idx] == '^':
                md_idx += 1
                while md_tag[md_idx] >= 65 and md_tag[md_idx] <= 90:
                    assert s[s_idx] == '-'
                    s[s_idx] = md_tag[md_idx]
                    s_idx += 1
                    md_idx += 1
{'fraction_non_alphanumeric': 0.105665087155187, 'fraction_numerical': 0.013475207310881705, 'mean_word_length': 2.0758670050100854, 'pattern_counts': {'":': 2, '<': 93, '<?xml version=': 0, '>': 82, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1853557', 'n_tokens_mistral': 15300, 'n_tokens_neox': 14279, 'n_words': 4551}