#! /usr/bin/env python

import sys
import os
import Gnuplot
import math
import argparse

sys.path.append("../../EvolutionSBP/")
import numpy as np
import simulation_data as sd

#sys.argv[1] = 'exp_-20xx.hdf' - hdf data file
#sys.argv[2] = '0', ..., '4' the component to display, '-1' indicates sum
#sys.argv[3] = 'gif' - output type (optional)
#sys.argv[4] = 'test.gif' - output file name (optional)

"""Produces a 3D graph of the behaviour of the functions."""

parser = argparse.ArgumentParser(description=\
"""This program produces 2D representations of 3D time varying data from hdf files. The z axis gives the value of specified component. The x axis gives the  domain of the component. The y axis gives the time.""")
parser.add_argument('file',help =\
"""The hdf file, produced by simulation_data.py that contains the data to be
plotted""")
parser.add_argument('-dg','-dgtype',default="raw",help=\
"""The data group type to be plotted. See simulation_data for information on
the different data groups. This option is currently not implemented.""")
parser.add_argument('-c','-component',type=int,\
nargs='+',help=\
"""A list of the components for which you want graphs produced. If not given this defaults to all possible components.""", default=None)
parser.add_argument('-t','-terminal',default=None,help=\
"""Terminal commands for gnuplot. WARNING: Setting the terminal without setting the output can have odd effects.""")
parser.add_argument('-o','-outputfile',default=None,help=\
"""The file name for output of the graph. WARNING: Setting the output without setting the terminal can have odd effects.""")
parser.add_argument('-s','-saveintermediate',action='store_true',default=False,\
help=\
"""If option ensures that the gnuplot intermediate files which store the grid overlay and the plotted surface are not deleted.""")
parser.add_argument('-zrange',nargs='+',help=\
"""The range for the z axis of the plots given in Gnuplot style, e.g. "[0:0.1]". If not specified Gnuplot will choose an appropriate range. If one range is specified this range is used for each component. Other wise for each specified component one range must be given. The zranges are assumed to correspond to the components at the same position. This option has not been implemented yet.""")

args = parser.parse_args()

if args.o is not None:
    name,sep,ext = args.o.rpartition('.')
    args.oname = name
    args.oext = ext

with sd.SimulationHDF(args.file) as file:
    print "Producing graphs..."
    
    
    # Collect groups and datagroups
    sims = file.getSims()
    
    #set defaul components
    numOfComps = sims[0].raw[0].shape[0]
    if args.c is None:
        args.c = [i for i in range(numOfComps)]
    
    if args.zrange is not None:
        if len(args.zrange)==1:
            args.zrange = [args.zrange for comp in args.c]
        elif len(args.zrange)!=len(args.c):
            print "You have specified the ranges, "+args.zrange+" and the components "+args.c+". This is invalid. Either specify one zrange or specify one zrange per component."
            file.close()
            exit(1)
            
    # Setting up gnuplot
    g = Gnuplot.Gnuplot()
    #g('set terminal '+sys.argv[3])
    
    
    for sim in sims:
        print "Doing %s"%sim.name
    
        #Reset Gnuplot settings
        g.reset()
    
        #collate needed data
        x = be.array([sim.time[i][0] for i in range(len(sim.time))])
        y = sim.domain[0]
        z = be.empty((len(x),numOfComps,len(y)))
        for j in range(len(x)):
          z[j] = sim.raw[j] #[phi_index,:]
        
        #calculate data for minimal grid overlay
        numTimeLines = 10
        tskip = (len(x)-1)/numTimeLines
        a = x[::tskip]
        numDomainLines = 10
        dskip = (len(y)-1)/numDomainLines
        b = y[::dskip]
        c = z[::tskip,:,::dskip]        
        
        #Output data file with grid overlay
        g('set style data lines')
        for comp in args.c:
            Gnuplot.GridData(c[:,comp,:],xvals=a,yvals=b,\
                filename="grid%i.gp"%comp,binary=0)
        
        #os.system(r'awk "NF<2{printf\"\n\"}{print}" <overlay.gp >overlay1.gp')
        
        #Initialising Gnuplot settings for pm3d surface plot
        g.reset()
        if args.t is not None:
            g('set terminal '+args.t)
        g('set ylabel "radius" offset 0.6,-0.2')
        g('set xlabel "time" offset -0.3,-0.3')
        g('set xyplane relative 0.1')
        g('unset key')
        g('set grid')
        g('set xtics offset -0.3,-0.3 autofreq 0,0.2')
        g('set ytics offset 0.6,-0.2 autofreq 0,0.2')
        g('set ztics autofreq 0.1')
        
        
        #Plot pm3d surface with grid overlay
        for comp in args.c:
            g(r'set title "Phi %i from \nSimulation %s from\n%s" enhanced' %\
            (comp, sim.name,args.file))
            if args.o is not None:
                g('set out \"'+args.oname+"-%s"%sim.name+"-%i."%comp\
                    +args.oext+"\"")
            g.splot(Gnuplot.GridData(z[:,comp,:],x,y,\
                with_='pm3d',filename="surface%i.gp"%comp),\
              Gnuplot.File("grid%i.gp"%comp,binary=0,with_="lines")\
              )
    g.close()
    if not args.s:  
        os.system("rm *.gp")
    print '...Done.-'
