##################################
#                                #
# Last modified 10/15/2015       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
from sets import Set
import matplotlib
matplotlib.use('Agg')
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.path as mpath
import matplotlib.lines as mlines
import matplotlib.patches as mpatches
from matplotlib.collections import PatchCollection

def label(xy, text):
    y = xy[1] + 0.05
    plt.text(xy[0], y, text, ha="center", family='sans-serif', size=14)

def run():

    if len(sys.argv) < 4:
        print 'usage: python %s config bar_height figure_wdith outputfileprefix' % sys.argv[0]
        print '\tconfig file format'
        print '\t#Name\tdomain\tstart\tend\tcolor'
        print '\tThe protein itself should be listed as a domain with the same name as in the first color'
        print '\tThe proteins will be plotted in the order listed'
        print '\tbar height and figure widht in inches'
        sys.exit(1)

    config = sys.argv[1]
    BH = float(sys.argv[2])
    FW = float(sys.argv[3])
    outfileprefix = sys.argv[4]

    ProteinDict = {}

    maxLength = 0

    Plist = []
    Dlist = []

    linelist = open(config)
    for line in linelist:
        if line.startswith('#') or line.strip() == '':
            continue
        fields = line.strip().split('\t')
        protein = fields[0]
        domain = fields[1]
        start = int(fields[2])
        end = int(fields[3])
        c = fields[4]
        if ProteinDict.has_key(protein):
            pass
        else:
            Plist.append(protein)
            ProteinDict[protein] = {}
            ProteinDict[protein]['domains'] = []
        if domain == protein:
            ProteinDict[protein]['p'] = (start,end,c)
            if end > maxLength:
                maxLength = end
        else:
            Dlist.append((domain,c))
            ProteinDict[protein]['domains'].append((domain,start,end,c))

    print 'plotting proteins'

    fig = plt.figure()
    fig.set_size_inches(FW,BH*len(ProteinDict.keys()))
    ax = fig.add_subplot(111)

    step = 1./(len(ProteinDict.keys()) + 1)
    print step
    pos = 1 - 0.5*step

    for protein in Plist:
#        print protein, ProteinDict[protein]
        (pstart,pend,c) = ProteinDict[protein]['p']
        length = 0.9*(pend/(maxLength + 0.0))
        print protein, pos, length, step/1.5, c
        rect = matplotlib.patches.Rectangle((0.05,pos), length, 0.66*step, color=c, linestyle = 'solid', edgecolor = 'black')
        ax.add_patch(rect)
        for (domain,dstart,dend,c) in ProteinDict[protein]['domains']:
            if max(dstart,dend) > maxLength:
               print 'domain outside of maximal protein length range, exiting'
               sys.exit(1)
            if max(dstart,dend) > pend:
               print 'domain outside of protein length range, exiting'
               sys.exit(1)
            S = 0.9*(dstart/(maxLength + 0.0))
            E = 0.9*(dend/(maxLength + 0.0))
            rect = matplotlib.patches.Rectangle((0.05 + S,pos), E-S, 0.66*step, color=c)
            ax.add_patch(rect)
        rect = matplotlib.patches.Rectangle((0.05,pos), length, 0.66*step, fill=None, color = 'black')
        ax.add_patch(rect)
        plt.text(0.05, pos + (2.1/3.)*step, protein, fontsize=BH*12)
        pos -= step

    plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
    plt.axis('equal')
    plt.axis('off')
    plt.savefig(outfileprefix + '.proteins.png', dpi=200)

    print 'plotting domains'

    Dlist = list(Set(Dlist))

    fig = plt.figure()
    fig.set_size_inches(FW,BH*len(Dlist))
    ax = fig.add_subplot(111)

    step = 1./(len(Dlist) + 1)
    pos = 1 - 0.5*step
    for (domain,c) in Dlist:
        rect = matplotlib.patches.Rectangle((0.1,pos), 0.5, 0.66*step, color=c, linestyle = 'solid', edgecolor = 'black')
        ax.add_patch(rect)
        rect = matplotlib.patches.Rectangle((0.1,pos), 0.5, 0.66*step, fill=None, color = 'black')
        ax.add_patch(rect)
        plt.text(0.625, pos + 0.3*step, domain, fontsize=BH*12)
        pos -= step

    plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
    plt.axis('equal')
    plt.axis('off')
    plt.savefig(outfileprefix + '.domains.eps', format='eps')
    
run()