#! /usr/bin/env python 
###################################
#    Davi Ortega 7/1/2012 
###################################
import sys
import bitk
import os
import multiprocessing as multip
if '-h' in sys.argv:
	print 'get_histogram.py hmm_model.hmm genome_list fasta.fa -n Number of processors'
	sys.exit()


hmm_model = sys.argv[1]

def find_mcp(genome_list='', fasta=''):
	sel = []
	datafile = open(genome_list,'r')
	for line in datafile:
		sel.append(line.replace('\n',''))
	datafile.close()
	
	seq_dic, seq_list = bitk.fastareader(fasta)
	results = []	
	file_list = []
	wrong_list = []

	for genome in sel:
		seq_dic_sel = {}
		output = ''
		for name in seq_list:
			if name.split('-')[0].split('.')[-1] == genome:
				output += '>' + name + '\n' + seq_dic[name] + '\n'
				seq_dic_sel[name] = seq_dic[name]
		outfile = open(genome + '_mcp.fa', 'w')
		outfile.write(output)
		outfile.close()
		results.append(seq_dic_sel)
		if len(seq_dic_sel.keys()) <= 3 :
			wrong_list.append(genome + '_mcp.fa')
		else:
			file_list.append(genome + '_mcp.fa')
	datafile = open('genomes_problems_finding_mcp.dat', 'w')
	datafile.write('\n'.join(wrong_list))
	datafile.close()
	return file_list

def chain_of_fire(fasta_file=''):
	#trimming by hmm
	os.system('hmmsearch --noali --cut_tc ' + hmm_model + ' ' + fasta_file + ' > output_' + fasta_file[:-3] + '.dat')
	
	seq_dic, seq_list = bitk.fastareader(fasta_file)
	datafile = open('output_' + fasta_file[:-3] + '.dat','r')

	new_seq_dic = {}
	output = ''
	for line in datafile:
        	if '>>' in line:
                	name = line.split(' ')[1]
	                for line in datafile:
        	                if line[60:66].replace(' ','').isdigit():
                	                start = int(line[60:66])
                        	        end = int(line[68:74])
					new_seq_dic[name] = seq_dic[name][start:end]
					output += '>' + name + '\n' + seq_dic[name][start:end] + '\n'
	                                #print name + '\t' + str(start) + '\t' + str(end)
        	                        break
	
	datafile.close()
	
	outfile = open(fasta_file[:-3] + '_hmmtrim.fa', 'w')
	outfile.write(output)
	outfile.close()
	hmm_filename = fasta_file[:-3] + '_hmmtrim.fa'
	#run mafft linsi
	linsi_filename = hmm_filename[:-3] + '.linsi.fa'
	os.system('linsi ' + hmm_filename + ' > ' + linsi_filename )
	#Transform fasta in phy
	os.system('fa2phy ' + linsi_filename)
	phy_filename = linsi_filename[:-3] + '.phy'
	#Run phyml
	os.system('phyml -i ' + phy_filename + ' -d aa -m JTT')
	tree_filename = phy_filename + '_phyml_tree.txt'
	os.system('rectaxontree2 ' + tree_filename)
	tree_filename = phy_filename + '_phyml_tree.rec.txt'
	#writing customized R script
	r_filename = linsi_filename[:-3] + '.r'
	r_output = linsi_filename[:-3] + '.R.dat'
	r_script = open(r_filename,'w')
	r_script.write("library('ape')\ntree <- read.tree('"+tree_filename+"')\nD <- cophenetic.phylo(tree)\nwrite(as.vector(D),file='" + r_output + "', ncolumns=length(D))")
	r_script.close()
	os.system('R CMD BATCH ' + r_filename)
	#reading results
	r_file = open(r_output,'r')
	diversity = []
	for line in r_file:
		diversity += line.split(' ')
	while 0 in diversity:
		diversity.remove('0')
	div_output = ' '.join(diversity)
	r_file.close()
	r_out = open(r_output,'w')
	r_out.write(div_output)
	r_out.close() 
	print diversity
	return diversity

file_list = find_mcp(sys.argv[2], sys.argv[3])
hmmtrim_file_list = []
Dmatrix_file_list = []

if '-n' in sys.argv:
	i = sys.argv.index('-n')
	Np = int(sys.argv[i + 1])
	pool = multip.Pool(processes=Np)
	pool.map(chain_of_fire, file_list)
	for fasta in file_list:
		Dmatrix_file_list.append(fasta[:-3]+'_hmmtrim.linsi.R.dat') 
	
else:
	for fasta in file_list:
		print 'working here on ' + fasta
		Dmatrix_file_list.append(fasta[:-3]+'_hmmtrim.linsi.R.dat')
		chain_of_fire(fasta)

Dm_all = []

for Dm in Dmatrix_file_list:
	datafile = open(Dm, 'r')
	for line in datafile:
		Dm_all += line.replace('\n','').split(' ')
	datafile.close()

output = ''
for value in Dm_all:
#	if value != '0':
	if float(value) > 0.5:
		output += value + '\n'

final_file = open('final.dat','w')
final_file.write(output)
final_file.close()
		
		


