#!/usr/bin/env python

# Copyright (c) 2016 Ryan Collins <rcollins@chgr.mgh.harvard.edu>
# Distributed under terms of the MIT license.

"""
Estimate original molecule sizes and coordinates from 10X linked-read WGS barcodes
"""

import argparse
from collections import defaultdict, Counter, namedtuple
from subprocess import call
import pysam

def get_gemcode_regions(ibam, dist):
    """
    Estimates molecule coordinates from colinear reads with overlapping 10X barcodes

    Parameters
    ----------
    ibam : pysam.AlignmentFile
        Input 10X bam
    dist : int
        Partitioning distance (bp) for separating reads with overlapping
        barcodes into independent fragments

    Yields
    ------
    region : namedtuple
        chr, start, end, barcode, readcount
    """

    #Create namedtuples for storing read coordinate and molecule info
    coords = namedtuple('coords', ['chr', 'pos'])
    molecule = namedtuple('molecule', ['chr', 'start', 'end', 'barcode', 'readcount'])

    #Create defaultdict for storing gemcode tuples
    #Key is gemcodes, value is coords namedtuple
    gemcodes = defaultdict(list)

    #Set current_chr as reference contig of first read
    current_chr = next(ibam).reference_name
    ibam.seek(0)

    #Iterate over reads in bamfile
    for read in ibam:
        #Save 10X barcode as gem
        gem = read.get_tag('RX')
        #If the read is from a new contig, write out all molecules held in memory
        #then add read to emptied dictionary
        if read.reference_name != current_chr:

            for gem in gemcodes:
            
                yield molecule(gemcodes[gem][0].chr, min([pos for chr,
                               pos in gemcodes[gem]]), 
                               max([pos for chr, pos in gemcodes[gem]]),
                               gem, len(gemcodes[gem]))

            gemcodes = defaultdict(list)
            
            gemcodes[gem].append(coords(read.reference_name, read.reference_start))
        
        #If barcode has been seen previously and new read is colinear but beyond 
        #dist, yield old barcode as interval before adding new read to list
        elif gem in gemcodes and read.reference_start - gemcodes[gem][-1].pos > dist:

            yield molecule(gemcodes[gem][0].chr, 
                           min([pos for chr, pos in gemcodes[gem]]), 
                           max([pos for chr, pos in gemcodes[gem]]), 
                           gem, len(gemcodes[gem]))

            gemcodes[gem] = [coords(read.reference_name, read.reference_start)]

        else:
            #Else just add read to preexisting dictionary 
            gemcodes[gem].append(coords(read.reference_name, read.reference_start))

        #Save read contig as current_chr
        current_chr = read.reference_name

    #Write out all remaining molecules at end of bam
    for gem in gemcodes:

        yield molecule(gemcodes[gem][0].chr, min([pos for chr, pos in gemcodes[gem]]), 
                       max([pos for chr, pos in gemcodes[gem]]), gem, len(gemcodes[gem]))

#Run function
def main():
    #Add arguments
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('ibam', type=pysam.AlignmentFile,
                        help='Input bam')
    parser.add_argument('outfile', help='Output bed file')
    parser.add_argument('-d', '--dist', type=int, default=50000,
                        help='Molecule partitioning distance in bp (default: 50000)')
    args = parser.parse_args()

    #Open outfile
    fout = open(args.outfile, 'w')

    #Get gemcode regions
    for bed in get_gemcode_regions(args.ibam, args.dist):

        #Turn molecule object into string
        bed_str = '{0}\t{1}\t{2}\t{3}\t{4}'.format(bed.chr, bed.start, bed.end,
                                                   bed.barcode, bed.readcount)

        #Write to file
        fout.write(bed_str + '\n')

    #Close outfile
    fout.close()

    #Sort outfile by chr then start (system call)
    call('sort -Vk1,1 -k2,2n -o ' + args.outfile + ' ' + args.outfile, shell=True)


#Main block
if __name__ == '__main__':
    main()
@@ -1538,16 +1538,28 @@
 nt_chr:%0A
+            
 %0A       
@@ -1559,35 +1559,39 @@
             for 
-gem
+barcode
  in gemcodes:%0A  
@@ -1588,28 +1588,16 @@
 mcodes:%0A
-            
 %0A       
@@ -1621,35 +1621,39 @@
 lecule(gemcodes%5B
-gem
+barcode
 %5D%5B0%5D.chr, min(%5Bp
@@ -1703,35 +1703,39 @@
 pos in gemcodes%5B
-gem
+barcode
 %5D%5D), %0A          
@@ -1785,27 +1785,31 @@
 in gemcodes%5B
-gem
+barcode
 %5D%5D),%0A       
@@ -1824,35 +1824,39 @@
                 
-gem
+barcode
 , len(gemcodes%5Bg
@@ -1846,35 +1846,39 @@
 e, len(gemcodes%5B
-gem
+barcode
 %5D))%0A%0A           
@@ -2885,19 +2885,23 @@
     for 
-gem
+barcode
  in gemc
@@ -2931,35 +2931,39 @@
 lecule(gemcodes%5B
-gem
+barcode
 %5D%5B0%5D.chr, min(%5Bp
@@ -2965,32 +2965,55 @@
 in(%5Bpos for chr,
+%0A                      
  pos in gemcodes
@@ -3005,35 +3005,39 @@
 pos in gemcodes%5B
-gem
+barcode
 %5D%5D), %0A          
@@ -3087,19 +3087,50 @@
 des%5B
-gem%5D%5D), gem
+barcode%5D%5D),%0A                       barcode
 , le
@@ -3140,24 +3140,27 @@
 emcodes%5B
-gem
+barcode
 %5D))%0A
-%0A
 #Run fun

Fix variable collision