#!/usr/bin/python
# -*- coding: utf-8 -*-

# Standalone script to create a movie from 2D HDF5 output. See "h5movie.py -h" for help.
# This script is put in "Public domain" by it's authors:
#   Tanja Bode <tanja.bode@physics.gatech.edu>
#   Roland Haas <rhaas@tapir.caltech.edu>
#   Frank Löffler <knarf@cct.lsu.edu>


import getopt
import sys
import os
import shutil
import re
import glob
import numpy
import h5py
import string
sys.path.insert(0, "/home/knarf/utils/visit2_5_0.linux-x86_64/2.5.0/linux-x86_64/lib/site-packages/")
import visit

# get started with VisIt (this needs to be global)
visit.AddArgument('-nowin')
visit.Launch()
#LaunchNowin()
visit_dir = dir()


def usage():
        print 'usage: h5movie <rundir> <var>.<plane>'
        print 'Standalone script to create a movie from 2D HDF5 output.  <var> should include'
        print 'the plane (e.g., rho.xy).  We assume the HDF5 files for chained runs have'
        print 'already been merged.'
        print ''
        print 'Default movie:  Pseudocolor map, linear scaling, no BH masking, orangehot colortable'
        print ''
        print '        -h       Display this help'
        print '        -v       Print status information while running'
        print ''
        print '        Plotted Quantity Options:'
        print '                -e EXPR  Feed in expression for quantity on the command line'
        print '                     * Ex. -e "sqrt(1 - 1/<WHISKY--w_lorentz>^2)"'
        print '                     * Masking and atmosphere thresholds are done on top of this expression, but'
        print '                       can be referenced via the quantities Mask and Atmo'
        print '                     * Reference to other variables must use full conn_cmfe() statement.'
        print '                -m       Mask variable with AHMask'
        print '                -a       Extra threshold to plot where above atmosphere only'
        print '                -w ATMO  Set density of what is considered atmosphere'
        print '                -d THR   Express (minimum) extra threshold on expression itself'
        print ''
        print '        Plot Specification Options:'
        print '                -b BDY   Spatial extent (Default: 10)'
        print '                -n ZMIN  sets min z-value (Default: none)'
        print '                -x ZMAX  sets max z-value (Default: none)'
        print ''
        print '                -u       Colortable'
        print '                -l       Choose logarithmic scaling'
        print '                -c       Add contours'
        print ''
        print '        Movie Options:'
        print '                -f MBASE Specify different base for frames (e.g. rhoMasked_xz_Wide)'
        print '                -o TITLE Override default title on frames (default <var>(<plane>))'
        print '                -t VTYPE Type of movie (gif, mp4, ogg ... Default: mp4)'

def main():
        ############################################################
        ## User-specified variables currently not set in options. ##
        ############################################################
        verbose_opts="/dev/null"
        outputsize=(800,800)
        
        max=""          # spatial boundaries (default 10)
        var_min="none"  # lower limit of plotted variable  (auto searches for max over all time)
        var_max="none"  # upper limit of plotted variable
        
        VarScaling="Linear"     # Scaling: bulk variable <var> (Linear/Log)
        ContourScaling="Linear" # Scaling: contours (Linear/Log)
        OFudge=-0.001           # Fudge for the origin of the Clip operator
        MaxIt=10000000000000
        
        # Color table for the pseudoplot:
        #       orangehot, hot, bluehot
        colortable="orangehot"
        
        # Minimum number of frames that most be present for mp4 output to succeed (found by trial-and-error)
        MinNumFrames=66
        
        
        ####################################################
        ## Defaults and acquisition of basic information. ##
        ####################################################
        Supported_Movietypes=["mp4", "gif", "ogg"]
        ## We "support" certain thorns here because VisIt is case-sensitive while the variable names
        ## taken from the output files generally does not give the proper cases.
        Supported_Variables={
                "rho":       "HYDROBASE--rho",
                "press":     "HYDROBASE--press",
                "eps":       "HYDROBASE--eps",
                "velx":      "HYDROBASE--vel_lb_0_rb_",
                "vely":      "HYDROBASE--vel_lb_1_rb_",
                "velz":      "HYDROBASE--vel_lb_2_rb_",
                "dens":      "GRHYDRO--dens",
                "tau":       "GRHYDRO--tau",
                "sx":        "GRHYDRO--scon_lb_0_rb_",
                "sy":        "GRHYDRO--scon_lb_1_rb_",
                "sz":        "GRHYDRO--scon_lb_2_rb_",
                "w_lorentz": "HYDROBASE--w_lorentz",
        }
        CONTOURS=0
        
        variable=""
        TYPE="mp4"
        verbose_opts="/dev/null"
        AtmoVal="rho_min"
        ExternalExpr=""
        MASKED="no"
        OnlyOverAtmo="no"
        FrameBase=""
        newTitle=""
        range_max="30"
        ExprThresh=""
        
        opts, args = getopt.gnu_getopt(sys.argv[1:], "hlmct:n:x:b:u:f:ae:d:o:vw:")
        for o, a in opts:
                if o == "-h":
                        usage()
                        sys.exit(0)
                elif o == "-m":
                        MASKED="yes"
                elif o == "-c":
                        CONTOURS=1
                elif o == "-l":
                        VarScaling="Log"
                        ContourScaling="Log"
                elif o == "-t":
                        TYPE=a
                elif o == "-u":
                        colortable=a
                elif o == "-n":
                        var_min=a
                elif o == "-x":
                        var_max=a
                elif o == "-b":
                        range_max=a
                elif o == "-f":
                        FrameBase=a
                elif o == "-a":
                        OnlyOverAtmo=a
                elif o == "-w":
                        AtmoVal=a
                elif o == "-e":
                        ExternalExpr=a
                elif o == "-d":
                        ExprThresh=a
                elif o == "-o":
                        newTitle=a
                elif o == "-v":
                        verbose_opts="/dev/stdout"
                else:
                        print "ERROR: Invalid parameter '%s'." % o
                        usage()
                        sys.exit(1)
        
        if len(args) != 2:
                usage()
                sys.exit(1)

        sys.stdout = open(verbose_opts, 'w')
        
        dirname=args[0]
        filebase=args[1]
        
        PFILE=glob.glob('%s/output-0000/*.par' % dirname)
        if not PFILE:
                print "ERROR: Could not find parameter file in output directory."
                sys.exit(1)
        PFILE=PFILE[-1]

        pardir = re.match('.*/(.*)\.par', PFILE).group(1);
        
        m = re.match('([^.]*).*\.([^.]*)$', filebase)
        if not m:
                print "Could not parse %s into variable and symmetry part" % filebase
                sys.exit(1)
        var=m.group(1)
        plane=m.group(2)
        planex=plane[0]
        planey=plane[1]

        if not TYPE in Supported_Movietypes:
                print "ERROR: The movie type %s is not currently supported." % TYPE
                sys.exit(1)
        
        # h5 variable name
        if not var in Supported_Variables:
                print "ERROR: Variable %s comes from a thorn that is currently not supported.  Add the thorn's variable list!" % var
                sys.exit(1)
        variable = Supported_Variables[var]
        
        # Create expression to be plotted (e.g. masked vs not masked)
        if not ExternalExpr:
               Vararg="<%s>" % variable
        else:
               Vararg=ExternalExpr
        
        if MASKED == 'yes':
               Maskvar="AHMASK--AHmasked_weight"
               Maskarg="conn_cmfe(<AHmasked_weight.%s.h5[0]id:%s>, <Carpet AMR-grid>)" % (plane,MaskVar)
               Maskexpr="if( gt(%s,0), 1, 0 )" % MaskArg
               expression="%s*Mask" % Vararg
               exprname="%sMasked_%s" % (var,plane)
        else:
               expression=Vararg
               exprname="%s_%s" % (var, plane)
        if OnlyOverAtmo == 'yes':
                AtmoExpr="conn_cmfe(<rho.%s.h5[0]id:HYDROBASE--rho>, <Carpet AMR-grid>)" % plane
                if AtmoVal == "rho_min":
                        atmo=1e-13 # TODO read from parameter file
                else:
                        atmo=AtmoVal

        if newTitle:
           plotTitle=newTitle
        else:
           plotTitle="%s(%s)" %(var, plane)
        
        symmetry=''
        if planex == 'x' and 0 == os.system("egrep -i --quiet '^[[:space:]]*ReflectionSymmetry::reflection_x[[:space:]]*=[[:space:]]*\"?yes\"?' %s" % PFILE):
           symmetry='x'
        if 0 == os.system("egrep -i --quiet '^[[:space:]]*ReflectionSymmetry::reflection_y[[:space:]]*=[[:space:]]*\"?yes\"?' %s" % PFILE):
           if planex == 'y' or planey == 'y':
              symmetry += 'y'
        if planey == 'z' and 0 == os.system("egrep -i --quiet '^[[:space:]]*ReflectionSymmetry::reflection_z[[:space:]]*=[[:space:]]*\"?yes\"?' %s" % PFILE):
           symmetry += 'y'
        
        #if egrep -i --quiet '^[[:space:]]*RegridBoxes::rotatingsymmetry180[[:space:]]*=[[:space:]]*"?yes"?' $PFILE ; then
        #   if [ $plane == "xy" ]; then
        #      symmetry=180
        #   elif [ $plane == "xz" ]; then
        #      symmetry=x${symmetry}
        #   fi
        #fi
        #if egrep -i --quiet '^[[:space:]]*RegridBoxes::rotatingsymmetry90[[:space:]]*=[[:space:]]*"?yes"?' $PFILE ; then
        #   if [ $plane == "xy" ]; then
        #      symmetry=90
        #   elif [ $plane == "yz" ]; then
        #      symmetry=x${symmetry}
        #   fi
        #fi
        
        print "variable=%s, symmetry=%s" % (variable, symmetry)
        print "We are to plot the following expression: \"%s\"" % expression
        
        moviedir="%s/frames_%s" % (dirname,exprname)
        moviefile="%s/%s" % (dirname,exprname)
        try:
                os.makedirs(moviedir)
        except os.error:
                pass
        if not FrameBase:
           frames="%s/%s" % (moviedir, exprname)
        else:
           frames="%s/%s" % (moviedir, FrameBase)
           moviefile="%s/%s" % (dirname, FrameBase)
                
        #################################

        for dirname in sorted(glob.glob('%s/output-[0-9][0-9][0-9][0-9]/%s' % (dirname,pardir))):
                h5file="%s/%s.h5" % (dirname,filebase)
                print "Processing %s..." % h5file
                
                if not os.path.exists(h5file):
                    h5file="%s/data/H5_2d/%s.h5" % (dirname,filebase)
                    print "Processing %s..." % h5file
                    if not os.path.exists(h5file):
                        print "ERROR: There is no %s.h5 in directory %s.  Skipping." % (filebase, dirname)
                        continue
                
                # find out which frames are new
                filespresent=glob.glob("%s.[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9].png" % frames)
                f = h5py.File(h5file)
                its = dict()
                for obj in f:
                        m = re.search(' it=0*([0-9]+)', obj)
                        if m:
                                it = int(m.group(1))
                        else:
                                m = re.search('^it0*([0-9]+)', obj)
                                if m:
                                        it = int(m.group(1))
                        if not "%s.%09d.png" % (frames,it) in filespresent:
                                its[it] = 1
                f.close()
                its = sorted(its.keys())
                
                if its:
                        print "We found stored frames until %d. Continuing to create frames until %d." % (its[0], its[-1])
                        ## If h5file is of the format <thorn>::<var>.x[yz].h5, then we need to add a localhost: for visit to not give up
                        h5file="localhost:" + h5file
                        RenderFrames(symmetry, range_max, planex, planey, h5file,
                                        MASKED, OnlyOverAtmo, expression, exprname,
                                        ExprThresh, var_min, var_max, VarScaling,
                                        colortable, plotTitle, CONTOURS, outputsize,
                                        its, frames)
                else:
                  print "No new frames in %s." % h5file
        MakeMovie(frames, TYPE, moviefile, MinNumFrames)
        print "All done."

def RenderFrames(symmetry,range_max, planex, planey, h5file, MASKED, OnlyOverAtmo, expression, exprname, ExprThresh, var_min, var_max, VarScaling, colortable, plotTitle, CONTOURS, outputsize, its, frames):
        global visit_dir

        # This is done early since SetOperatorAttributes seems to ignore the all flag
        if symmetry != '':
           clip = ClipAttributes()
           clip.planeInverse = 1
           r = ReflectAttributes()
           r.useXBoundary = 0
           r.useYBoundary = 0
           r.useZBoundary = 0
           if symmetry == 'x':
              clip.plane1Status = 1
              clip.plane2Status = 0
              clip.plane1Origin = (OFudge,0,0)
              r.reflections = (1,1,0,0,0,0,0,0)
           elif symmetry == 'y':
              clip.plane1Status = 0
              clip.plane2Status = 1
              clip.plane2Origin = (0,OFudge,0)
              r.reflections = (1,0,1,0,0,0,0,0)
           elif symmetry == 'xy':
              clip.plane1Status = 1
              clip.plane2Status = 1
              clip.plane1Origin = (OFudge,0,0)
              clip.plane2Origin = (0,OFudge,0)
              r.reflections = (1,1,1,1,0,0,0,0)
           elif symmetry == '180':
              clip.plane1Status = 1
              clip.plane2Status = 0
              clip.plane1Origin = (OFudge,0,0)
              r.reflections = (1,0,0,1,0,0,0,0)
        
        sl = visit.SliceAttributes()
        sl.axisType = sl.ZAxis
        
        mvmax_str = range_max
        if(mvmax_str.find(':') == -1): # just a single number
            mvmax_list = [[0,float(mvmax_str)]]
        else:
            ranges = mvmax_str.split(',')
            mvmax_list = []
            for range_str in ranges:
                mvmax_list.append(map(float, range_str.split(':')))
        # later on it is more convenient to have the list sorted from the last timestep
        # to the first (so as to easier find the first entry smaller or equal to the
        # current)
        mvmax_list.reverse()
        
        u = visit.AnnotationAttributes()
        u.userInfoFlag=0
        u.axes2D.xAxis.title.title='%s/M' % planex
        u.axes2D.yAxis.title.title='%s/M' % planey
        u.axes2D.xAxis.label.font.scale=0.75
        u.axes2D.yAxis.label.font.scale=0.75
        u.axes2D.xAxis.title.font.scale=0.75
        u.axes2D.yAxis.title.font.scale=0.75
        u.axes2D.xAxis.title.font.font=u.axes2D.xAxis.title.font.Times
        u.axes2D.yAxis.title.font.font=u.axes2D.yAxis.title.font.Times
        u.databaseInfoFont.scale=0.75
        u.databaseInfoFont.font=u.databaseInfoFont.Times
        visit.SetAnnotationAttributes(u)
        
        pth = visit.ThresholdAttributes()
        
        visit.OpenDatabase(h5file)
        # Create expression
        if ( MASKED == "yes" ):
           maskform=Maskexpr
           visit.DefineScalarExpression("Mask",maskform)
           pth.listedVarNames=("Mask")
           pth.lowerBounds=(1e-20)
           pth.zonePortions=(0) # only include cells if fully above threshold
        
        # Set up atmosphere
        if ( OnlyOverAtmo == "yes" ):
           atmoform=AtmoExpr
           visit.DefineScalarExpression("Atmo",atmoform)
           pth.listedVarNames = pth.listedVarNames + ("Atmo",)
           pth.lowerBounds = pth.lowerBounds + (atmo,)
           visit.ActivateDatabase(h5file)
        
        exprform=expression
        visit.DefineScalarExpression(exprname,exprform)
        
        # Set up a (min) threshold on the expression itself
        if ExprThresh:
           numthresh=len(pth.listedVarNames)
           pth.listedVarNames = pth.listedVarNames + (variable,)
           pth.lowerBounds = pth.lowerBounda + (ExprThresh,)
        
        visit.AddPlot("Pseudocolor", exprname)
        visit.AddOperator("Project")
        visit.AddOperator("Threshold")
        visit.SetOperatorOptions(pth)
        visit.AddOperator("Box") # will be modified later during timestep loop
        mvmax = mvmax_list[-1][1]
        b = visit.BoxAttributes()
        b.maxx =  mvmax
        b.minx = -mvmax
        b.maxy =  mvmax
        b.miny = -mvmax
        b.maxz =  mvmax
        b.minz = -mvmax
        visit.SetOperatorOptions(b)
        v = visit.GetView2D()
        v.windowCoords = (-mvmax,mvmax,-mvmax,mvmax)
        visit.SetView2D(v)
        
        rho_min = var_min
        rho_max = var_max
        if rho_min == "auto" or rho_max == "auto":
           visit.DrawPlots() # activate plot for query below
           
           # find minimum and maximum of rho
           mymin = 1e37
           mymax = -1e37
           for state in range(0,TimeSliderGetNStates()-1,visit.TimeSliderGetNStates()/10):
              visit.SetTimeSliderState(state)
              visit.Query("MinMax",1,exprname) # 1=ActualData,0=OriginalData
              # sometimes the query will not find any data. We skip these occasions.
              try:
                tmpmin, tmpmax = visit.GetQueryOutputValue()
                if tmpmin <= mymin:
                   mymin = tmpmin
                if tmpmax >= mymax:
                   mymax = tmpmax
              except TypeError:
                continue
              except ValueError:
                continue
           visit.SetTimeSliderState(0)
           if rho_min == "auto":
              rho_min = str(mymin)
           if rho_max == "auto":
              rho_max = str(mymax)
        
        # modify the plot to look nicer
        p = visit.PseudocolorAttributes()
        p.centering = p.Zonal
        p.limitsMode = p.CurrentPlot
        if VarScaling == "Linear":
                p.scaling = p.Linear
        elif VarScaling == "Log":
                p.scaling = p.Log
        if rho_min != "none":
           p.minFlag = 1
           p.min     = float(rho_min)
        if rho_max != "none":
           p.maxFlag = 1
           p.max     = float(rho_max)
        p.colorTableName = colortable
        visit.SetPlotOptions(p)
        if symmetry != '':
           visit.AddOperator("Clip")
           visit.SetOperatorOptions(clip)
           visit.AddOperator("Reflect")
           visit.SetOperatorOptions(r)
        
        plottitle=visit.CreateAnnotationObject("Text2D")
        plottitle.position=(0.5,0.94)
        plottitle.fontFamily = plottitle.Times
        mintext=plotTitle
        plottitle.text=mintext[0:23]
        
        # Add contours if desired
        if CONTOURS:
           visit.AddPlot("Contour", exprname)
           visit.AddOperator("Threshold")
           visit.SetOperatorOptions(pth)
           c = visit.ContourAttributes()
           c.colorType   = c.ColorBySingleColor
           c.legendFlag  = 0
           if rho_min != "none":
              p.minFlag = 1
              p.min     = float(rho_min)
           if rho_max != "none":
              p.maxFlag = 1
              p.max     = float(rho_max)
           c.contourNLevels = 5
           if ContourScaling == 'Linear':
                   c.scaling     = c.Linear
           elif ContourScaling == 'Log':
                   c.scaling     = c.Log
           c.singleColor = (128,128,128,255)
           visit.SetPlotOptions(c)
           visit.AddOperator("Box")
           visit.SetOperatorOptions(b)
           if symmetry != '':
              visit.AddOperator("Clip")
              visit.SetOperatorOptions(clip)
              visit.AddOperator("Reflect")
              visit.SetOperatorOptions(r)
        
        visit.DrawPlots() # activate all plots for animation
        
        s = visit.SaveWindowAttributes()
        s.format = s.PNG
        s.outputDirectory = ""
        s.width, s.height = outputsize
        s.resConstraint = s.EqualWidthHeight
        s.screenCapture = 0
        s.family = 0
        visit.SetSaveWindowAttributes(s)
        
        times = visit.GetMetaData(visit.GetWindowInformation().activeSource).times
        cycles = visit.GetMetaData(visit.GetWindowInformation().activeSource).cycles
        old_mvmax = 0
        drawThePlots = False
        for state in range(0, visit.TimeSliderGetNStates()):
           if not cycles[state] in its:
                   continue
           # rescale plotting area if required
           for r in mvmax_list:
              if(r[0] <= times[state]):
                 mvmax = r[1]
                 break
           if mvmax != old_mvmax:
             b.maxx =  mvmax
             b.minx = -mvmax
             b.maxy =  mvmax
             b.miny = -mvmax
             b.maxz =  mvmax
             b.minz = -mvmax
             visit.SetOperatorOptions(b)
             v.windowCoords = (-mvmax,mvmax,-mvmax,mvmax)
             visit.SetView2D(v)
        
           # plot
           if ( visit.SetTimeSliderState(state)==0 ):
              drawThePlots = True
           if drawThePlots: # forcible replot when an error occured
              if not visit.DrawPlots():
                 print "VisIt could not draw plots of %s for state %d" % (exprname,state)
              else:
                 drawThePlots = False
        
           s.fileName = "%s.%09d" % (frames, cycles[state])
           visit.SetSaveWindowAttributes(s)
           visit.SaveWindow()
        visit.DeleteAllPlots()
        visit.CloseDatabase(h5file)
        visit.CloseComputeEngine()
        
def MakeMovie(frames, TYPE, moviefile, MinNumFrames):
        ############################################
        ## Check generated frames and make frames ##
        ############################################
        
        if TYPE == 'gif':
                for png in glob.glob("%s*.png" %frames):
                        gif=png[0:-4] + ".gif" # I am sure there is a nicer way...
                        if not os.path.exists(gif):
                                os.system("convert %s %s" % (png,gif))
        
        #####################
        ## Generate movie. ##
        #####################
        
        if TYPE == 'gif':
            os.system("gifsicle -l --delay 10 --colors 256 %s[0-9]*.gif > %s.gif" % (frames, moviefile))
        elif TYPE == 'mp4' or TYPE == 'ogg':
            loop=""
            if MinNumFrames:
                nframes=len(glob.glob("%s.[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9].png" % frames))
                if nframes > 0 and nframes < MinNumFrames:
                    # round up to nearest multiple of nframes from MinNumFrames
                    loopedframes=( (MinNumFrames + nframes - 1) / nframes) * nframes
                    loop="-loop_input -vframes %s" % loopedframes
            framefiles = sorted(glob.glob("%s.*.png" % frames))
            # create symlinks to be able to feed continuously numbered frames to ffmpeg
            # because some versions of ffmpeg don't cook piped pngs
            for i, frame in enumerate(framefiles):
              link = "%s_frame%05d.png" % (frames, i)
              if (os.path.exists(link)):
                os.remove(link)
              os.symlink(frame.split("/")[-1], link)
            enc_opts = "-pix_fmt yuv420p  -vcodec libx264"
            if TYPE == 'ogg':
              enc_opts = "-vcodec libtheora -qscale 6"
            if framefiles:
                command ="ffmpeg -y %s -vcodec png -f image2 -i %s_frame%%05d.png %s -crf 22 -threads 0 -preset slow %s.%s 2>/dev/null" % (loop,frames,enc_opts,moviefile,TYPE)
                os.system(command)
                # copy one frame "out of the middle" to represent movie
                shutil.copyfile(framefiles[len(framefiles)/2], moviefile + ".png")
                print "Encoding done to file %s.mp4" % moviefile
            # remove symlinks again
            for i, frame in enumerate(framefiles):
              link = "%s_frame%05d.png" % (frames, i)
              if (os.path.exists(link)):
                os.remove(link)

main()