#!/usr/bin/python2.7
# -*- Mode: Python -*-

import pysam
import sys
from optparse import OptionParser
import re
import os


cigarPattern   = '([0-9]+[MIDNSHP])'
cigarSearch    = re.compile(cigarPattern)
atomicCigarPattern   = '([0-9]+)([MIDNSHP])'
atomicCigarSearch    =  re.compile(atomicCigarPattern)

class refPos (object):
    """
    struct to store the end position of reference CIGAR operations
    """
    def __init__(self, rPos):
        self.rPos = int(rPos)


class queryPos (object):
    """
    struct to store the start and end positions of query CIGAR operations
    """
    def __init__(self, qsPos, qePos, qLen):
        self.qsPos = int(qsPos)
        self.qePos = int(qePos)
        self.qLen  = int(qLen)


class cigarOp (object):
    """
    sturct to store a discrete CIGAR operations
    """
    def __init__(self, opLength, op):
        self.length = int(opLength)
        self.op     = op

class SAM (object):
    """
    __very__ basic class for SAM input.
    """
    def __init__(self, samList = []):
        if len(samList) > 0:
            self.query    = samList[0]
            self.flag     = int(samList[1])
            self.ref      = samList[2]
            self.pos      = int(samList[3])
            self.mapq     = int(samList[4])
            self.cigar    = samList[5]
            self.matRef   = samList[6]
            self.matePos  = int(samList[7])
            self.iSize    = int(samList[8])
            self.seq      = samList[9]
            self.qual     = samList[10]
            self.tags     = samList[11:]#tags is a list of each tag:vtype:value sets
            self.valid    = 1
        else:
            self.valid = 0
            self.query = 'null'

    def extractCigarOps(self):
        if (self.cigar == "*"):
            self.cigarOps = []
        elif (self.flag & 0x0010):
            cigarOpStrings = cigarSearch.findall(self.cigar)
            cigarOps = []
            for opString in cigarOpStrings:
                cigarOpList = atomicCigarSearch.findall(opString)
                # "struct" for the op and it's length
                cigar = cigarOp(cigarOpList[0][0], cigarOpList[0][1])
                # add to the list of cigarOps
                cigarOps.append(cigar)
                self.cigarOps = cigarOps
            cigarOps.reverse()
            ##do in reverse order because negative strand##
        else:
            cigarOpStrings = cigarSearch.findall(self.cigar)
            cigarOps = []
            for opString in cigarOpStrings:
                cigarOpList = atomicCigarSearch.findall(opString)
                # "struct" for the op and it's length
                cigar = cigarOp(cigarOpList[0][0], cigarOpList[0][1])
                # add to the list of cigarOps
                cigarOps.append(cigar)
                self.cigarOps = cigarOps

    def getCigarOps(self):
        return self.cigarOps
    
    def isReverseStrand(self):
        if (self.flag & 0x0010):
            return True
        return False

    def extractTagValue (self, tagID):
        for tag in self.tags:
            tagParts = tag.split(':', 2);
            if (tagParts[0] == tagID):
                if (tagParts[1] == 'i'):
                    return int(tagParts[2]);
                elif (tagParts[1] == 'H'):
                    return int(tagParts[2],16);
                return tagParts[2];
        return None;

def processSam(opts):
    samFile = opts.samFile;
    if samFile == "stdin":
        s = sys.stdin
    else:
        s = open(samFile, 'r')
    blockList = []
    prevSam = SAM()
    inBlock = 0
    prefix_out = opts.prefix + "_telomeres.txt"
    out_f = open(prefix_out, "w")
    chrom = ""
    for line in s:
        if line[0] == "@":
            continue
        samList = line.strip().split('\t')
        currSam = SAM(samList)
        currSam.extractCigarOps()
        mapped_length = mappedLength(currSam.cigarOps)
        if mapped_length < 500:
            continue
        if currSam.query != prevSam.query:
            makeBedpe(blockList, opts)
            blockList = [currSam]
            prevSam = currSam
        else:
            blockList.append(currSam)
        if chrom != currSam.ref: 
            if chrom != "":
                # Generate for each chromosome the split
                outputBed_Tel(prev_row, out_f,start=False, output_prefix=prefix_out) 
            chrom = currSam.ref 
            outputBed_Tel(currSam,out_f,start=True,output_prefix=prefix_out) 
        prev_row = currSam 
    # Need to process the last block at end of file.
    makeBedpe(blockList, opts)
    out_f.close()


def calcOverlap(bed1, bed2):
    o = 1 + min(bed1[6], bed2[6]) - max(bed1[5], bed2[5])
    return max(0, o)

totalLines = 0;
def outputBedpe(bed1, bed2, opts):
    # set1 = set(range(bed1[5], bed1[6] + 1))
    # set2 = set(range(bed2[5], bed2[6] + 1))
    # nonOverlap1 = len(set1) - len(set1 & set2)
    # nonOverlap2 = len(set2) - len(set1 & set2)
    overlap = calcOverlap(bed1, bed2)
    nonOverlap1 = (1 + bed1[6] - bed1[5]) - overlap
    nonOverlap2 = (1 + bed2[6] - bed2[5]) - overlap
    #print(bed1)
    #print(bed2)
    minNonOverlap = min(nonOverlap1, nonOverlap2)  
    #print opts.MNO
    #print overlap
    #print minNonOverlap
    overlap_two = (bed2[1]-bed1[2])
    #print("OVERLAP", overlap_two)
    #print(minNonOverlap)
    #print(bed2[5] - bed1[6])
    #print(overlap_two - (bed2[5] - bed1[6]))
    #print(overlap)
    if (minNonOverlap >= opts.MNO):
        global totalLines;
        totalLines += 1;
        # sys.stderr.write(bed1[3] + ' ' + str(nonOverlap1) + ' ' + str(nonOverlap2) + ' ' + str(overlap) + ' ' + str(minNonOverlap) + '\n')
        print str(bed1[0]) + "\t" + str(bed1[1]) + "\t" + str(bed1[2]) + "\t" + str(bed2[0]) + "\t" + \
              str(bed2[1]) + "\t" + str(bed2[2]) + "\t" + str(bed1[3]) + "\t" + str(bed1[8]+bed2[8]) + "\t" + str(bed1[4]) + \
              "\t"  + str(bed2[4]) + "\t" + str(bed1[5]) + "\t" + str(bed1[6]) + "\t" + \
              str(bed2[5]) + "\t" + str(bed2[6])  + "\t" + str(minNonOverlap)  + "\t" + str(bed1[9]) + "\t" + \
              "MQ1="+str(bed1[7])+":"+"MQ2="+str(bed2[7])+":"+"AS1="+str(bed1[8])+":"+"AS2="+str(bed2[8])
        #print("WOW")

def checkforSplitAtStart(cigarOps, start=True):
    
    qsPos = 0
    qePos = 0
    qLen  = 0
    # if first op is a H, need to shift start position 
    # the opPosition counter sees if the for loop is looking at the first index of the cigar object    
    opPosition = 0
    cigar = cigarOps[0]
    if opPosition == 0 and (cigar.op == 'H' or cigar.op == 'S'):
        return cigar.length
    else:
        return None

def mappedLength(cigarOps): 
    read_length = 0
    for cigar in cigarOps:
        if cigar.op == "M":
            read_length += cigar.length
        elif cigar.op == "D":
            read_length += cigar.length
    return read_length

def get_Tel_seq(read_info, cigarOps, start=True): 
    read_length = 0
    print(cigarOps[0].length)

def outputBed_Tel(read_info,out_f, output_prefix ,start=True):
    """
        Extract split bed file with some telomere information.
    """
    cigar_string = read_info.cigar
    cigarOpStrings = cigarSearch.findall(cigar_string)
    cigarOps = []
    sample_name = os.path.basename(output_prefix)
    sample_name = (sample_name.split("_")[0])
    output_directory = os.path.dirname(output_prefix)
    telomere_folder = os.path.join(output_directory,sample_name, "event_fastas","telomeres")
    sequences_fasta = pysam.Fastafile(os.path.join(output_directory, sample_name +".fasta"))
    try:
        os.makedirs(telomere_folder)
    except OSError:
        pass
    for opString in cigarOpStrings:
        cigarOpList = atomicCigarSearch.findall(opString)
        # "struct" for the op and it's length
        cigar = cigarOp(cigarOpList[0][0], cigarOpList[0][1])
        # add to the list of cigarOps
        cigarOps.append(cigar)
        cigarOps = cigarOps
    if not start:
        cigarOps.reverse()
    how_much_split = checkforSplitAtStart(cigarOps, start=start)
    map_length = (mappedLength(cigarOps))
    if map_length < 10000:
        return 
    if how_much_split is not None:
        if how_much_split < 500:
            return
        # TODO BLAST 
        chrom =  read_info.ref
        idx = 1 
        strand = read_info.isReverseStrand()
        seq_fasta = sequences_fasta.fetch(region=read_info.query)
        if (strand):
            strand = "-"
        else:
            strand = "+"
        event_id = "telomeric_sv"
        if start == True:
            br1 = 0
            br2 = read_info.pos
            start_read = 0
            end_read = how_much_split
            deletion_id = chrom + ":" + "1-" + str(br2)
            out_string = "\t".join(map(str,[chrom, br1,br1+1, chrom,br2,br2+1, read_info.query, strand, start_read, end_read, end_read - start_read, event_id, deletion_id])) 
            with open(os.path.join(telomere_folder,deletion_id + "_" + event_id + ".fasta"),"w") as fasta_out:
                header = ">" + deletion_id +"_"+event_id +"\n"
                fasta_out.write(header + seq_fasta[:how_much_split] +"\n")
            #read_info
            #print(out_string)
        else:
            br2 = "END" 
            # END of chromosome
            br1 = read_info.pos + map_length 
            start_read = br1
            end_read = how_much_split
            idx_end = len(seq_fasta) - how_much_split
            idx_end = idx_end 
            deletion_id = chrom + ":" + str(br1) +"-END" 
            out_string = "\t".join(map(str,[chrom, br1,int(br1)+1, chrom,br2,br2, strand, int(br1) + 1, int(br1) + how_much_split, how_much_split,  event_id, deletion_id])) 
            #with open(os.path.join(telomere_folder,deletion_id + "_" + event_id + ".fasta")) as fasta_out:
            with open(os.path.join(telomere_folder,deletion_id + "_" + event_id + ".fasta"),"w") as fasta_out:
                header = ">" + deletion_id +"_"+event_id +"\n"
                fasta_out.write(header + seq_fasta[(idx_end):] +"\n")
                #header = ">" + deletion_id +"_"+event_id +"\n"
                #fasta_out.write(read_info.seq[:how_much_split] +"\n")
            #print(out_string)
        out_f.write(out_string + "\n")
    #check_for_split()
    #query = calcQueryPosFromCigar(blockList[i].cigarOps)


def makeBedpe(blockList, opts):
    if (len(blockList) < 2):
        return
    # First put the "primary" alignments into the pBedBlock.
    # And any remaining "secondary" alignments into the sBedBlock.
    pBedBlock=[]
    sBedBlock=[]
    for i in xrange(len(blockList)):
        blockList[i].extractCigarOps()
        ref = calcRefPosFromCigar(blockList[i].cigarOps, blockList[i].pos)
        #print blockList[i].cigar
        #print ref.rPos
        query = calcQueryPosFromCigar(blockList[i].cigarOps)
        #print query.qePos
        #print query.qLen 
        #print query.qsPos
        strand = "+"
        if blockList[i].isReverseStrand() == True:
            strand = "-"
        qual = int(blockList[i].mapq)
        score = blockList[i].extractTagValue('AS');
        # seqname, startRefOff, endRefOff, queryId, strand, startQueryOff, endQueryOff, mapQual, AS, queryLen
        bed = [blockList[i].ref, blockList[i].pos - 1, ref.rPos - 1, blockList[i].query, strand, query.qsPos, query.qePos, qual, score, query.qLen, blockList[i].cigar, blockList[i].query]
        yahaStatus = blockList[i].extractTagValue('YS')
        # For non-YAHA input, assume all alignments are primary.
        if (yahaStatus is None or ((yahaStatus & 0x20) != 0)):
            pBedBlock.append([bed])
        else:
            sBedBlock.append(bed)
    # Sort the primaries by starting query offset.
    pBedBlock.sort(cmp=lambda x,y: cmp(x[0][5],y[0][5]))
    # Now place secondaries into the corresponding primary list.
    for i in xrange(len(sBedBlock)):
        maxOverlap = -1;
        maxJvalue = -1;
        for j in xrange(len(pBedBlock)):
            overlap = calcOverlap(sBedBlock[i], pBedBlock[j][0])
            if (overlap > maxOverlap):
                maxOverlap = overlap
                maxJvalue = j
        if (maxOverlap > 0):
            pBedBlock[maxJvalue].append(sBedBlock[i])
    # Now output all pairs from adjoining similarity sets.
    for i in xrange(len(pBedBlock)-1):
        set1 = pBedBlock[i]
        duplicate_alignments = []
        for j in xrange(i+1, len(pBedBlock)):
            set2 = pBedBlock[j];
            # We need to check for the "span" feature.
            # Let the adjacent ones goes through unmolested.
            p1 = set1[0];
            p2 = set2[0];
            # sys.stderr.write(str(p1[6]) + ' ' + str(p2[5]) + '\n');
            #print p2
            #print opts.span
            #print opts
            #print p1[10]
            #print p2[11]
            #print p1[11]
             #   print "Non-matching reference assembly" 
            if ((opts.span == 0 and j > i+1) or (opts.span > 0 and p1[6] + opts.span -1 < p2[5])):
                #   print p1
              #  print p2
               # print opts.span
                #print "breaking"
                break;

            for ii in xrange(len(set1)):
                for jj in xrange(len(set2)):
                    outputBedpe(set1[ii], set2[jj], opts)
    # Check for telomeric 3' deletion.

def calcRefPosFromCigar(cigarOps, alignmentStart):
    rPos = alignmentStart
    for cigar in cigarOps:
        if cigar.op == 'M':
            rPos += cigar.length
        elif cigar.op == 'D':
            rPos += cigar.length
        elif cigar.op == 'I':
            continue
        elif cigar.op == 'N':
            raise ValueError('Unexpected Cigar Operation')
    d = refPos(rPos)
    return d


def calcQueryPosFromCigar(cigarOps):
    qsPos = 0
    qePos = 0
    qLen  = 0
    # if first op is a H, need to shift start position 
    # the opPosition counter sees if the for loop is looking at the first index of the cigar object    
    opPosition = 0  
    for cigar in cigarOps:
        if opPosition == 0 and (cigar.op == 'H' or cigar.op == 'S'):
            qsPos += cigar.length
            qePos += cigar.length
            qLen  += cigar.length
        elif opPosition > 0 and (cigar.op == 'H' or cigar.op == 'S'):
            qLen  += cigar.length
        elif cigar.op == 'M' or cigar.op == 'I':
            qePos += cigar.length
            qLen  += cigar.length
        # elif cigar.op == 'D' or cigar.op == 'N':
            # Do nothing.
        opPosition += 1
    d = queryPos(qsPos, qePos, qLen);
    return d


def main():

    usage = """%prog -i <samFile or stdin>>

splitReadSamToBedpe

Author: Michael Lindberg, Aaron Quinlan & Ira Hall	
Description: reports split read mappings in a SAM file to bedpe; 
IMPORTANT NOTE: this replaces previous versions splitReadSamToBedPe and splitReadSamToBedPe_ih (Ira completed final version)
OUTPUT: in addition to standard bedpe format, 
col8=sum of alignment scores; 
col11=queryStart1; 
col12=queryEnd1;
col13=queryStart2; 
col14=queryEnd2; 
col15=minNonOverlap between the two alignments; 
col15=query length;
col16=mapping qualities and alignment scores

Modified by GGF on Jan 12, 2012 to handle secondary alignments.
Modified by GGF on Feb 10, 2012 to add minNonOverlap parameter.
Modified by GGF on March 27, 2012 to add span parameter.
    """

    parser = OptionParser(usage)
    parser.add_option("-i", dest="samFile", help="sam filename sorted by read id (not enforced), or standard input (-i stdin)", metavar="FILE")
    parser.add_option("-m", dest="MNO", help="minimum non-overlap to allow in output (default 25)", metavar="MNO", type="int", default=25);
    parser.add_option("-s", dest="span", help="maximum span to allow between reported pairs (default 0)", metavar="SPAN", type="int", default=0);
    parser.add_option("-p", dest="prefix", help="Sample name used for outputting auxilary files", metavar="PREFIX", default="telomeres") 
    (opts, args) = parser.parse_args()

    if opts.samFile is None:
        parser.print_help()
        print
    else:
        try:
            processSam(opts);
        except IOError as err:
            sys.stderr.write("Output " + str(totalLines) + " total lines.\n");
            sys.stderr.write("IOError " + str(err) + "\n");
            return
            
if __name__ == "__main__":
    sys.exit(main())
    (END)
