#!/usr/bin/python
# filename: ab_analysis.py

###########################################################################
#
# Copyright (c) 2013 Bryan Briney.  All rights reserved.
#
# @version: 1.0.0
# @author: Bryan Briney
# @props: IgBLAST team (http://www.ncbi.nlm.nih.gov/igblast/igblast.cgi)
# @license: MIT (http://opensource.org/licenses/MIT) 
#
###########################################################################


import os
import time
import math
import glob
import platform
import argparse
import threading
from subprocess import Popen, PIPE
from multiprocessing import Pool, cpu_count

from Bio import SeqIO

from blast_parse import BlastParse


parser = argparse.ArgumentParser("Antibody annotation with IgBLAST.")
parser.add_argument('-i', '--in', dest='input', required=True, 
					help="The input file, to be split and processed in parallel. \
					If a directory is given, all files in the directory will be iteratively processed.")
parser.add_argument('-o', '--out', dest='output', required=True, 
					help="The output directory, which will contain JSON or tab-delimited output files.")
parser.add_argument('-l', '--log', dest='log', default='', 
					help="The log file, to which the BlastParse log info will be written. \
					Default is stdout.")
parser.add_argument('-t', '--temp', dest='temp_dir', default='', 
					help="The directory in which temp files will be stored.  \
					If the directory doesn't exist, it will be created.  \
					Defaults to './temp_files'.")
parser.add_argument('-p', '--threads', dest='num_threads', default=0, type=int, 
					help="Number of parallel igblastn instances to spawn. \
					Defaults to max available processors.")
parser.add_argument('-v', '--tsv', dest="tsv_out", action='store_true', default=False, 
					help="NOT YET IMPLEMENTED. If set, the final output (from BlastParse) will be in tab-delimited format.  \
					Defaults to JSON output.")
parser.add_argument('-m', '--merge', dest="merge", action='store_true', default=False, 
					help="Use if the input files are paired-end FASTQs (either gzip compressed or uncompressed) \
					from Illumina platforms. Prior to running IgBLAST, reads will be merged with pandaseq. \
					Requires that pandaseq is installed.")
parser.add_argument('-n', '--next_seq', dest="next_seq", action='store_true', default=False, 
					help="Use if the run was performed on a NextSeq sequencer. \
					Multiple lane files for the same sample will be merged.")
parser.add_argument('-u', '--uaid', dest="uaid", type=int, default=None, 
					help="Use if the input files contain unique antibody identifiers (UAIDs). \
					UAIDs will be identified and incorporated into the output JSON file.")
parser.add_argument('-b', '--basespace', dest="use_basespace", default=False, action='store_true', 
					help="NOT YET IMPLEMENTED. Use flag if files should be downloaded directly from BaseSpace. \
					Files will be downloaded into the directory provided with the '-i' flag, which should be empty.")
parser.add_argument('-d', '--debug', dest="debug", action='store_true', default=False, 
					help="If set, will write all failed/exception sequences to file and give more informative errors.")
parser.add_argument('-s', '--species', dest='species', default='human', 
					choices=['human', 'macaque', 'mouse'])
args = parser.parse_args()


class launch_thread(threading.Thread):

	def __init__(self, in_file, out_file):
		threading.Thread.__init__(self)
		self.in_file = in_file
		self.out_file = out_file
		binary = './igblastn_' + platform.system().lower()
		self.cmd = '{3} -germline_db_V database/{0}_gl_V -germline_db_J database/{0}_gl_J -germline_db_D database/{0}_gl_D ' + \
				   '-organism {0} -domain_system imgt -auxiliary_data optional_file/{0}_gl.aux ' + \
				   '-show_translation -outfmt 3 -num_alignments_V 1 -num_alignments_D 1 -num_alignments_J 1 ' + \
				   '-query {1} -out {2}'.format(args.species, self.in_file, self.out_file, binary)

	def run(self):
		p = Popen(self.cmd, shell=True, stdout=PIPE, stderr=PIPE)
		stdout, stderr = p.communicate()
		


#####################################################################
# 
#                        FILES AND DIRECTORIES
#
#####################################################################

def build_temp_dirs():
	if args.temp_dir != '': 
		temp_directory = args.temp_dir
	else: 
		temp_directory = "./temp_files"
	temp_out_directory = temp_directory + "/temp_output"
	if not os.path.exists(temp_directory): os.mkdir(temp_directory)
	if not os.path.exists(temp_out_directory): os.mkdir(temp_out_directory)
	return temp_directory, temp_out_directory

def build_output_dir():
	output_dir = args.output
	if not os.path.exists(output_dir): os.mkdir(output_dir)
	return output_dir

def list_files(d): 
	if os.path.isdir(d):
		expanded_dir = os.path.expanduser(d)
		return sorted(glob.glob(expanded_dir + '/*'))
	else:
		return [d,]

def file_length(file):
	c = 0
	with open(file) as f:
		for i, l in enumerate(f):
			if l[0] == '>': c += 1
	return c

def num_procs():
	if args.num_threads > 0: 
		return args.num_threads
	return cpu_count()

def clean_up(temp_files, temp_out_files, temp_directory, temp_out_directory):
	for file in temp_out_files:
		os.remove(file)
	os.rmdir(temp_out_directory)
	for file in temp_files:
		os.remove(file)
	os.rmdir(temp_directory)

def file_splitter(file, splitlen, num_seqs, temp_directory):
	counter = 1
	files_list = []
	lines = open(file, 'r').read().replace(' ', '_').split('>')
	for line in range(0, num_seqs+1, splitlen):
		output = lines[line:line+splitlen]
		temp_filename = temp_directory + "/tempfile_" + str(counter)
		files_list.append(temp_filename)
		open(temp_filename, "w").write("")
		temp_file = open(temp_directory + "/tempfile_" + str(counter), "a")
		if counter == 1:
			temp_file.write('>'.join(output))
			counter += 1
		else:
			temp_file.write('>' + '>'.join(output))
			counter += 1
		temp_file.close()
	return files_list
		


#####################################################################
# 
#                            PRINTING
#
#####################################################################

def print_input_info(i):
	print ''
	print ''
	print '========================================'
	print 'Parallel IgBLAST'
	print '========================================'
	print ''	
	if len(i) > 1:
		print 'Input is a directory of {} files.\n\n'.format(len(i))
	else:
		print 'Input is a single file.\n\n'

def print_infile(i):
	b = os.path.basename(i)
	print '-'*len(b)
	print b
	print '-'*len(b)

def print_summary_output(g, e, f, blast_time, parse_time):
	total_seqs = g+e+f
	print ''
	print 'Out of {} total sequences:'.format(total_seqs)
	print '{} sequences processed normally'.format(g)
	print '{} sequences passed sanity checks, but could not be processed'.format(e)
	print '{} sequences failed sanity checks are were not processed'.format(f)
	print ''
	print 'IgBLAST took {0} seconds ({1} sequences per second)'.format(blast_time, total_seqs/blast_time)
	print 'parsing took {0} seconds ({1} sequences per second)'.format(parse_time, total_seqs/parse_time)
	print ''



#####################################################################
# 
#                              PARSING
#
#####################################################################

def line_generator(blast_file):
	f = open(blast_file, 'r')
	for line in f:
		yield line