#!/usr/bin/env python

import subprocess
import tempfile
import shutil
import os
import re

import fasta_statter
from Bio import SeqIO


description = "Mimicking Intelligent Read Assembly - a multipass OLC assembler with advanced finishing tools."

core_load = 2 #number of cores this assembler will max out

supports = ('MiSeq', 'IonTorrent')

manifests = {'MiSeq':"""
#MIRA manifest file
#Autogenerated by Assembly Dispatch:mira_assembler
#For {accession}			

project = {accession}
job = denovo,genome,draft
parameters = -NW:somrnl=0

readgroup = {accession}_miseq
data = {reads1} {reads2}
technology = solexa
template_size = 200 650 infoonly
segment_placement = ---> <--- 

""",
			 'IonTorrent':"""
#MIRA manifest file
#Autogenerated by Assembly Dispatch:mira_assembler
#For {accession}

project = {accession}
job = denovo,genome,draft
parameters = -NW:somrnl=0

readgroup = {accession}_iontorrent
data = {reads1}
technology = iontor
template_size = 200 650 infoonly
segment_placement = ?
"""}

def assemble(reads1, data_type='MiSeq', reads2=None, fasta_file_name=None, callback=lambda s: None, update_callback=lambda d: None, debug=True, **kwargs):

	old_dir=os.getcwd()

	d = {'assembly_version':'',
		 'average_coverage':'',
		 'num_contigs':'',
		 'n50':'',
		 'num_bases':'',
		 'fasta_file':'',
		 'lib_insert_length':'Not determined',
		 'matched':''
		 }
	kwargs['temp_dir'] = temp_dir = tempfile.mkdtemp()
	if not fasta_file_name:
		kwargs['fasta_file_name'] = fasta_file_name = "{accession}.mira.fasta".format(**kwargs)
	else:
		kwargs['fasta_file_name'] = fasta_file_name
	#callback("Concatenating reads.")
	#subprocess.check_call("cat {reads1} {reads2} > {temp_dir}/{accession}_in.fastq".format(**kwargs), shell=True)
	results = ''
	try:
		#get version
		d['assembly_version'] = "MIRA v. {}".format(
			subprocess.check_output("mira --version", shell=True).replace(" ()\n", ""))
		callback("Starting " + d['assembly_version'])


		#convert reads to older quality scores
		for (ri, ro) in ((reads1, 1), (reads2, 2)):
			try:
				with open("{}/reads{}.fastq".format(temp_dir, ro), 'w') as reads_out, open(ri, 'r') as reads_in:
					callback("Converting reads {}...".format(ro))
					for r in SeqIO.parse(reads_in, 'fastq'):
						reads_out.write(r.format('fastq'))
				kwargs['reads{}'.format(ro)] = "{}/reads{}.fastq".format(temp_dir, ro)
			except Exception:
				pass

		#write manifest file
		callback("Writing manifest")
		with open(os.path.join(temp_dir, "manifest.txt"), 'w') as manifest:
			manifest.write(manifests[data_type].format(**kwargs))

		
		#assemble
		callback("Running MIRA ({})".format(temp_dir))
		os.chdir(temp_dir)
		results = subprocess.check_output("mira manifest.txt", shell=True)
		
		callback("Statting...")
		d.update(fasta_statter.stat_fasta("{temp_dir}/{accession}_assembly/{accession}_d_results/{accession}_out.unpadded.fasta".format(**kwargs)))
		shutil.copyfile("{temp_dir}/{accession}_assembly/{accession}_d_results/{accession}_out.unpadded.fasta".format(**kwargs), "{path}/{fasta_file_name}".format(**kwargs))
		d['fasta_file'] = kwargs['fasta_file_name']
		
	
	#d.update(fasta_statter.stat_fasta(...))
	except subprocess.CalledProcessError as e:
		if debug:
			with open(os.path.join(temp_dir, "manifest.txt"), 'r') as f:
				for l in f:
					print l,
			print results
			print type(e), e, e.output
			import traceback
			import sys
			traceback.print_exc(sys.stdout)
		raise ValueError("MIRA assembly failed.")
	
	finally:
		callback("Cleaning up " + temp_dir)
		os.chdir(old_dir)
		shutil.rmtree(temp_dir)
		
		
	return d
	
	
if __name__ == "__main__":
	#debug
	import datetime
	def cb(d):
		print "[{}] {}".format(datetime.datetime.today().ctime(), d)
	def bcb(d):
		for (k, v) in d.items():
			cb("{} : {}".format(k, v))
	bcb(assemble(path='/home/justin.payne', 
			 reads1='/shared/gn2/CFSANgenomes/CFSAN003966/CFSAN003966_01/CFSAN003966_S6_L001_R1_001.fastq', 
			 reads2='/shared/gn2/CFSANgenomes/CFSAN003966/CFSAN003966_01/CFSAN003966_S6_L001_R2_001.fastq', 
			 accession='CFSAN003966_01', 
			 callback=cb,
			 update_callback=bcb,
			 debug=False))
	bcb(assemble(path='/home/justin.payne', 
			 reads1='/shared/gn2/CFSANgenomes/CFSAN003966/CFSAN003966_01/CFSAN003966_S6_L001_R1_001.fastq', 
			 #reads2='/shared/gn2/CFSANgenomes/CFSAN003966/CFSAN003966_01/CFSAN003966_S6_L001_R2_001.fastq', 
			 accession='CFSAN003966_01', 
			 data_type='IonTorrent',
			 callback=cb,
			 update_callback=bcb,
			 debug=True))