import Gnuplot, Gnuplot.funcutils
import os
from funcs import *

#
# Exporters
#
def latex(f):
    g('set terminal push')
    #g('set terminal latex 10')
    g('set terminal epslatex "default" 10')
    g('set format xy "$%g$"')
    g.set_string('output', f.replace(".tex",".eps"))
    g('set size .675, .675')
    #g('set size 3/5, 3/5')
    g.refresh()
    g('set terminal pop')
    g.set_string('output')
    # Darnit, fix the path in the output file
    os.system("echo ',s,{%(b)s},{figures/%(b)s},\nw' | ed %(e)s" %
              {'b': f.replace(".tex",""), 'e': f})

def png(f):
    g('set terminal push')
    g('set terminal png')
    g.set_string('output', f)
    g.refresh()
    g('set terminal pop')
    g.set_string('output')

#
# 3d plot of k versus K
#
def plot3d(mr):
    global g
    g = Gnuplot.Gnuplot(debug=1)
    Kr = arange(1, 9)
    data = Gnuplot.funcutils.compute_GridData(mr, Kr, I, with="lines 2")
    g('set data style lines')
    g('set hidden3d offset -1')
    g.xlabel('$m$')
    g.ylabel('$K$')
    g('set zlabel "$I(m)$"')
    g('set logscale z 2')
    g.splot(data)

#
# Sampled 2d plot of k versus K = [1..4, 50]
#
def plot2d(mr):
    global g
    g = Gnuplot.Gnuplot(debug=1)
    def foo(x):
        if x < 50:
            return x
        return "\\\\infty"
    data = [Gnuplot.Data(mr, [I(m, K) for m in mr],
                         title='$K = %s$'%foo(K))
            for K in [1, 2, 3, 4, 50]]
    g('set data style linespoints')
    g('set logscale y 2')
    g('set key left Left reverse')
    g('set lmargin 5')
    g('set ylabel "$I(m)$" 1.5, 0')
    g('set xlabel "$m$" 0, .5')
    g.plot(*data)

def main():
    os.chdir("../figures")
    mr = arange(1, 10)
    plot3d(mr)
    latex("egraph3d.tex")
    #png("egraph3d.png")
    plot2d(mr)
    latex("egraph2d.tex")
    #png("egraph2d.png")
    #raw_input()

if __name__ == "__main__":
    main()