##################################
#                                #
# Last modified 09/18/2013       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import os
import pysam

def run():

    if len(sys.argv) < 5:
        print 'usage: python %s samtools chrom.sizes BAM1 BAM2 outfileprefix' % sys.argv[0]
        print '\tNote: use the phase parameter specifies which line to print out, i.e if you want every odd line, enter 1, if you want '
        print '\tNote: use this script will files that contain UNIQUE ALIGNMENTS ONLY!!!'
        print '\tNote: the script will produce unsorted BAM files, which it will then sort and delete, and index the sorted files'
        sys.exit(1)

    samtools = sys.argv[1]
    BAM1 = sys.argv[3]
    BAM2 = sys.argv[4]
    outprefix = sys.argv[5]
    chrominfo = sys.argv[2]
    chromInfoList=[]
    linelist=open(chrominfo)
    for line in linelist:
        fields=line.strip().split('\t')
        chr=fields[0]
        start=0
        end=int(fields[1])
        chromInfoList.append((chr,start,end))

    i=0
    samfile = pysam.Samfile(BAM1, "rb" )
    outfile = pysam.Samfile(outprefix + '.unsorted1.bam', "wb", template=samfile)
    for (chr,start,end) in chromInfoList:
        try:
            for alignedread in samfile.fetch(chr, 0, 100):
                a='b'
        except:
            continue
        for alignedread in samfile.fetch(chr, start, end):
            i+=1
            if i % 5000000 == 0:
                print str(i/1000000) + 'M alignments processed processed', chr,start,alignedread.pos,end
            if i % 2 == 1:
                outfile.write(alignedread)
    samfile = pysam.Samfile(BAM2, "rb" )
    for (chr,start,end) in chromInfoList:
        try:
            for alignedread in samfile.fetch(chr, 0, 100):
                a='b'
        except:
            continue
        for alignedread in samfile.fetch(chr, start, end):
            i+=1
            if i % 5000000 == 0:
                print str(i/1000000) + 'M alignments processed processed', chr,start,alignedread.pos,end
            if i % 2 == 1:
                outfile.write(alignedread)

    outfile.close()

    i=0
    samfile = pysam.Samfile(BAM1, "rb" )
    outfile = pysam.Samfile(outprefix + '.unsorted2.bam', "wb", template=samfile)
    for (chr,start,end) in chromInfoList:
        try:
            for alignedread in samfile.fetch(chr, 0, 100):
                a='b'
        except:
            continue
        for alignedread in samfile.fetch(chr, start, end):
            i+=1
            if i % 5000000 == 0:
                print str(i/1000000) + 'M alignments processed processed', chr,start,alignedread.pos,end
            if i % 2 == 0:
                outfile.write(alignedread)
    samfile = pysam.Samfile(BAM2, "rb" )
    for (chr,start,end) in chromInfoList:
        try:
            for alignedread in samfile.fetch(chr, 0, 100):
                a='b'
        except:
            continue
        for alignedread in samfile.fetch(chr, start, end):
            i+=1
            if i % 5000000 == 0:
                print str(i/1000000) + 'M alignments processed processed', chr,start,alignedread.pos,end
            if i % 2 == 0:
                outfile.write(alignedread)

    outfile.close()

    cmd = samtools + ' sort ' + outprefix + '.unsorted2.bam ' + outprefix + '.pseudorep2'
    os.system(cmd)
    cmd = samtools + ' sort ' + outprefix + '.unsorted1.bam ' + outprefix + '.pseudorep1'
    os.system(cmd)
    cmd = 'rm ' + outprefix + '.unsorted2.bam'
    os.system(cmd)
    cmd = 'rm ' + outprefix + '.unsorted1.bam'
    os.system(cmd)
    cmd = samtools + ' index ' + outprefix + '.pseudorep2.bam'
    os.system(cmd)
    cmd = samtools + ' index ' + outprefix + '.pseudorep1.bam'            
    os.system(cmd)

run()