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

import sys
import string
import pandas as pd
import os
import math
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import logomaker

def run():

    if len(sys.argv) < 2:
        print 'usage: python %s meme outfile [-bits]' % sys.argv[0]
        sys.exit(1)

    meme1 = sys.argv[1]
    outfilename = sys.argv[2]

    doBits = False
    if '-bits' in sys.argv:
        doBits = True

    Xticks = []
    X = 0

    MOT1 = []

    lineslist = open(meme1)
    InMotif = False
    for line in lineslist:
        if line.startswith('letter-probability matrix:'):
            InMotif = True
            continue
        if InMotif:
            newline = line
            while ' ' in newline:
                newline = newline.replace(' ','\t')
            newline = newline.strip()
            while '\t\t' in newline:
                newline = newline.replace('\t\t','\t')
            fields = newline.split('\t')
            A = float(fields[0])
            C = float(fields[1])
            G = float(fields[2])
            T = float(fields[3])
            print A,C,G,T
            if doBits:
                E = 2 - (-A*math.log(A+1e-6,2) - C*math.log(C+1e-6,2) - G*math.log(G+1e-6,2) - T*math.log(T+1e-6,2))
                print A, C, G, T, E, A*E,C*E,G*E,T*E
                MOT1.append((A*E,C*E,G*E,T*E))
            else:
                MOT1.append((A,C,G,T))
            X+=1
            Xticks.append(X)

    outfile = open(outfilename,'w')
    outline = 'Pos\tA\tC\tG\tT'
    outfile.write(outline + '\n')
    i=1
    for (A,C,G,T) in MOT1:
        outline = str(i) + '\t' + str(A) + '\t' + str(C) + '\t' + str(G) + '\t' + str(T)
        outfile.write(outline + '\n')
        i+=1

    outfile.close()

    MOT1_pd = pd.read_csv(outfilename, sep="\t", index_col=0)

    print MOT1_pd

    logo = logomaker.Logo(MOT1_pd, font_name='Arial Rounded MT Bold')

#    logo.ax.set_ylabel('bits')
    Xticks = []
    for i in range(len(MOT1)):
        Xticks.append(i+1)

    logo.style_spines(visible=False)
    logo.ax.set_xticks(Xticks)
#    logo.ax.set_xticklabels(Xticks)
 
    plt.savefig(outfilename)

run()