##################################
#                                #
# Last modified 2017/03/31       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import os
import random
import pysam
import string

def run():

    if len(sys.argv) < 1:
        print 'usage: python %s BAM' % sys.argv[0]
        sys.exit(1)

    BAM = sys.argv[1]

    samfile = pysam.Samfile(BAM, "rb" )

    outFileDict = {}

    for alignedread in samfile.fetch(until_eof=True):
#        fields = str(alignedread).split('\t')
#        cellID = fields[0].split(':::')[0]
        cellID = alignedread.query_name.split(':::')[0]
        alignedread.query_name = alignedread.query_name.split(':::')[1]
        if outFileDict.has_key(cellID):
            pass
        else:
            outFileDict[cellID] = pysam.Samfile(BAM.split('.bam')[0] + '.' + cellID + '.bam', "wb", template=samfile)
        outFileDict[cellID].write(alignedread)

    for cellID in outFileDict.keys():
        outFileDict[cellID].close()

run()