##################################
#                                #
# Last modified 2017/06/19       # 
#                                #
# Georgi K. Marinov              #
#                                # 
##################################

import sys
import os
import random
import string

def run():

    if len(sys.argv) < 4:
        print 'usage: python %s input.meme 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=''
    CurrentMotif = ''
    PWM=''
    for line in lineslist:
        if line.strip() == '':
            continue
        if line.startswith('URL'):
            continue
        if line.startswith('MEME version'):
            continue
        if line.startswith('strands'):
            continue
        if line.startswith('Background'):
            continue
        if line.startswith('A '):
            continue
        if line.startswith('ALPHABET'):
#            letters = line.strip().split(' ')[-1]
#            i = 0
#            for L in letters:
#                PosIDDict[L] = i
#                i+=1
            continue
        if line.startswith('MOTIF '):
            if CurrentMotif != '':
                (MOTIF_ID,MOTIF_Name) = CurrentMotif
                MotifList.append((PWM,MOTIF_ID,MOTIF_Name))
            MOTIF_ID = line.strip().split(' ')[1]
            MOTIF_Name = line.strip().split(' ')[2]
            CurrentMotif = (MOTIF_ID,MOTIF_Name)
            PWM = []
            continue
        if line.startswith('letter'):
            continue
        newline = line.strip()
        while '  ' in newline:
            newline = newline.replace('\t',' ').replace('  ',' ')
        fields = newline.strip().split(' ')
#        print fields
        weights=[]
        for w in fields:
            weights.append(float(w))
        PWM.append(tuple(weights))

    m=0
    for (PWM,MOTIF_ID,MOTIF_Name) in MotifList:
        print MOTIF_ID,MOTIF_Name
        m+1
        outfile_name = outfile_prefix + '.' + MOTIF_ID + '.' + MOTIF_Name
        outfile_name = outfile_name.replace(')','').replace('(','')
        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()