#! /usr/bin/env python 
###################################
#    Davi Ortega 12/9/2014 
###################################
import sys
import os
import bitk
if '-h' in sys.argv:
	print 'Pipeline to build weblogos of top N COG_finder groups. PASE cog.dat fasta.fa N \n \
		1) Se - Separate the sequences of each group in fasta files \n \
		2) Al - Align each file using linsi \n \
		3) Pa - Sequential profile alignment \n \
		4) Se - Separate the files again \n \
		5) W - Weblogo \n\n \
                Options: -c F force CD-HIT to filter sequences F*100 % identical (F -> { 0,1} ) \n \
                         -n N Number of stacks per line in the weblogo\n \
                         --trimMSA N1 N2 It will trim the final profile\n \
                         --noWL          no weblogo will be produced'
	sys.exit()

if "-c" in sys.argv:
    try:
        C = float(sys.argv[sys.argv.index("-c") + 1])
    except:
        print " value passed to -c must be between 0 and 1"
        sys.exit()

N = int(sys.argv[3])

fastainfile = sys.argv[2]
cogdatinfile = sys.argv[1]

os.system('fa2fa ' + fastainfile + ' -u')
print 'getseqfromcogoutput ' + cogdatinfile + ' ' + fastainfile + ' all --cog-flag'
os.system('getseqfromcogoutput ' + cogdatinfile + ' ' + fastainfile + ' all --cog-flag')
fastainfile = cogdatinfile[:-4] + ".cogs.all.fa"
os.system('IIintag ' + fastainfile + ' 0 spaceholder')
fastainfile = cogdatinfile[:-4] + ".cogs.all.IIintag.fa"

print "Phase 1 and 2 - Se + Al "
print fastainfile
for i in range(1,N+1):
	print "grep 'COG" + str(i) + "-spaceholder' -A1 " + fastainfile + " > " + fastainfile[:-3] + ".cluster" + str(i) + ".fa"
	os.system("grep 'COG" + str(i) + "-spaceholder' -A1 " + fastainfile + " > " + fastainfile[:-3] + ".cluster" + str(i) + ".fa")
        if "-c" in sys.argv:
            os.system("cd-hit -i " + fastainfile[:-3] + ".cluster" + str(i) + ".fa -o " + fastainfile[:-3] + ".cluster" + str(i) + ".cdhit." + str(C) + ".fa -c " + str(C))
            os.system("mv " + fastainfile[:-3] + ".cluster" + str(i) + ".cdhit." + str(C) + ".fa" + " " + fastainfile[:-3] + ".cluster" + str(i) + ".fa")
	os.system("linsi --thread 12 " + fastainfile[:-3] + ".cluster" + str(i) + ".fa > " + fastainfile[:-3] + ".cluster" + str(i) + ".linsi.fa")
#	os.system("fa2fa " + fastainfile[:-3] ".cluster" + str(i) + ".linsi.fa -u")

print "Phase 3 - Pr"
tmpprof = "PASEfile.tmp.0.fa"
os.system("mafft-profile " + fastainfile[:-3] + ".cluster1.linsi.fa " + fastainfile[:-3] + ".cluster2.linsi.fa > " + tmpprof)


for i in range(3,N+1):
	newtmp = "PASEfile.tmp." + str(i-2) + ".fa"
	os.system("mafft-profile " + tmpprof + " " + fastainfile[:-3] + ".cluster" + str(i) + ".linsi.fa > " + newtmp)
	tmpprof = newtmp


os.system('fa2fa ' + newtmp + ' -u')

if "--trimMSA" in sys.argv:
    Nx = sys.argv.index("--trimMSA") + 1
    N1, N2 = [ int(i) for i in sys.argv[Nx: Nx+2] ]
    print Nx
    print N1
    print N2
    seq_dic, seq_list = bitk.fastareader(newtmp)
    seq_dic = bitk.trimMSA(seq_dic, N1, N2, seq_list)
    output = ""
    for seq in seq_list:
        output += ">" + seq + "\n" + seq_dic[seq] + "\n"
    with open(newtmp, "w") as f:
        f.write(output)

os.system('fa2fa ' + newtmp + ' -u')

os.system('cp ' + newtmp + ' ' + fastainfile[:-3] + '.fullprofile.fa')

if "--noWL" in sys.argv:
    sys.exit()

print "Phase 4 and 5 - Se + We"

if "-n" in sys.argv:
    Nopt = "-n " + sys.argv[sys.argv.index("-n")+ 1]
else:
    Nopt = ""

for i in range(1,N+1):
        os.system("grep 'COG" + str(i) + "-spaceholder' -A1 " + newtmp + " > " + fastainfile[:-3] + ".profile" + str(i) + ".fa")
	os.system("weblogo -f " + fastainfile[:-3] + ".profile" + str(i) + ".fa -o " + fastainfile[:-3] + ".profile" + str(i) + ".pdf " + Nopt + " -A protein -s large -F pdf -c chemistry --errorbars NO --resolution 300")

print "DONE"


	


