##################################
#                                #
# Last modified 2017/12/13       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import os
import random
import string

def run():

    if len(sys.argv) < 4:
        print 'usage: python %s HOMER_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)
    n=''
    PWM=''
    for line in lineslist:
        fields = line.strip().split('\t')
        if line.startswith('>'):
            if PWM != '':
                MotifList.append((PWM,MOTIF))
            PWM=[]
            MOTIF = fields[0].split('>')[1] + '_' + fields[5]
            continue
        weights = []
        for w in fields:
            weights.append(float(w))
        PWM.append(tuple(weights))
    MotifList.append((PWM,MOTIF))

    m=0
    for (PWM,MOTIF) in MotifList:
        m += 1
        outfile_name = outfile_prefix + '.motif' + str(m) + '.' + MOTIF.replace(':','_').replace('(',',fraction=').replace('%','').replace(')','')
        print outfile_name
        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()