#!/usr/bin/env python

import os, subprocess, shlex, time
from multiprocessing import Pool, cpu_count
from subprocess import Popen, PIPE #debug only this should only need to be imported into run_pipe


def run_pipe(steps, outfile=None):
    #break this out into a recursive function
    #TODO:  capture stderr
    from subprocess import Popen, PIPE
    p = None
    p_next = None
    first_step_n = 1
    last_step_n = len(steps)
    for n,step in enumerate(steps, start=first_step_n):
        print "step %d: %s" %(n,step)
        if n == first_step_n:
            if n == last_step_n and outfile: #one-step pipeline with outfile
                with open(outfile, 'w') as fh:
                    print "one step shlex: %s to file: %s" %(shlex.split(step), outfile)
                    p = Popen(shlex.split(step), stdout=fh)
                break
            print "first step shlex to stdout: %s" %(shlex.split(step))
            p = Popen(shlex.split(step), stdout=PIPE)
            #need to close p.stdout here?
        elif n == last_step_n and outfile: #only treat the last step specially if you're sending std
            with open(outfile, 'w') as fh:
                print "last step shlex: %s to file: %s" %(shlex.split(step), outfile)
                p_last = Popen(shlex.split(step), stdin=p.stdout, stdout=fh)
                p.stdout.close()
                p = p_last
        else: #handles intermediate steps and, in the case of a pipe to stdout, the last step
            print "intermediate step %d shlex to stdout: %s" %(n,shlex.split(step))
            p_next = Popen(shlex.split(step), stdin=p.stdout, stdout=PIPE)
            p.stdout.close()
            p = p_next
    out,err = p.communicate()
    return out,err

def resolve_reference():
        # assume the reference file is the only .fa or .fna file
        return next((f for f in os.listdir(".") \
                                if f.endswith('.fa') or \
                                        f.endswith('.fna') or \
                                        f.endswith('.fa.gz') or \
                                        f.endswith('.fna.gz')), None)
class Bwa(object):
	#@dxpy.entry_point("process")
	def __init__(self, job_input ):
		# reads_file, reference_tar should be links to file objects.
		# reference_tar should be a tar of files generated by bwa index and
		# the tar should be uncompressed to avoid repeating the decompression.
                self.reads_file    = job_input['LocalFile'].init(job_input['reads_file'])
                self.reference_tar = job_input['LocalFile'].init(job_input['reference_tar'])
                self.bwa_aln_params= job_input['bwa_aln_params']
                self.bwa_version   = job_input['LocalFile'].init(job_input['bwa_version'])

		print "In process"

		if self.bwa_version == "0.7.7":
			self.bwa = "bwa0.7.7"
		elif self.bwa_version == "0.7.10":
			self.bwa = "bwa0.7.10"
		else:
			self.bwa = "bwa" #we don't have different versions on the bibs
			#bwa = "bwa0.7.7"
		print "Using bwa version %s" %(self.bwa_version)
                return None

        def process(self, File, downloader, uploader, linker):
		# Generate filename strings and download the files to the local filesystem
		reads_filename = File.describe(self.reads_file)['name']
		reads_basename = reads_filename
		# the order of this list is important.  It strips from the right inward, so
		# the expected right-most extensions should appear first (like .gz)
		for extension in ['.gz', '.fq', '.fastq', '.fa', '.fasta']:
			reads_basename = reads_basename.rstrip(extension)
		self.reads_file = downloader.download(self.reads_file.get_id(),reads_filename)

		reference_tar_filename = File.describe(self.reference_tar)['name']
		reference_tar_file = downloader.download(self.reference_tar.get_id(),reference_tar_filename)
		# extract the reference files from the tar
		if reference_tar_filename.endswith('.gz') or reference_tar_filename.endswith('.tgz'):
			tar_command = 'tar -xzv --no-same-owner --no-same-permissions -f %s' %(reference_tar_filename)
		else:
			tar_command = 'tar -xv --no-same-owner --no-same-permissions -f %s' %(reference_tar_filename)
		print "Unpacking %s" %(reference_tar_filename)
		print tar_command
		print subprocess.check_output(shlex.split(tar_command))
		reference_filename = resolve_reference()
		print "Using reference file: %s" %(reference_filename)

		print subprocess.check_output('ls -l', shell=True)

		#generate the suffix array index file
		sai_filename = '%s.sai' %(reads_basename)
		with open(sai_filename,'w') as sai_file:
			# Build the bwa command and call bwa
			bwa_command = "%s aln %s -t %d %s %s" \
				%(self.bwa, self.bwa_aln_params, cpu_count(), reference_filename, reads_filename)
			print bwa_command
			subprocess.check_call(shlex.split(bwa_command), stdout=sai_file)

		print subprocess.check_output('ls -l', shell=True)

		# Upload the output to the DNAnexus project
		print "Uploading %s" %(sai_filename)
		sai_dxfile = uploader.upload(sai_filename)
		process_output = { "output": linker.link(sai_dxfile) }
		print "Returning from process:"
		print process_output
		return process_output


	def postprocess(self, fn_input):

		print "In postprocess with:"

		if fn_input['samtools_version'] == "0.1.19":
			self.samtools = "/usr/local/bin/samtools-0.1.19/samtools"
		elif fn_input['samtools_version'] == "1.0":
			self.samtools = "/usr/local/bin/samtools-1.0/bin/samtools"
		else:
			self.samtools = "samtools" # no version necessary
#			samtools = "/usr/local/bin/samtools-0.1.19/samtools"

		if fn_input['bwa_version'] == "0.7.7":
			self.bwa = "bwa0.7.7"
		elif fn_input['bwa_version'] == "0.7.10":
			self.bwa = "bwa0.7.10"
		else:
			self.bwa = "bwa" #as above
#			print "BWA version %s not supported, defaulting to 0.7.7"
#			bwa = "bwa0.7.7"

		print "samtools version: %s" %(self.samtools)
		print "bwa version %s" %(self.bwa)

		indexed_reads_filenames = []
		unmapped_reads_filenames = []
		for i,reads in enumerate(fn_input['indexed_reads']):
			read_pair_number = i+1

			tmp_reads=fn_input['LocalFile'](reads,"./")

			fn = fn_input['LocalFile'].describe(tmp_reads)['name']
			print "indexed_reads %d: %s" %(read_pair_number, fn)
			indexed_reads_filenames.append(fn)
			#fn_input['LocalDownloader'].download(tmp_reads.get_id(),fn)

			unmapped = fn_input['LocalFile'](fn_input['unmapped_reads'][i],"./")
			fn = fn_input['LocalFile'].describe(unmapped)['name']
			print "unmapped reads %d: %s" %(read_pair_number, fn)
			unmapped_reads_filenames.append(fn)
			#fn_input['LocalDownloader'].download(unmapped.get_id(),fn)

                tmp_reference=fn_input['LocalFile'](fn_input['reference_tar'],"./")
		reference_tar_filename = fn_input['LocalFile'].describe(tmp_reference)['name']
		print "reference_tar: %s" %(reference_tar_filename)
		#fn_input['LocalDownloader'].download(tmp_reference.get_id(), reference_tar_filename)
		# extract the reference files from the tar
		if reference_tar_filename.endswith('.gz') or reference_tar_filename.endswith('.tgz'):
			tar_command = 'tar -xzv --no-same-owner --no-same-permissions -f %s' %(reference_tar_filename)
		else:
			tar_command = 'tar -xv --no-same-owner --no-same-permissions -f %s' %(reference_tar_filename)
		print "Unpacking %s" %(reference_tar_filename)
		print tar_command
		print subprocess.check_output(shlex.split(tar_command))
		reference_filename = resolve_reference()

		paired_end = len(fn_input['indexed_reads']) == 2

		if paired_end:
			r1_basename = unmapped_reads_filenames[0].rstrip('.gz').rstrip('.fq').rstrip('.fastq')
			r2_basename = unmapped_reads_filenames[1].rstrip('.gz').rstrip('.fq').rstrip('.fastq')
			reads_basename = r1_basename + r2_basename
		else:
			reads_basename = unmapped_reads_filenames[0].rstrip('.gz').rstrip('.fq').rstrip('.fastq')
		raw_bam_filename = '%s.raw.srt.bam' %(reads_basename)
		raw_bam_mapstats_filename = '%s.raw.srt.bam.flagstat.qc' %(reads_basename)

		if paired_end:
			reads1_filename = indexed_reads_filenames[0]
			reads2_filename = indexed_reads_filenames[1]
			unmapped_reads1_filename = unmapped_reads_filenames[0]
			unmapped_reads2_filename = unmapped_reads_filenames[1]
			raw_sam_filename = reads_basename + ".raw.sam"
			badcigar_filename = "badreads.tmp"
			steps = [ "%s sampe -P %s %s %s %s %s" %(bwa, reference_filename, reads1_filename, reads2_filename, unmapped_reads1_filename, unmapped_reads2_filename),
					  "tee %s" %(raw_sam_filename),
					  r"""awk 'BEGIN {FS="\t" ; OFS="\t"} ! /^@/ && $6!="*" { cigar=$6; gsub("[0-9]+D","",cigar); n = split(cigar,vals,"[A-Z]"); s = 0; for (i=1;i<=n;i++) s=s+vals[i]; seqlen=length($10) ; if (s!=seqlen) print $1"\t" ; }'""",
					  "sort",
					  "uniq" ]
			out,err = run_pipe(steps,badcigar_filename)
			if err:
				print "sampe error: %s" %(err)

			steps = [ "cat %s" %(raw_sam_filename),
					  "grep -v -F -f %s" %(badcigar_filename)]
		else: #single end
			reads_filename = indexed_reads_filenames[0]
			unmapped_reads_filename = unmapped_reads_filenames[0]
			steps = [ "%s samse %s %s %s" %(self.bwa, reference_filename, reads_filename, unmapped_reads_filename) ]
		if fn_input['samtools_version'] == "0.1.9":
			steps.extend(["%s view -Su -" %(self.samtools),
						  "%s sort - %s" %(self.samtools, raw_bam_filename.rstrip('.bam')) ]) # samtools adds .bam
		else:
			steps.extend(["%s view -@%d -Su -" %(self.samtools, cpu_count()),
						  "%s sort -@%d - %s" %(self.samtools, cpu_count(), raw_bam_filename.rstrip('.bam')) ]) # samtools adds .bam
		print "Running pipe:"
		print steps
		out,err = run_pipe(steps)

		if out:
			print "samtools output: %s" %(out)
		if err:
			print "samtools error: %s" %(err)

		with open(raw_bam_mapstats_filename, 'w') as fh:
			subprocess.check_call(shlex.split("%s flagstat %s" \
				%(self.samtools, raw_bam_filename)), stdout=fh)

		print subprocess.check_output('ls', shell=True)
		mapped_reads = fn_input['LocalUploader'].upload(raw_bam_filename)
		mapping_statistics = fn_input['LocalUploader'].upload(raw_bam_mapstats_filename)

		output = { "mapped_reads": fn_input['LocalLinker'].link(mapped_reads),
				   "mapping_statistics": fn_input['LocalLinker'].link(mapping_statistics) }
		print "Returning from post with output: %s" %(output)
		return output

