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 reversed([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)

#
# Plot of total FreeDB index size vs K
#
def plotISvK():
    global g
#    g = Gnuplot.Gnuplot(debug=1)
    
    dist = [22540694330399L, 134403379L, 560909744L, 1494688373L,
    3145467366L, 5776420479L, 9920593844L, 16847568836L, 29469690851L,
    53978982224L, 102553131654L, 197309295436L, 375069445703L,
    691292573588L, 1220066348633L, 2046380965390L, 3248410730672L,
    4871796452576L, 6903578407753L, 9257149947916L, 11777989750936L,
    14272274479420L, 16550118375625L, 18467921705896L,
    19954563688498L, 21013863178925L, 21706291519518L,
    22120514415903L, 22346645717764L, 22458924476538L,
    22509429313498L, 22529915160548L, 22537367332994L,
    22539782730829L, 22540474817747L, 22540648448837L,
    22540686135563L, 22540693105018L, 22540694181232L,
    22540694316205L, 22540694329409L, 22540694330354L,
    22540694330398L, 22540694330399L, 22540694330399L,
    22540694330399L, 22540694330399L, 22540694330399L,
    22540694330399L, 22540694330399L, 22540694330399L]

    scaleddist = Gnuplot.Data(list(x/dist[1] for x in dist))
    g('set data style linespoints')
    g('set logscale y 10')
    g('set key left Left reverse')
    g('set ylabel "Index size increase\\\\\\\\vs $K=1$" 1.5, 0')
    g('set xlabel "$K$" 0, .5')
    g('set xrange [1:15]')
    g.plot(scaleddist)
    
def main():
    os.chdir("../figures")
    mr = arange(1, 10)
#    plot3d(mr)
#    latex("egraph3d.tex")
#    png("egraph3d.png")
    plot2d(mr)
    latex("egraph2d.tex")
    plotISvK()
    latex("index-size-increase-vs-k.tex")
#    png("egraph2d.png")
#    raw_input()
    import sys; sys.exit(0)
    
if __name__ == "__main__":
    main()