import os
import sys
import glob
import pysam
import argparse
import multiprocessing

def get_args():
    '''Parse sys.argv'''
    parser = argparse.ArgumentParser()
    parser.add_argument('--cores', type=int,
                        help='the number of avalible processor cores')
    parser.add_argument('-i','--input-dir', required=True,
                        help='The input directory containing the bam files.')
    parser.add_argument('-CN',type=str,
                        help="Name of sequencing center producing the read. \
                        GATK not required.")
    parser.add_argument('-DS',type=str,
                        help="Description. GATK Not Required. ")
    parser.add_argument('-DT',type=str, required=True,
                        help="Date the run was produced (ISO8601 date or date/time). \
                        GATK Not Required. ")
    parser.add_argument('-PI',type=int, required=True,
                        help="Predicted median insert size. GATK Not Required.")
    parser.add_argument('-PL',type=str, required=True,
                        choices = ['CAPILLARY', 'LS454', 'ILLUMINA', 'SOLID', 
                                    'HELICOS', 'IONTORRENT', 'PACBIO'],
                        help="Platform/technology used to produce the reads.")

    args = parser.parse_args()
    return args

def addRG2Header(filename, sam_info, args):
    """Add read group info to a header."""
    # CREATE TEMPLATE
    # Read group. Unordered multiple @RG lines are allowed.
    RG_template = { 'ID': '',           # Read group identifier. e.g., Illumina flowcell + lane name and number
                    'CN': '',           # GATK Not Required. Name of sequencing center producing the read.
                    'DS': '',           # GATK Not Required. Description
                    'DT': '',           # GATK Not Required. Date the run was produced (ISO8601 date YYYY-MM-DD or YYYYMMDD)
                    'PI': '',           # GATK Not Required. Predicted median insert size.
                    'PU': '',           # GATK Not Required. Platform unit (e.g. flowcell-barcode.lane for Illumina or slide for SOLiD).
                    'SM': '',           # Sample. Use pool name where a pool is being sequenced.
                    'PL': 'ILLUMINA'}   # Platform/technology used to produce the reads.

    samfile = pysam.Samfile(filename, 'r')
    new_header = samfile.header.copy()
    samfile.close()

    # ADD INFO TO TEMPLATE
    RG_template = RG_template.copy()
    RG_template['ID'] = sam_info["sample_name"]
    if args.CN: RG_template['CN'] = args.CN.upper()
    if args.DS: RG_template['DS'] = args.DS
    RG_template['DT'] = args.DT
    RG_template['LB'] = sam_info["sample_name"]
    RG_template['SM'] = sam_info["sample_name"]
    RG_template['DS'] = "{0}.{1}".format(sam_info['sample_name'], sam_info['locality'])
    RG_template['PI'] = args.PI
    RG_template['PU'] = '{0}.{1}'.format(sam_info['flowcell_id'], sam_info['lane'])
    new_header['RG'] = [RG_template]
    return new_header


def add_RGs_2_BAMs_runner(data):
    """Generates the correct @RG header and adds a RG field to a bam file."""
    # Make Bam of Sam if it doesn't exist
    filename, new_RG_header = data
    if filename.endswith('sam'):
        sam_handle = pysam.Samfile(filename)
        bam_name = os.path.splitext(filename)[0]+".bam"
        bam_handle = pysam.Samfile( bam_name, "wb", template = sam_handle )
        for s in sam_handle:
            bam_handle.write(s)
        filename = bam_name

    # Massage paths and make outputfiles
    pysam.sort(filename, os.path.splitext(filename)[0]+".sorted")
    pysam.index(os.path.splitext(filename)[0]+".sorted.bam")
    filename = os.path.splitext(filename)[0]+".sorted.bam"
    path, filename = os.path.split(filename)
    name, ext = os.path.splitext(filename)
    new_name = name + '.wRG.' + 'bam'
    outfile_name =  os.path.join(path,new_name)
    outfile = pysam.Samfile( outfile_name, 'wb', header = new_RG_header )

    # Step 2: Process Samfile adding Read Group to Each Read
    samfile = pysam.Samfile(os.path.join(path, filename))
    samfile.fetch()
    for count, read in enumerate(samfile.fetch()):
        name = read.qname
        read_group = os.path.split(filename)[1].split(".")[0]
        new_tags = read.tags
        new_tags.append(('RG', read_group))
        read.tags = new_tags
        outfile.write(read)
    outfile.close()

    # Step 3: Make index of read group enabled samfile
    pysam.index(outfile_name)
    sys.stdout.write(".")
    sys.stdout.flush()
    return

def parseFileName(filepath):
    path, filename = os.path.split(filepath)
    sample_name, locality, inline_tag, third_read_tag = filename.split(".")[:-1]
    sam_handle = pysam.Samfile(filepath,'r')
    sam_line = sam_handle.next()
    read_info = sam_line.qname
    sam_handle.close()
    instrument, run_id, flowcell_id, lane = read_info.split(":")[:4]
    info = {'sample_name': sample_name,
            'locality': locality,
            'inline_tag': inline_tag,
            'third_read_tag': third_read_tag,
            'instrument':instrument,  
            'run_id': run_id,
            'flowcell_id': flowcell_id,
            'lane': lane}
    return info

def add_RGs_2_BAMs(pool, args):
    data_for_map = []
    print 'Making RG headers.'
    for count, filename in enumerate(glob.glob(os.path.join(args.input_dir,'*'))):
        if filename.endswith('sam') or filename.endswith('bam'):
            sam_info = parseFileName(filename)
            new_RG_header = addRG2Header(filename, sam_info, args)
            data_for_map.append([filename, new_RG_header])

    sys.stdout.write("\nAdding RGs and making BAMs")
    sys.stdout.flush()
    pool.map(add_RGs_2_BAMs_runner, data_for_map)
    return

def main():
    args = get_args()
    cores = args.cores
    pool = multiprocessing.Pool(args.cores)
    add_RGs_2_BAMs(pool, args)

if __name__ == '__main__':
    main()

#!/usr/bin/env python
# encoding: utf-8

import multiprocessing
from textwrap import dedent
from itertools import izip_longest

def process_chunk(d):
	"""Replace this with your own function
	that processes data one line at a
	time"""

	d = d.strip() + ' processed'
	return d 

def grouper(n, iterable, padvalue=None):
	"""grouper(3, 'abcdefg', 'x') -->
	('a','b','c'), ('d','e','f'), ('g','x','x')"""

	return izip_longest(*[iter(iterable)]*n, fillvalue=padvalue)

if __name__ == '__main__':

	# test data
	test_data = """\
	1 some test garbage
	2 some test garbage
	3 some test garbage
	4 some test garbage
	5 some test garbage
	6 some test garbage
	7 some test garbage
	8 some test garbage
	9 some test garbage
	10 some test garbage
	11 some test garbage
	12 some test garbage
	13 some test garbage
	14 some test garbage
	15 some test garbage
	16 some test garbage
	17 some test garbage
	18 some test garbage
	19 some test garbage
	20 some test garbage"""
	test_data = dedent(test_data)
	test_data = test_data.split("\n")

	# Create pool (p)
	p = multiprocessing.Pool(4)

	# Use 'grouper' to split test data into
	# groups you can process without using a
	# ton of RAM. You'll probably want to 
	# increase the chunk size considerably
	# to something like 1000 lines per core.

	# The idea is that you replace 'test_data'
	# with a file-handle
	# e.g., testdata = open(file.txt,'rU')

	# And, you'd write to a file instead of
	# printing to the stout

	for chunk in grouper(10, test_data):
		results = p.map(process_chunk, chunk)
		for r in results:
			print r 	# replace with outfile.write()

#!/usr/bin/env python
# encoding: utf-8

import gzip
import argparse

def interface():
	args = argparse.ArgumentParser()
	args.add_argument('-l', '--left-input',
		help='The first input file in fastq format.')

	args.add_argument('-r', '--right-input', 
		help='The first input file in fastq format.')

	args.add_argument('-o', '--output', 
		help='The output file in fastq format.')

	args = args.parse_args()
	return args


def process_reads(args):

	if args.left_input.endswith('.gz') or args.right_input.endswith('.gz'):

		left = gzip.open(args.left_input,'rb')
		right = gzip.open(args.right_input,'rb')
		fout = gzip.open(args.output,'wb')

	else:
		left = open(args.left_input,'rU')
		right = open(args.right_input,'rU')
		fout = open(args.output,'wb')


	# USING A WHILE LOOP MAKE THIS SUPER FAST
	# Details here: 
	#   http://effbot.org/zone/readline-performance.htm

	while 1: 

		# process the first file
		left_line = left.readline()
		if not left_line: break
		fout.write(left_line)

		left_line = left.readline()
		fout.write(left_line)

		left_line = left.readline()
		fout.write(left_line)

		left_line = left.readline()
		fout.write(left_line)


		# process the second file
		right_line = right.readline()
		fout.write(right_line)

		right_line = right.readline()
		fout.write(right_line)

		right_line = right.readline()
		fout.write(right_line)

		right_line = right.readline()
		fout.write(right_line)

	left.close()
	right.close()
	fout.close()
	return 0

if __name__ == '__main__':
	args = interface()
	process_reads(args)