"""

clawplotting:  Clawpack plotting routines in python
Uses matplotlib which requires numpy and pylab

"""
# ----------------------------------------------------------


import os,sys
from numpy import *

# to work with EagleClaw:
import matplotlib
matplotlib.use('Agg')  # Use an image backend
#import pylab
from pylab import *

from clawtools import *


    
#==================================================================
def plotframe1(frame=0, clawframe=None, outdir='.', refresh=True):
#==================================================================
    """

    Makes a 1d plot of corresponding frame and data clawframe

    If clawframe (of class ClawFrame) is already loaded then the data in
    this frame is plotted (for which frame = clawframe.frame).

    If clawframe is None, then clawframe is created and data read in from
    files fort.t000N and fort.q000N where N = frame value passed in.
    In this case the value of outdir is needed to tell where to find the data.

    If refresh == False then the new plot is added to an existing plot,
    otherwise a clf() command is given before plotting.

    """

    if clawframe is not None:
        # frame data has already been read into clawframe of class ClawFrame
        frame = clawframe.frame   # frame number
    else:
        # need to read in data for this frame
        clawframe = read_clawframe(frame,outdir)
        
    # For 1d, no AMR so there is only one grid:
    grid = clawframe.grids[0]

    t = clawframe.t

    if refresh:
        clf()
    else:
        hold(True)

    # set default plot parameter values:
    clawplotdata = setplot1()

    # values may be redefined in local setplot1user.py:
    if os.path.exists("setplot1user.py"):
        file = open('setplot1user.py','r')
        exec(file)
        file.close()
        if (setplot1user):
            clawplotdata = setplot1user(clawplotdata)
        else:
            print "Local file setplot1user.py exists but does not define"
            print "  a function setplot1user"

    # the following plot parameters should now be set:

    mq = clawplotdata.mq
    axis_limits = clawplotdata.axis_limits
    user_variable = clawplotdata.user_variable
    user_variable_name = clawplotdata.user_variable_name
    mapped_grid = clawplotdata.mapped_grid
    plotstyle = clawplotdata.plotstyle


    if os.path.exists("beforeframe.py"):
        file = open("beforeframe.py",'r')
        exec(file)
        file.close()


    if mapped_grid:
        fname = os.path.join(os.getcwd(),"mapc2p.py")
        print "fname = ",fname
        try:
            file = open(fname,'r')
            exec(file)
            file.close()
        except:
            print "mapped_grid = True but no mapc2p.py file found in directory"
            print "     ",os.getcwd()
            sys.exit(1)
        xp = mapc2p(grid.xcenter)
    else:
        xp = grid.xcenter

    if user_variable:
        try:
            file = open("%s.py" % user_variable_name, 'r')
            exec(file)
            file.close()
        except:
            print "user_variable = True and user_variable_name = ",\
                   user_variable_name
            print "  but file %s.py not found" % user_variable_name
            sys.exit(1)
        exec('q = %s(grid.q,xp,t)'  % user_variable_name)
        varname = user_variable_name
        mq = [-1]

    nplots = len(mq)
    for n in range(nplots):
        subplot(nplots,1,n+1)
        if not user_variable:
            q = grid.q[:,mq[n]-1]
            varname = 'q(%g)'  % mq[n]
        plotcommand = "plot(xp,q,'%s')"  % plotstyle
        exec(plotcommand)
        title("%s at time %8.4f" % (varname,t))
        axis(axis_limits)

    if os.path.exists("afterframe.py"):
        execfile("afterframe.py")

    return clawframe


#==================================================================
def plotframe2(frame=0, clawframe=None, outdir='.', refresh=True):
#==================================================================
    """

    Makes a 2d plot of corresponding frame and data clawframe

    If clawframe (of class ClawFrame) is already loaded then the data in
    this frame is plotted (for which frame = clawframe.frame).

    If clawframe is None, then clawframe is created and data read in from
    files fort.t000N and fort.q000N where N = frame value passed in.
    In this case the value of outdir is needed to tell where to find the data.

    If refresh == False then the new plot is added to an existing plot,
    otherwise a clf() command is given before plotting.

    """

    if clawframe is not None:
        # frame data has already been read into clawframe of class ClawFrame
        frame = clawframe.frame   # frame number
    else:
        # need to read in data for this frame
        clawframe = read_clawframe(frame,outdir)
        

    t = clawframe.t

    if refresh:
        clf()
    hold(True)

    # set default plot parameter values:
    clawplotdata = setplot2()

    # values may be redefined in local setplot2user.py:
    if os.path.exists("setplot2user.py"):
        file = open('setplot2user.py','r')
        exec(file)
        file.close()
        if (setplot2user):
            clawplotdata = setplot2user(clawplotdata)
        else:
            print "Local file setplot2user.py exists but does not define"
            print "  a function setplot2user"

    # the following plot parameters should now be set:

    mq = clawplotdata.mq
    axis_limits = clawplotdata.axis_limits
    plot_grid = clawplotdata.plot_grid
    grid_color = clawplotdata.grid_color
    plot_grid_edges = clawplotdata.plot_grid_edges
    grid_edge_color = clawplotdata.grid_edge_color
    user_variable = clawplotdata.user_variable
    user_variable_name = clawplotdata.user_variable_name
    mapped_grid = clawplotdata.mapped_grid
    plot_contour = clawplotdata.plot_contour
    contour_levels = clawplotdata.contour_levels
    min_contour = clawplotdata.min_contour
    max_contour = clawplotdata.max_contour
    contour_colors = clawplotdata.contour_colors
    plot_pcolor = clawplotdata.plot_pcolor
    pcolor_map = clawplotdata.pcolor_map
    c_min = clawplotdata.c_min
    c_max = clawplotdata.c_max
    plot_colorbar = clawplotdata.plot_colorbar

    if os.path.exists("beforeframe.py"):
        file = open("beforeframe.py",'r')
        exec(file)
        file.close()

    if (not plot_contour) & (not plot_pcolor):
        print "The solution will not be plotted since"
        print "   plot_contour and plot_pcolor are both False"

    if mapped_grid:
        try:
            file = open("mapc2p.py",'r')
            exec(file)
            file.close()
        except:
            print "mapped_grid = True but no mapc2p.py file found in directory"
            print "     ",os.getcwd()
            sys.exit(1)


    for gridno in range(clawframe.ngrids):
        grid = clawframe.grids[gridno]
        level = grid.level   # AMR level of this grid

        if user_variable:
            try:
                execfile(user_variable_name)
            except:
                print "user_variable = True and user_variable_name = ",\
                       user_variable_name
                print "  but file %s.py not found" % user_variable_name
                sys.exit(1)


        xc_center = grid.xcenter
        yc_center = grid.ycenter
        [X_center,Y_center] = meshgrid(xc_center,yc_center)
        if mapped_grid:
            X_center,Y_center = mapc2p(X_center,Y_center)

        xc_edge = grid.xedge
        yc_edge = grid.yedge
        [X_edge,Y_edge] = meshgrid(xc_edge,yc_edge)
        if mapped_grid:
            X_edge,Y_edge = mapc2p(X_edge,Y_edge)

        if user_variable:
            exec('q = %s(grid.q,X_center,Y_center,t)'  % user_variable_name)

        if plot_grid[level-1]==2:
            for i in range(X_edge.shape[0]):
                plot(X_edge[i,:], Y_edge[i,:], grid_color)
            for i in range(X_edge.shape[1]):
                plot(X_edge[:,i], Y_edge[:,i], grid_color)

        if plot_grid_edges[level-1]:
            for i in [0, X_edge.shape[0]-1]:
                X1 = X_edge[i,:]
                Y1 = Y_edge[i,:]
                plot(X1, Y1, grid_edge_color)
            for i in [0, X_edge.shape[1]-1]:
                X1 = X_edge[:,i]
                Y1 = Y_edge[:,i]
                plot(X1, Y1, grid_edge_color)

        if not user_variable:
            q = grid.q[:,:,mq-1]
            varname = 'q(%g)'  % mq

        if plot_pcolor:
            if plot_grid[level-1]:
                pobj = pcolor(X_edge, Y_edge, transpose(q), \
                    cmap=pcolor_map, edgecolors=grid_color, shading='faceted')
            else:
                pobj = pcolor(X_edge, Y_edge, transpose(q), \
                    cmap=pcolor_map, edgecolors='None', shading='flat')
            clim(c_min,c_max)

        if plot_contour:
            contour(X_center, Y_center, transpose(q), \
                    contour_levels, colors=contour_colors)

    if plot_pcolor:
        if plot_colorbar:
            #print "colorbar doesn't work reliably.... not plotting"
            colorbar(pobj)
            #colorbar(pad=0.1)

    title("%s at time %8.4f" % (varname,t))
    axis(axis_limits)

    hold(False)

    if os.path.exists("afterframe.py"):
        execfile("afterframe.py")

    return clawframe



#--------------
def setplot1():
#--------------
    from clawtools import ClawData
    clawplotdata = ClawData()
    clawplotdata.mq = [1] 
    clawplotdata.axis_limits = 'auto'

    clawplotdata.user_variable = False
    clawplotdata.user_variable_name = ""
    clawplotdata.mapped_grid = False
    clawplotdata.plotstyle = 'ro'
    return clawplotdata

#------------------------------
def setplot1user(clawplotdata):
#------------------------------

    """
    Default library version does nothing.
    Create a file setplot1user.py containing a modified function if you
    want to change parameters.  This will be called after setplot1.
    """
    return clawplotdata

#--------------
def setplot2():
#--------------
    from pylab import cm    # colormaps
    from clawtools import ClawData
    clawplotdata = ClawData()

    clawplotdata.mq = 1 
    clawplotdata.axis_limits = 'auto'

    clawplotdata.user_variable = False
    clawplotdata.user_variable_name = ""
    clawplotdata.mapped_grid = False

    maxlevels = 10
    clawplotdata.plot_grid = zeros(maxlevels) # don't plot grid lines
    clawplotdata.grid_color = 'k'       # color to plot grid lines
    clawplotdata.plot_grid_edges = ones(maxlevels)  # plot grid edges
    clawplotdata.grid_edge_color = 'k'       # color to plot grid edges

    clawplotdata.plot_pcolor = True       # pcolor plot?
    clawplotdata.pcolor_map = cm.RdYlBu   # for pcolor plots - set a colormap
    clawplotdata.c_min  = 0
    clawplotdata.c_max = 4
    clawplotdata.plot_colorbar = False

    clawplotdata.plot_contour = True      # contour plot?
    clawplotdata.contour_levels = 30      # for specified number of contour lines
                                          # can instead set to a vector of values
    clawplotdata.min_contour = 0.0        # or can set these to min and
    clawplotdata.max_contour = 1.0        # max values, then set contour_levels
                                          # based on these
    clawplotdata.contour_colors = 'k'     # colors parameter in contour plot

    return clawplotdata

#------------------------------
def setplot2user(clawplotdata):
#------------------------------

    """
    Default library version does nothing.
    Create a file setplot2user.py containing a modified function if you
    want to change parameters.  This will be called after setplot2.
    """
    return clawplotdata