#! /usr/bin/env python 
###################################
##    Davi Ortega                ##
###################################
import sys
import bitk

if '-h' in sys.argv:
	print "Update fasta file with class information retrived from chew-class-library.hmm output file"
	print "Sintax: hmmchewparser.py fasta_file.fa output_hmm.dat"
	print "outfile = fasta_file.class.fa"
	sys.exit()

DIFF_THRES = 15
MIN_SCORE = 20


seq_dic, seq_list = bitk.fastareader(sys.argv[1])

hmmfile = open(sys.argv[2],'r')

new_seq = 1
for line in hmmfile:
	if "Query sequence" in line:
		tag = line.split(':')[1].replace(' ','').replace('\n','')
		new_seq = 0
	if ("------" in line) and (new_seq == 0):
		new_seq = 1
		chew_class = []
		for line2 in hmmfile:
			line_info = line2.split(' ')
			while '' in line_info:
				line_info.remove('')
			#print tag
			chew_class.append([line_info[0].replace('chew-',''),float(line_info[1])])
			if len(chew_class) == 2:
				break
		if (chew_class[0][1] >= MIN_SCORE) and (chew_class[0][1] - DIFF_THRES >= chew_class[1][1]):
			seq_dic[tag + '-' + chew_class[0][0]] = seq_dic[tag]
			del seq_dic[tag]
		else:
			seq_dic[tag + '-' + 'Uncat'] = seq_dic[tag]
			del seq_dic[tag]

hmmfile.close()

outfile = open(sys.argv[1][:-2] + 'class.fa', 'w')
output = ''
for tag, seq in seq_dic.iteritems():
	output += '>' + tag + '\n' + seq + '\n'

outfile.write(output)
outfile.close()



