#!/usr/bin/env python
#import fileinput
from bitk import *
import sys

if '-h' in sys.argv:
	print 'The table as the first file, alignment of chews as a second file. make sure the chews has as name the loci number and a \
		mark "/" before the sequence numbers\n \n \
		Flags: \n\n \
		-ns	No Subgroup\n'
	sys.exit()


#reader of CheW and classes

rawtable = open(sys.argv[1],'r')
classes = {}
fields = range(5)
old_class = ''
chew = ''
for line in rawtable:
	end = 0
	start = 0
	for field in fields:
		start = line[end:].find(',') + end + 1
		end = line[start:].find(',') + start
		if end == start - 1:
			end = len(line) - 2
		if field == 3:
			new_class = line[start:end]
		if field == 0:
			chew = line[start:end]
	print chew + '-----' + new_class
	if new_class == '':
		new_class = old_class
	else:
		old_class = new_class
	if chew != '':
		classes[chew] = new_class

rawtable.close()

print classes

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

for line in seq:
	if line[0] != '<':
		end = line.find('/')
		if line[:end] not in list:
			list.append(line[:end])

dataout = ''

if '-ns' in sys.argv:
	for v, k in classes.iteritems():
		k = k.replace('-1','')
		k = k.replace('-2','')
		k = k.replace('-3','')
		k = k.replace('-4','')
		k = k.replace('-5','')
		k = k.replace('-6','')
		classes[v] = k

for locus in list:
	if locus != ' ' or locus.find('CLUSTAL') == -1:
		try:
			dataout = dataout + locus + ' ' + classes[locus] + '\n'
		except KeyError:
			dataout = dataout + locus + ' ' + 'NT\n'
	else:
		print 'Taking CLUSTAL out'
	
output = open('color.chews.classes.dat','w')
output.write(dataout)
output.close()

print 'Done...'
