##################################
#                                #
# Last modified 05/10/2011       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import os
import random
import string

def run():

    if len(sys.argv) < 4:
        print 'usage: python %s DREME_output N_sequences weblogo_location oufile_prefix ' % sys.argv[0]
        sys.exit(1)
    
    input = sys.argv[1]
    Nseq = int(sys.argv[2])
    logoPath = sys.argv[3]
    outfile_prefix = sys.argv[4]

    tempfile=sys.argv[4]+'.temp'

    MotifList=[]

    lineslist = open(input)
    inMOTIF=False
    n=''
    PWM=''
    for line in lineslist:
        if line.startswith('MOTIF '):
            inMOTIF=True
            PWM=[]
            continue
        if inMOTIF: 
            if line=='\n':
                inMOTIF=False
                MotifList.append((PWM,n,E))
                n=''
                PWM=''
                continue
            if line.startswith('letter'):
                n=int(line.split('nsites= ')[1].split(' ')[0])
                E=line.split('E= ')[1].split(' ')[0]
                continue
            else:
                fields=line.strip().split(' ')
                weights=[]
                for w in fields:
                    weights.append(float(w))
                PWM.append(tuple(weights))

    m=0
    for (PWM,n,E) in MotifList:
        print PWM,n
        m+1
        outfile_name=outfile_prefix+'.motif'+str(m)+'.nsites='+str(n)+'.E='+E
        PWMPool={}
        j=0
        for (A,C,G,T) in PWM:
            PWMPool[j]=[]
            for i in range(int(Nseq*A)):
                PWMPool[j].append('A')
            for i in range(int(Nseq*C)):
                PWMPool[j].append('C')
            for i in range(int(Nseq*G)):
                PWMPool[j].append('G')
            for i in range(int(Nseq*T)):
                PWMPool[j].append('T') 
            j+=1
        motifLength=j
        seqfile = open(tempfile,'w')
        for i in range(Nseq):
            sequence=''
            for j in range(motifLength):
                sequence+=random.sample(PWMPool[j],1)[0]
            seqfile.write('>seq'+str(i)+'\n')
            seqfile.write(sequence+'\n')
        seqfile.close()
        cmd = logoPath + ' -f ' + tempfile + ' -F PNG -c ' + '-a -Y -n -M -k 1 -o ' + outfile_name
        contents = os.system(cmd)
        os.remove(tempfile)
   
run()