#
#  plotprofile.py
#  ENRAGE
#

try:
    import psyco
    psyco.full()
except:
    pass

import sys
import optparse
from pylab import *
from math import *
import matplotlib
from commoncode import getConfigParser, getConfigOption, getConfigIntOption, getConfigBoolOption


print "plotprofile: version 2.3"

def main(argv=None):
    if not argv:
        argv = sys.argv

    usage = "usage: python %s infile outfile.png [--scale] [--max weightMax] [--ymin bottom] [--ymax top] [--subtractEvens]"

    parser = makeParser(usage)
    (options, args) = parser.parse_args(argv[1:])

    if len(args) < 2:
        print usage
        sys.exit(1)

    infile = args[0]
    pngfilename = args[1]

    plotprofile(infile, pngfilename, options.doScale, options.weightMax, options.ymin, options.ymax, options.subtractEvens)


def makeParser(usage=""):
    parser = optparse.OptionParser(usage=usage)
    parser.add_option("--scale", action="store_true", dest="doScale")
    parser.add_option("--max", type="float", dest="weightMax")
    parser.add_option("--ymin", type="float", dest="ymin")
    parser.add_option("--ymax", type="float", dest="ymax")
    parser.add_option("--subtractEvens", action="store_true", dest="subtractEvens")

    configParser = getConfigParser()
    section = "plotprofile"
    doScale = getConfigBoolOption(configParser, section, "doScale", False)
    weightMax = getConfigIntOption(configParser, section, "weightMax", -1)
    ymin = getConfigOption(configParser, section, "ymin", None)
    ymax = getConfigOption(configParser, section, "ymax", None)
    subtractEvens = getConfigBoolOption(configParser, section, "subtractEvens", False)

    parser.set_defaults(doScale=doScale, weightMax=weightMax, ymin=ymin, ymax=ymax, subtractEvens=subtractEvens)

    return parser


def plotprofile(inFileName, pngfilename, doScale=False, weightMax=-1, ymin=None, ymax=None, subtractEvens=False):
    infile = open(inFileName)
    limitYscale = False
    if ymax is not None:
        limitYscale = True
    else:
        ymax = 0.

    if ymin is not None:
        limitYscale = True
    else:
        ymin = 0.

    matplotlib.use("Agg")

    labelList = []
    dataList = []
    plotList = []
    xmin = 10**20
    xmax = -10**20

    xcoordList = []
    datapointList = []
    weightList = []
    line = infile.readline()
    fields = line.strip().split()
    for data in fields[1:-1]:
        datapoint = float(data)
        if datapoint < xmin:
            xmin = datapoint

        if datapoint > xmax:
            xmax = datapoint

        xcoordList.append(datapoint)

    index = 1
    for line in infile:
        fields = line.strip().split()
        datapointList = []
        for data in fields[1:-1]:
            datapointList.append(float(data))

        if subtractEvens and index % 2 == 0:
            for dataIndex in range(len(datapointList)):
                dataList[-1][dataIndex] -= datapointList[dataIndex]
        else:
            dataList.append(datapointList)

        weight = float(fields[-1])
        if subtractEvens and index % 2 == 0:
            pass
        else:
            labelList.append(fields[0])
            if weight > weightMax:
                weightMax = weight

            weightList.append(weight)

        index += 1

    for index in range(len(dataList)):
        newList = []
        if doScale:
            scale = weightList[index] / weightMax
            print weightList[index], weightMax, scale
            for val in dataList[index]:
                newList.append(val * scale)
        else:
            newList = dataList[index]

        plotList.append(plot(xcoordList, newList, linewidth=3.0))

    xticks(xcoordList, rotation="vertical")
    xlim(xmin - 0.1, xmax + 0.1)
    if limitYscale:
        ylim(ymin, ymax)

    legend(plotList, labelList)
    savefig(pngfilename)


if __name__ == "__main__":
    main(sys.argv)