#!/usr/bin/env python
#
# Python script to run xclaw with several different parameter choices.
#

import shelve
import os
from clawplotting import *
from clawtools import *

t0 = os.times()   # keep track of CPU time 

#-------------------------------------
def runtests(itest=1, plotsoln=False):
#-------------------------------------
    '''
    ================================
    Main code to run a set of tests:
    ================================
    '''

    if plotsoln:
        from plotclaw import plotclaw
    


    # set data values:
    data = ClawData()
    data.tfinal = 0.75     # final time

    if plotsoln:
        data.nout = 3          # output at several times for plotting
        data.plot_contour = True
        data.plot_grid = 0
        data.write('setplot2.data')
    else:
        data.nout = 1          # only output at final time
    
    if itest==1:
        # Test only coarse Cartesian grids, runs quickly
        grids = [1]            # set of grids to test
        limiters = [0,3]         # set of limiters (mthlim values) to test
        mxvals_polar = array([10,20])        # mx values on polar grid
        myvals_polar = array([75,150])      # my values on polar grid
        mxvals_other = array([30,60])       # mx values on other grids
        myvals_other = array([30,60])       # my values on other grids
    elif itest==2:
        # Do a few more tests
        grids = [1,3]        # set of grids to test
        limiters = [0,3]       # set of limiters (mthlim values) to test
        mxvals_polar = array([10,20,40])       # mx values on polar grid
        myvals_polar = array([75,150,300])    # my values on polar grid
        mxvals_other = array([30,60,120])     # mx values on other grids
        myvals_other = array([30,60,120])     # my values on other grids
    elif itest==3:
        # Full test of all grids
        grids = [1,2,3]        # set of grids to test
        limiters = [0,3]       # set of limiters (mthlim values) to test
        mxvals_polar = array([10,20,40,80])       # mx values on polar grid
        myvals_polar = array([75,150,300,600])    # my values on polar grid
        mxvals_other = array([30,60,120,240])     # mx values on other grids
        myvals_other = array([30,60,120,240])     # my values on other grids
    else:
        print "Unrecognized input:  itest = ", itest
        sys.exit(1)

    print "------------------------------------------------------------------"

    
    table = {}  # dictionary of data and results for each test
    
    for mthlim in limiters:
        data.mthlim = mthlim
        for igrid in grids:
            data.igrid = igrid
        
            if igrid==2:
                # for polar grid:
                mxvals = mxvals_polar
                myvals = myvals_polar
                data.xlower = 0.2
                data.xupper = 1.0
                data.ylower = 0
                data.yupper = 2*pi
                data.mthbc_ylower = 2
                data.mthbc_yupper = 2
                area = (1 - data.xlower**2)*pi 
            else:
                # for other grids:
                mxvals = mxvals_other
                myvals = myvals_other
                data.xlower = -1.0
                data.xlower = -1.0
                data.xupper = 1.0
                data.ylower = -1.0
                data.yupper = 1.0
                data.mthbc_ylower = 1
                data.mthbc_yupper = 1
                area = pi
        
            table[(mthlim,igrid)] = {} # dictionary to hold data
                                       # and results for this test
            this_table = table[(mthlim,igrid)]  # short name
            this_table['mxvals'] = mxvals       # grid resolutions to test
            this_table['myvals'] = myvals
            this_table['ave_cell_area'] = area / (mxvals*myvals)
            this_table['tfinal'] = data.tfinal
            this_table['errors'] = empty(len(mxvals))
            
            for itest in range(len(mxvals)):
                mx = mxvals[itest]
                my = myvals[itest]
                data.mx = mx
                data.my = my
                data.write('claw2ez.data')
                data.write('setprob.data')

                print "-------------------------------------------------------"
                print 'Running code with igrid = %i, mx = %g, my = %g ...' \
                                      % (igrid,mx,my)
            
                odir = '.'
                # or use next line to put each set of output in separate subdir:
                # odir = 'output_lim%ggrid%gmx%g' % (mthlim,igrid,mx)

                # run fortran code:
                t1 = os.times()
                runclaw(outdir=odir, overwrite=True)
		t2 = os.times()
                CPUtime = t2[2]-t1[2] + t2[3]-t1[3]
                print "CPU time = ",CPUtime
            
                # compute errors:
                errors = compute_errors(odir, frame=data.nout)

                # max-norm of error:
                errormax = abs(errors).max()   

                # approx 1-norm of error:
                errorsum = abs(errors).sum() 
                error1 = errorsum * this_table['ave_cell_area'][itest]

                this_table['errors'][itest] = error1
            
                if plotsoln:
                    # put each set of plots in a different directory:
                    pdir = 'plots_lim%ggrid%gmx%g' % (mthlim,igrid,data.mx)
                    plotclaw(outdir=odir, plotdir=pdir, overwrite=True, \
                         indexentry='mthlim = %2i, igrid = %2i, mx = %4i, \
                         my = %4i' % (mthlim,igrid,mx,my))


    # Save results by shelving them.
    # The results can later be accessed from a different Python process
    # if desired:
    db = shelve.open('table.db')
    db['tfinal'] = data.tfinal
    db['limiters'] = limiters
    db['grids'] = grids
    db['table'] = table
    db.close
    
    # output tables:
    make_text_table(table, limiters, grids, fname='errortables.txt')
    make_latex_table(table, limiters, grids, fname='errortables.tex')

    # log-log plot of errors:
    make_error_plots(table, limiters, grids, fname='errors.png')


#------------------
def qtrue(x,y,t):
#------------------
    # true solution, rotating flow and cosine hump:

    # parameters:
    data = ClawData('setprob.data')
    x1 = data.x1
    x2 = data.x2
    y1 = data.y1
    y2 = data.y2
    xm = (x1 + x2) / 2
    ym = (y1 + y2) / 2
    xscale = 2*pi / (x2-x1)
    yscale = 2*pi / (y2-y1)

    # rotate (x,y) back to time 0:
    tr = 2*pi*t
    x0 = cos(tr)*x + sin(tr)*y
    y0 = -sin(tr)*x + cos(tr)*y

    # evaluate q(x,y,t=0):
    coshump = (1 + cos((x0-xm)*xscale)) * (1 + cos((y0-ym)*yscale))
    inhump = (x0 > x1) & (x0 < x2) & (y0 > y1) & (y0 < y2)
    qtrue = where(inhump, coshump, 0)
    return qtrue



#------------------------------------
def compute_errors(odir='.',frame=0):
#------------------------------------
    # change to output directory and read in solution from fort.q000N
    # for frame N.

    # read in clawpack solution for this frame:
    os.chdir(odir)
    clawframe = read_clawframe(frame)
    
    t = clawframe.t
    
    # Assume there is only one grid (AMR not used):
    grid = clawframe.grids[0]
    
    print "Computing errors for mx, my = ",grid.mx,grid.my
    
    xc = grid.xcenter
    yc = grid.ycenter
    Xc,Yc = meshgrid(xc,yc)
    try:
        file = open('mapc2p.py')
        exec(file)
        file.close()
    except:
        print "Error: mapc2p file not found"
        sys.exit(1)
    
    Xp,Yp = mapc2p(Xc,Yc)
    qq = qtrue(Xp,Yp,t)
    q = transpose(grid.q[:,:,0])
    errors = qq - q
    return errors

    
#-------------------------------------------------------------------
def make_text_table(table, limiters, grids, fname='errortables.txt'):
#-------------------------------------------------------------------
    errfile = open(fname,'w')
    errfile.write('Errors computed by clawtest.py on %s \n\n' \
                   % current_time())
    for mthlim in limiters:
        errfile.write('\n\nLimiter = %g\n===========\n'  % mthlim)
        for igrid in grids:
            errfile.write('\n\nGrid Number %g\n-------------\n'  % igrid)
            errfile.write('  mx    my      ave Delta x         error  ')
            errfile.write('      observed order\n\n')

            # table values for this table:
            te = table[(mthlim,igrid)]   
            for itest in range(len(te['mxvals'])):
                mx = te['mxvals'][itest]
                my = te['myvals'][itest]
                ave_cell_area = te['ave_cell_area'][itest]
                ave_dx = sqrt(ave_cell_area)
                errornorm = te['errors'][itest]
                if itest>0:
                    E1 = te['errors'][itest-1]
                    E2 = te['errors'][itest]
                    A1 = te['ave_cell_area'][itest-1]
                    A2 = te['ave_cell_area'][itest]
                    porder = 2*log(E1/E2)/log(A1/A2)
                else:
                    porder = nan
                errfile.write('%4i  %4i  %14.4e  %16.4e  %16.4e\n' \
                            % (mx, my, ave_dx, errornorm, porder))

    final_time = os.times()
    python_user_time = final_time[0] - t0[0]
    python_system_time = final_time[1] - t0[1]
    fortran_user_time = final_time[2] - t0[2]
    fortran_system_time = final_time[3] - t0[3]
    python_CPUtime = python_user_time + python_system_time
    fortran_CPUtime = fortran_user_time + fortran_system_time
    errfile.write('\n\nTotal python CPU time: %g seconds\n\n' % python_CPUtime)
    errfile.write('\n\nTotal fortran CPU time: %g seconds\n\n' % fortran_CPUtime)
    errfile.close()

    print "------------------------------------------------------------------"
    print "Table of errors is in file ",   fname
    print "------------------------------------------------------------------"
    

#-------------------------------------------------------------------
def make_latex_table(table, limiters, grids, fname='errortables.tex'):
#-------------------------------------------------------------------

    errfile = open(fname,'w')
    errfile.write(r'% ')
    errfile.write(' Errors computed by clawtest.py on %s \n\n' \
                   % current_time())
    
    for igrid in grids:
        t0 = table[(0,igrid)]   # limiter mthlim = 0
        t3 = table[(3,igrid)]   # limiter mthlim = 3
        errfile.write(r'''

          \begin{table}[th]
          \caption{Errors on Grid %i }
          \begin{center}
          \begin{tabular}{|ccc|c|c|c|c|}
          \hline
          &&&\multicolumn{2}{c|}{Limiter = 0}
            &\multicolumn{2}{c|}{Limiter = 3}\\ \hline
          mx & my & Ave $\Delta x$& Error & Observed order
                                  & Error & Observed order\\
          \hline
        '''  % igrid)
        errfile.write('\n')

        # lines in table for each grid resolution:
        for itest in range(len(t0['mxvals'])):
            mx = t0['mxvals'][itest]
            my = t0['myvals'][itest]
            ave_cell_area = t0['ave_cell_area'][itest]
            ave_dx = sqrt(ave_cell_area)
            errornorm0 = t0['errors'][itest]
            if itest>0:
                E1 = t0['errors'][itest-1]
                E2 = t0['errors'][itest]
                A1 = t0['ave_cell_area'][itest-1]
                A2 = t0['ave_cell_area'][itest]
                porder0 = 2*log(E1/E2)/log(A1/A2)
            else:
                porder0 = nan
            errornorm3 = t3['errors'][itest]
            if itest>0:
                E1 = t3['errors'][itest-1]
                E2 = t3['errors'][itest]
                A1 = t3['ave_cell_area'][itest-1]
                A2 = t3['ave_cell_area'][itest]
                porder3 = 2*log(E1/E2)/log(A1/A2)
            else:
                porder3 = nan
            errfile.write('%4i & %4i & %14.4f & %16.4f & %16.2f' \
                        % (mx, my, ave_dx, errornorm0, porder0))
            errfile.write(' & %16.4f & %16.2f \\\\ \n' \
                        % (errornorm3, porder3))


        errfile.write(r'''
          \hline
          \end{tabular}
          \end{center}
          \end{table}

        ''')

    errfile.close()

    print "------------------------------------------------------------------"
    print "Latex table of errors is in file ",  fname
    print "------------------------------------------------------------------"
    
    
#-------------------------------------------------------------------
def make_error_plots(table, limiters, grids, fname='errors.png'):
#-------------------------------------------------------------------
    clf()
    hold(True)
    legendtext = []
    # perform plots in proper order for legend, from worst to best:
    for (igrid,mthlim) in [(3,0),(1,0), (3,3),(1,3), (2,0),(2,3)]:
        try:
            # table entries for this test case:
            te = table[(mthlim,igrid)]   
            dx = sqrt(te['ave_cell_area'])
            loglog(dx,te['errors'],'o-')
	    legendtext.append('Grid %i, Limiter %i' % (igrid,mthlim))
        except:
            pass   # if this case not present
    axis([2e-3,1e-1,1e-4,1])
    xlabel('Average Delta x', fontsize=15)
    ylabel('1-norm of Error', fontsize=15)
    title('Log-log plot of errors', fontsize=15)
    legend(legendtext,'upper left')
    hold(False)
    savefig(fname)


    print "------------------------------------------------------------------"
    print "Plot of errors is in file ",   fname
    print "------------------------------------------------------------------"

if __name__ == '__main__':
    import string
    # execute runtests with any command line arguments:
    cmd = 'runtests(%s)' % string.join(sys.argv[1:])
    exec(cmd)