#
# plotclaw.py  -- python plotting routine for CLAWPACK.
# Assumes that matplotlib is used for plotting and the required
# modules are on the path specified by the environment variable
# PYTHONPATH.
#
# Usage:
#   plotclaw()                                 creates plots in current directory
#   plotclaw(plotdir="NEWDIR")                 creates plots in new directory
#   plotclaw(plotdir="DIR", overwrite=True)    creates plots in DIR, replacing
#                                               previous plots if present 
# Other optional keyword arguments:
#   outdir="OUTDIR"                reads fort.* files from directory OUTDIR
#   movie=True                     also creates an animated gif using 'convert'
#   savehtmls=True                 saves html versions of source files
#   savecode=True                  saves code files in directory outdir
#   latex=True                     also creates plots.tex
#   indexentry=''                  label for entry in index.html
#

import sys,os,glob,time,shutil
from clawplotting import *


#------------------------------------------------------------
# Plotting commands for each frame:
#
#==========================================
def create_plotframe(outdir,frame,pngfile):
#==========================================
    # Create a png figure for the given frame 
    clawframe = read_clawframe(frame, outdir=outdir)
    ndim = clawframe.ndim
    clf() 
    hold(False)
    if ndim==1:
        plotframe1(clawframe=clawframe)
    elif ndim==2:
        plotframe2(clawframe=clawframe)
    else:
        print "Cannot handle ndim = ",ndim
    savefig(pngfile)
    return float(clawframe.t)


#======================================================================
def plotclaw(outdir='.', plotdir='.', overwrite=False, movie=False, \
             savecode=False, savehtmls=False, latex=False, indexentry=''):
#======================================================================

    #print "outdir = ",outdir
    #print "plotdir = ",plotdir
    pyclaw = os.getenv('CLAW') + '/python'
    #sys.path.append(pyclaw)
    
    print " "
    print "Running    ", pyclaw + "/plotclaw.py"
    
    
    rootdir = os.path.abspath(os.getcwd())
    
    if plotdir != '.':
        if os.path.isfile(plotdir):
            print "Error: plotdir specified is a file"
            sys.exit(1)
        elif (os.path.isdir(plotdir) & overwrite):
            print "Directory ", plotdir, " already exists, "
            print "   will be removed and recreated..."
            try:
                shutil.rmtree(plotdir)
            except:
                print "Cannot remove directory ",plotdir
                sys.exit(1)
        elif (os.path.isdir(plotdir) & (not overwrite)):
            print "Directory ", plotdir, " already exists."
            print "Remove directory with 'rm -r ",plotdir,"' and try again,"
            print "  or use overwrite=True in call to plotclaw"
            sys.exit(1)

        try:
            os.mkdir(plotdir)
        except:
            print "Cannot make directory ",plotdir
            sys.exit(1)
    
    if savecode:
        try:
            os.system('make program')
        except:
            print "Couldn't make program.f"
    
    
    
    if plotdir != '.':
        # copy files to plotdir:
        if savehtmls:
            os.system('make htmls')
            for file in (glob.glob('*.f.html') + glob.glob('*.py.html') + \
                    glob.glob('README.html') + glob.glob('Makefile') + \
                    glob.glob('*.data.html')):
                shutil.copy(file,plotdir)
        if savecode:
            for file in (glob.glob('*.f') + glob.glob('*.py') + \
                        glob.glob('README') + \
                        glob.glob('*.data')):
                shutil.copy(file,plotdir)
    
    
    # determine number of space dimensions:
    clawframe = ClawFrame(0,outdir)
    clawframe.read_fortt()
    ndim = clawframe.ndim
    
    
    
    # Put path into form for printing
    rootPath = os.path.abspath('./')
    rootPath = rootPath.replace("\\","/")   # in case of cygwin
    
    
    # determine current time and reformat:
    time1 = time.asctime()
    year = time1[-5:]
    day = time1[:-14]
    hour = time1[-13:-5]
    creationtime = day + year + ' at ' + hour
    
    # remove any old versions:
    for file in glob.glob(os.path.join(plotdir,'frame*.png')) + \
                glob.glob(os.path.join(plotdir,'frame*.html')):
        os.remove(file)
    
    
    fortfile = {}
    pngfile = {}
    htmlfile = {}
    pdffile = {}
    frametime = {}
    
    # Create each of the figures
    os.chdir(outdir)
    for file in glob.glob('fort.q*'):
        frame = int(file[7:10])
        fortfile[frame] = file
        pngfile[frame] = 'frame' + file[-4:] + '.png'
        pdffile[frame] = 'frame' + file[-4:] + '.pdf'
        htmlfile[frame] = 'frame' + file[-4:] + '.html'
    
    if len(fortfile) == 0:
        print('No fort.q files found')
        sys.exit(1)
    
    # Figure out which files are from latest run:
    numFrames = 0
    mtime = 0
    for frame in range(len(fortfile)):
        mtimeprev = mtime
        mtime = os.path.getmtime(fortfile[frame])
        # sometimes later fort files are closed a few seconds after
        # earlier ones, so include a possible delaytime:
        delaytime = 5  # seconds
        if mtime < mtimeprev-delaytime:
            break
        numFrames = numFrames + 1
    
    os.chdir(rootdir)
    
    # Make png files:
    print "Making png files for figures..."
    for frame in range(numFrames):
        pngloc = os.path.join(plotdir,pngfile[frame])
        frametime[frame] = float(create_plotframe(outdir,frame,pngloc))
    
    
    # Write the index page
    
    # PlotIndex.html
    #---------------
    
    html = open(plotdir+'/PlotIndex.html','w')
    
    html.write('')
    html.write('
CLAWPACK Results\n' )
    html.write('
\n')
    html.write('
Plots created: %s    ' % creationtime + \
               ' On machine: %s\n'  % os.getenv('HOST'))
    html.write('
\n')
    html.write('
From directory: %s    ' % rootPath+ \
            'index.html' '\n' % rootPath)
    html.write('
\n')
    if plotdir != '.':
        html.write('
Saved to directory %s\n' % (plotdir))
        html.write('
\n')
    
    if savecode:
        html.write('
Code saved... full program: ' + \
                   ' clawprogram.f.html\n')
        html.write('
\n')
        
    html.write('
README.html
\n')
    html.write('
\n')
    
    if movie:
        html.write('\n Movie
\n')
    
    html.write('\n ALL Frames
\n')
    
    for frame in range(numFrames):
        ftime = frametime[frame]
        if ((ftime == 0) | ((ftime > 0.001) & (ftime < 1000))):
            html.write('Frame %i      ' \
                   'at time t = %9.5f 
\n' % \
                   (htmlfile[frame],frame,ftime))
        else:
            html.write('Frame %i      ' \
                   'at time t = %12.5e 
\n' % \
                   (htmlfile[frame],frame,ftime))
    
    html.write('
\n')
    html.write('\n ALL Frames
\n')
    html.write('')
    
    # PlotAllFrames.html
    #-------------------
    html = open(plotdir+'/PlotAllFrames.html', 'w')
    html.write('')
    html.write('
CLAWPACK Results\n' )
    html.write('
\n')
    html.write('
Plots created: %s\n' % (creationtime))
    html.write('
\n')
    html.write('
In directory %s\n' % (rootPath))
    html.write('
\n')
    if plotdir != '.':
        html.write('
Saved to directory %s\n' % (plotdir))
        html.write('
\n')
    html.write('
Plot Index ' +
            '    README.html\n')
    html.write('
\n')
    html.write('
All frames... Click on a figure to enlarge\n')
    html.write('
\n')
    html.write('
\n')
    
    for frame in range(0,numFrames,1):
        html.write('  \n' \
                % (htmlfile[frame], pngfile[frame]))
    
    html.write('\n')
    html.write('\n')
    html.close()
    
    
    # frameN.html
    #------------
    
    for frame in range(0,numFrames,1):
        html = open(os.path.join(plotdir,htmlfile[frame]),'w')
        html.write('\n')
        html.write('
Frame %i' % frame)
        #html.write('  from Directory %s\n' % (rootPath))
    
        # Write link commands
        html.write('''
                &#060; <    ''')
	if len(htmlfile)==1:
            html.write('''
                   Index  
                   ''' )
        elif frame==0:
            html.write('''
               &#060;     Index  
                    &#062;  ''' % htmlfile[frame+1])
        elif frame==numFrames-1:
            html.write('''
                &#060;       
               Index      &#062; ''' \
               % htmlfile[frame-1])
        else:
            html.write('''
                &#060;       
               Index      
                &#062;  ''' % (htmlfile[frame-1],htmlfile[frame+1]))
        html.write('''
                    &#062; > 
                \n''' % htmlfile[numFrames-1])
    
        html.write('''
                
                \n '''  % pngfile[frame])
        for frame in range(0,numFrames,1):
            html.write('''
               %i      ''' % (htmlfile[frame],frame))
        html.write('\n ALL   \n')
    
        html.write('')
        html.close()
    
    
    PlotPath = os.path.abspath(plotdir)
    if PlotPath[0] != '/':
        PlotPath = '/' + PlotPath
    PlotPath.replace('\\','/') # for windows
    print "Created png and html pages...  "
    print " "
    print "Point your browser to:"
    print "    ", "file://" + PlotPath + "/PlotIndex.html"
    print " "
    
    
    # index.html file listing all plot directories:
    # ---------------------------------------------

    if os.path.isfile('index.html'):
        # index.html already exists and we want to append to end:
        index = open('index.html','r')
        oldlines = index.read()
        index.close()
        index = open('index.html','w')
        index.write(oldlines)
    else:
        # index.html does not yet exist.  Make header:
        index = open('index.html','w')
        index.write('\n
\n')
        index.write('
CLAWPACK Results\n' )
        index.write('
\n')
        index.write('
In directory %s\n' % (rootPath))
        index.write('
\n')
        index.write("""
Note that some directories may have
                    been later overwritten.
\n
                    Remove plotindex.html before running plotclaw
                    to start a fresh index.\n""")

    index.write('
\n
\n' +\
                 plotdir+'/PlotIndex.html    ')
    index.write('(Created: %s)\n' % creationtime)
    index.write('
 '+indexentry + ' \n')
    index.close()

    
    
    # latex file of plots:
    # --------------------
    
    os.chdir(plotdir)
    if latex:
        #for frame in range(numFrames):
        #    os.system('convert %s %s' % (pngfile[frame],pdffile[frame]))
    
        latexfile = open('plots.tex', 'w')
    
        latexfile.write('''
    \\documentclass[11pt]{article}
    \\usepackage{graphicx}
    \\setlength{\\textwidth}{7.5in}
    \\setlength{\\oddsidemargin}{-0.5in}
    \\setlength{\\evensidemargin}{-0.5in}
    \\setlength{\\textheight}{9.2in}
    \\setlength{\\voffset}{-1in}
    \\setlength{\\headsep}{26pt}
    \\begin{document}
    \n''')
        latexfile.write('~\n\\vskip -0.5in')
        latexfile.write('\\centerline{\\Large\\bf CLAWPACK Results}\\vskip 5pt\n')
        latexfile.write('\\centerline{\\bf Plots created {\\tt %s}' \
                        % (creationtime))
        latexfile.write(' in directory:} \n\\vskip 5pt')
        latexfile.write('\\centerline{{\\tt %s}}\n'  % (rootPath))
        if plotdir != '.':
            plotdirtex = plotdir.replace('_','\\_')  # for _ in latex
            latexfile.write('\\vskip 5pt\n\\centerline{Saved to directory: {\\tt %s}}\n' \
                        % (plotdirtex))
        latexfile.write('\\vskip 5pt\n')
            
    
        for frame in range(0,numFrames,2):
            latexfile.write('\\includegraphics[width=3.5in]{%s}\n' \
                            % pngfile[frame])
            if frame+1 < numFrames:                  
                latexfile.write('\\includegraphics[width=3.5in]{%s}\n' \
                                % pngfile[frame+1])
            latexfile.write('\\vskip 10pt \n')                       
        latexfile.write('\\end{document}\n')                         
        latexfile.close()
        print("Latex file %s/plots.tex created (use pdflatex to make pdf file)" \
               % plotdir)
    
    
    # Movie:
    #-------
    
    if movie:
        print('Making movie.gif.  This may take some time....')
        os.system('convert -delay 20 frame*.png movie.gif')
    
    os.chdir(rootdir)
    
    
#================================================

if __name__ == '__main__':
    plotclaw()