"""
Plotting script for energies based on precomputed values
"""

# Copyright (C) 2013 Marie E. Rognes

import os
import sys
import pylab
import numpy

if len(sys.argv) < 2:
    print "Usage: python plot_energies.py results-directory"
    sys.exit()

resultsdir = sys.argv[1]
resultspath = os.path.join(os.getcwd(), resultsdir)
print "Treating data from %s" % resultspath
sys.path.insert(0, resultspath)

# Increase fonts
myfontsize=22
pylab.rc('lines', linewidth=2)
pylab.rc('legend', fontsize=myfontsize)
pylab.rc("lines", markeredgewidth=2.0)

pylab.figure(figsize=(10, 8))

# Import each of the cases here
from energies import times, kinetics, potentials

pylab.plot(times, kinetics, label="$E_k$")
pylab.plot(times, potentials, label="$E_p$")
pylab.plot(times, numpy.array(kinetics) + numpy.array(potentials),
           label="$E_T$")

pylab.xlabel("$t$", fontsize=myfontsize)
pylab.legend(loc='best')
pylab.grid(True)
print "Storing plot to '%s/energies.pdf'" % resultsdir
pylab.savefig("%s/energies.pdf" % resultsdir)
pylab.show()

