#!/bin/env python

import sys,pysam,argparse

# TODO: add readgroup mongering

# fixes unmapped reads that are marked as 'reverse'
def cleanup(read):
    if read.is_unmapped and read.is_reverse:
        read.is_reverse = False
    return read

def main(args):
    targetbam = pysam.Samfile(args.targetbam, 'rb')
    donorbam  = pysam.Samfile(args.donorbam, 'rb')
    outputbam = pysam.Samfile(args.outputbam, 'wb', template=targetbam)

    namechange = None

    if args.namechange:
        namechange = args.namechange

    # load reads from donorbam into dict 
    sys.stderr.write("loading donor reads into dictionary...\n")
    nr = 0
    rdict = {}
    for read in donorbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        rdict[extqname] = read
        nr += 1

    sys.stderr.write("loaded " + str(nr) + " reads, reading " + args.targetbam + "...\n")


    used = {}
    for read in targetbam.fetch(until_eof=True):
        pairname = 'F' # read is first in pair
        if read.is_read2:
            pairname = 'S' # read is second in pair
        if not read.is_paired:
            pairname = 'U' # read is unpaired

        extqname = ','.join((read.qname,pairname))
        if extqname in rdict: 
            if args.keepqual:
                if namechange: # must set name _before_ setting quality (see pysam docs)
                    rdict[extqname].qname = args.namechange + rdict[extqname].qname    
                rdict[extqname].qual = read.qual
            rdict[extqname] = cleanup(rdict[extqname])
            outputbam.write(rdict[extqname])  # write read from donor .bam
            used[extqname] = True
        else:
            if namechange:
                qual = read.qual # temp
                read.qname = args.namechange + read.qname
                read.qual = qual
            read = cleanup(read)
            outputbam.write(read) # write read from target .bam

    nadded = 0
    # dump the unused reads from the donor if requested with --all
    if args.all:
        for extqname in rdict.keys():
            if extqname not in used:
                if namechange:
                    qual = rdict[extqname].qual # temp
                    rdict[extqname].qname = args.namechange + rdict[extqname].qname
                    rdict[extqname].qual = qual
                rdict[extqname] = cleanup(rdict[extqname])
                outputbam.write(rdict[extqname])
                nadded += 1
        sys.stderr.write("added " + str(nadded) + " reads due to --all\n")

    targetbam.close()
    donorbam.close()
    outputbam.close()

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='replaces aligned reads in bamfile1 with aligned reads from bamfile2')
    parser.add_argument('-b', '--bam', dest='targetbam', required=True,
                        help='original .bam')
    parser.add_argument('-r', '--replacebam', dest='donorbam', required=True,
                        help='.bam with reads to replace original bam')
    parser.add_argument('-o', '--outputbam', dest='outputbam', required=True)
    parser.add_argument('-i', '--ignoresize', dest='ignoresize', default=0, help="don't replace reads with reads having insert size > ignoresize")
    parser.add_argument('-n', '--namechange', dest='namechange', default=None, help="change all read names by prepending string (passed as -n [string])")
    parser.add_argument('--all', action='store_true', default=False, help="append reads that don't match target .bam")
    parser.add_argument('--keepqual', action='store_true', default=False, help="keep original quality scores, replace read and mapping only")
    args = parser.parse_args()
    main(args)

@@ -37,16 +37,43 @@
 argparse
+%0Afrom random import randint
 %0A%0A# TODO
@@ -99,17 +99,49 @@
 gering%0A%0A
-#
+def cleanup(read,RG):%0A    '''%0A   
  fixes u
@@ -187,26 +187,110 @@
 se'%0A
-def cleanup(read):
+    fill in read group at random from existing RGs if %0A    RG tags are present in .bam header %0A    '''
 %0A   
@@ -367,23 +367,501 @@
 lse%0A
+%0A
     
-return read
+if RG:%0A        hasRG = False%0A        for tag in read.tags:%0A            if tag%5B0%5D == 'RG':%0A                hasRG = True%0A%0A        if not hasRG:%0A            # add random read group from list in header%0A            newRG = RG%5Brandint(0,len(RG)-1)%5D%0A            read.tags = read.tags + %5B(%22RG%22,newRG)%5D%0A    return read%0A%0Adef getRGs(bam):%0A    '''return list of RG IDs'''%0A    RG = %5B%5D%0A    if 'RG' in bam.header:%0A        for headRG in bam.header%5B'RG'%5D:%0A            RG.append(headRG%5B'ID'%5D)%0A    return RG
 %0A%0Ade
@@ -1050,16 +1050,58 @@
 etbam)%0A%0A
+    RG = getRGs(targetbam) # read groups%0A%0A
     name
@@ -2401,32 +2401,35 @@
 (rdict%5Bextqname%5D
+,RG
 )%0A            ou
@@ -2726,16 +2726,19 @@
 nup(read
+,RG
 )%0A      
@@ -3237,32 +3237,35 @@
 (rdict%5Bextqname%5D
+,RG
 )%0A              

read group fill-in