#!/usr/bin/python

''' Function to summarize mapping of reads to contigs, from sam file generated by mapping program.Report the following:
    1. length of each reference contig
    2. percentage of coverage of each reference contig
    3. average coverage depth of each reference contig
    4. number of reads mapped to each reference contig
    5. error rate, in base pairs, including insertion, deletion, and substitutions. (without reference contig, if the sam file only reports M instead of =X, it's impossible to know which one is a match and which one is a substitution.)
    6. total number of base pairs mapped to the reference conig
    5. what else?
    
'''
import re

## ======================================================================
## From CIGAR string, find number of indels and substitutions  
## also the length of the read, the sequence, and the mapped reference region 
## arguments: CIGAR string
## ======================================================================
def cigar(cigar_string):
    """ parse CIGAR string from .sam file, find number of matches, mismatches, insertion (to ref), deletion,
        soft clipping, hard clipping, padding
        return: 1-level dictionary
    """
    char = re.findall('\D',cigar_string) # operation characters, MIDNSHP=X
    char = [x.upper() for x in char]  # convert to upper case
    count = map(int,re.findall('\d+',cigar_string))
    read_len = sum(count[i] for i in xrange(len(char)) if char[i] in 'MIS=XH') # length of the mapped read (including clipped part)
    seq_len = sum(count[i] for i in xrange(len(char)) if char[i] in 'MI=X') # length of the mapped read (excluding lipped part)
    ref_region_len = sum(count[i] for i in xrange(len(char)) if char[i] in 'MD=X') # length of the region on the reference sequence corresponding to the read
    ins_len = sum(count[i] for i in xrange(len(char)) if char[i] == 'I') # insert bps
    del_len = sum(count[i] for i in xrange(len(char)) if char[i] == 'D') # deletion bps

    match_len = None # these are only available with X= in the cigar string, sam 1.4 and after
    sub_len = None
    if '=' in char or 'X' in char: # if either '=' or 'X' is detected, indicating sam 1.4 format for CIGAR string
        match_len = sum(count[i] for i in xrange(len(char)) if char[i] == '=') # matched bps
        sub_len = sum(count[i] for i in xrange(len(char)) if char[i] == 'X') # substitution bps
        align_len = match_len + sub_len
    else:   # sam 1.3, aligned bps from 'M'
        align_len = sum(count[i] for i in xrange(len(char)) if char[i] == 'M') # aligned bps(including matching and substitution)

    pad_len = sum(count[i] for i in xrange(len(char)) if char[i] == 'P') # padded bps, inserted both in reference sequence and the read sequence)
    return {'read_len':read_len, 'seq_len': seq_len, 'ref_len':ref_region_len, 'ins_len':ins_len, 'del_len':del_len, 'match_len':match_len, 'sub_len':sub_len, 'align_len':align_len, 'pad_len':pad_len}

    
##################################################################
## From MD tag, find the number of matched, deletion, and substitution bps 
## arguments: MD tag
##################################################################

def md(MD_tag):
    """
    Given MD tag and a sequence, find the number of matched, deletion, and substitution bps 

    Return: 1-level dictionary 
    """
    #matchLens = re.sub("\D"," ",MD_tag).split() #replace non-digit chars by blanks
    #matchLens = map(int, matchLens) #length of the matches before a mismatch/deletion
    match_len = sum(map(int,re.findall("\d+",MD_tag))) #Get all the numbers in MD tag, these are the number of matches
    nonMatches = re.findall("\D+",MD_tag) #mismatch/deletion, replace digits by blanks
    del_len = 0
    sub_len = 0
    for string in nonMatches:
        if string.startswith('^'): # deletion
            del_len += (len(string) - 1)
        else: # substitution
            sub_len += len(string)
    return {'match_len':match_len, 'del_len':del_len, 'sub_len':sub_len}
## ======================================================================
## parse a single line in sam file with an aligned read record
## extract information
## ======================================================================
def readAlign(alignRecord):
    """ From alignment record for a read, get the following information:
        1. qname - query read name
        2. flag
        3. rname - reference sequence name
        4. pos - 1-based leftmost mapping position on the reference sequence
        5. mapQ - mapping quality
        6. cigarstring - CIGAR string
        the files following are: rnext, pnext, tlen, seq, qual
        10. seq - segment sequence
        
        Output:
        Dictionary with these fields, including some other inferred ones:
        1. qstart, qend, rstart(pos), rend, positions
        2. what else?
    """
    fields = alignRecord.split("\t") # split by tabs
    qname = fields[0]
    flag = format(int(fields[1]),'016b') # binary format
    flag = map(int,flag) # convert flag from string to integers

    # interpret bitwise flags
    if flag[-1] == 1:
        is_onlymap = False
    else:
        is_onlymap = True

    if flag[-2] == 1:
        all_seg_proper = True
    else:
        all_seg_proper = False

    if flag[-3] == 1:
        is_unmapped = True
    else:
        is_unmapped = False

    if flag[-4] == 1:
        is_next_seg_unmapped = True
    else:
        is_next_seg_unmapped = False

    if flag[-5] == 1:
        is_reversecomplement = True
    else:
        is_reversecomplement = False

    if flag[-6] == 1:
        is_next_seg_reversecomplement = True
    else:
        is_next_seg_reversecomplement = False

    if flag[-7] == 1:
        is_first_seg = True
    else:
        is_first_seg = False

    if flag[-8] == 1:
        is_last_seg = True
    else:
        is_last_seg = False

    if flag[-9] == 1:
        is_secondary_alignment = True
    else:
        is_secondary_alignment = False

    if flag[-10] == 1:
        fail_quality = True
    else:
        fail_quality = False

    if flag[-11] == 1:
        is_pcr = True
    else:
        is_pcr = False

    if flag[-12] == 1:
        is_supplementary_alignment = True
    else:
        is_supplementary_alignment = False
    
    rname = fields[2] # reference sequence name
    rstart = int(fields[3]) # starting mapping position on the reference sequence, 1-based
    mapQ = int(fields[4]) # mapping quality
    cigarstring = fields[5] # CIGAR string
    qSeq = fields[9] # segment sequence

    # search for NM tag
    NM_pos = alignRecord.find('NM:i:')
    if NM_pos != -1: # if NM tag exists
        NM = int(alignRecord[(NM_pos+5):alignRecord.find('\t',NM_pos+5)])
    else:
        NM = None
    
    char = re.findall('\D',cigarstring) # operation characters, MIDNSHP=X
    char = [x.upper() for x in char]  # convert to upper case
    count = map(int,re.findall('\d+',cigarstring))
    #print cigarstring, '\t',len(char), '\t', len(count)

    positions = []
    if cigarstring == '*':
        cigar = [(0,'M')]
        rend = rstart
        qstart = 0
        qend = 0
    else:
        cigar = [(count[x], char[x]) for x in xrange(len(char))]

        lastBp = rstart-1
        for bp, ch in cigar:
            if ch in 'MX=D':# matching or deletion
                for i in xrange(bp):
                    positions.append(lastBp+1 + i)
                lastBp = positions[-1]
            elif ch == 'N':
                lastBp += bp # skipped region
        rend = positions[-1]

        qstart = 1
        bp, ch = cigar[0]
        if ch in 'HS':
            qstart = bp + 1
        qlen = sum([x[0] for x in cigar if x[1] in 'MIX='])
        qend = qstart + qlen - 1

    return {'qname':qname, 'flag':int(fields[1]),'is_onlymap':is_onlymap,'all_seg_proper':all_seg_proper,'is_unmapped':is_unmapped,'is_next_seg_unmapped':is_next_seg_unmapped,'is_reversecomplement':is_reversecomplement,'is_next_seg_reversecomplement':is_next_seg_reversecomplement,'is_first_seg':is_first_seg,'is_last_seg':is_last_seg,'is_secondary_alignment':is_secondary_alignment,'is_pcr':is_pcr,'is_supplementary_alignment':is_supplementary_alignment,'rstart':rstart, 'mapQ':mapQ, 'cigarstring':cigarstring,'cigar':cigar,'positions':positions,'rend':rend,'rname':rname,'qSeq':qSeq,'qstart':qstart,'qend':qend,'NM':NM}
