"""Plot the errordata generated by test_analytic_light"""
#import io
import matplotlib.pyplot as plt
import cPickle
from dolfin import *
from matplotlib.backends.backend_pdf import PdfPages

__author__ = "Gabriel Balaban"
__copyright__ = "Copyright (C) 2012 Simula Research Laboratory and %s"%__author__
__license__  = "GNU GPL Version 3 or any later version"


PLOTSTYLES = {"U_F":'bD',"P_F":'gp',"L_U":'k2',"D_S":'r*',"U_S":'co',"D_F":'mv',"L_D":'y1'}
def generate_plots(test,elem_order,bctype,start_refine,end_refine,solve):
    """Generate L2 convergance plots"""
    from fsinewton.convergence_tests.test_analytic import PROBLEMFOLDER
    basepath = "../results/%s/%sdegree%i/" %(PROBLEMFOLDER,bctype,elem_order)
    datapath = basepath + "%sdata/"%test
    plottitle = "%s %s %s "%(test,elem_order,bctype)
    xaxis = []

    allerrors = []
    times = []
    spaceerrors = []
    for i in range(start_refine,end_refine + 1):
        store = datapath + "refinement" + str(i)
        try:
            errorfile = open(store + "/errordata")
            allerrors.append(cPickle.load(errorfile))
            
            meshsizefile = open(store + "/meshsize")
            xaxis.append(cPickle.load(meshsizefile))

            spaceerrorfile = open(store + "/spaceerrordata")
            spaceerrors.append(cPickle.load(spaceerrorfile))
            
            timesfile = open(store + "/times")
            times.append(cPickle.load(timesfile))
        except:
            info_blue("Pickle load error, has all the error data been generated?")
            exit()

    #rearrange to get a sequence of errors per function
    L2errorsperfunc = {}
    for k in allerrors[0].keys():
        L2errorsperfunc[k] = [allerrors[i][k] for i in range(len(allerrors))]
    
    #save loglogdata
    save_loglogdata(L2errorsperfunc, datapath + test, xaxis)

    #Save convergencedata
    save_convergencedata(L2errorsperfunc,datapath + test,start_refine,end_refine)

    #Create convergenceplots
    create_convergenceplots(basepath + test,allerrors,L2errorsperfunc,plottitle,xaxis)

    #Create space L2 error plots
    for i in range(start_refine,end_refine + 1):
        store = datapath + "refinement" + str(i)
        plot_L2_errors(spaceerrors[i], times[i],
                       store + "/L2errors", plottitle + "ref %i"%i)

def plot_L2_errors(L2_errors,times,path,title):
    """Create a plot of L2 error vs time (fixed refinement level)"""
    plt.figure() 
    ax = plt.gca()
    ax.set_yscale('log')
    for f in L2_errors.keys():
        plt.plot(times,L2_errors[f],PLOTSTYLES[f],label = latex_convert(f), linestyle = '-')
    plt.ylabel("Solution L2 error")
    plt.xlabel("Time")
##    plt.title(title)
    plt.legend(loc=0)
    plt.savefig(path)

def create_convergenceplots(store,allerrors,L2errorsperfunc,plottitle,xaxis):
    """Plot the convergence of errors"""

    #Plot the various data points
    for refinement_level in allerrors:                                                                
        plt.figure()
        plt.xlabel("mesh size hmin")
        pdf = PdfPages(store)
        for f in L2errorsperfunc.keys(): 
            plt.loglog(xaxis, L2errorsperfunc[f], PLOTSTYLES[f],
                       label = latex_convert(f), linestyle = '-')
        ax = plt.gca()
        ax.set_xlim(ax.get_xlim()[::-1])
        ax.grid()
        plt.ylabel("L2 error")
        #plt.title(plottitle)
        plt.legend(loc=0)
        plt.savefig(pdf, format ='pdf')
        pdf.close()

def plot_lmerror(store,lmerrors,plottitle,xaxis,title):
    """Plot the convergence of Lagrange multiplier
       condition errors"""

    #Plot the various data points                                                                
    plt.figure()
    plt.xlabel("mesh size hmin")
    pdf = PdfPages(store)

    labels = {"L_U":"U_F - U_S","L_D":"D_F - D_S"}
    
    for f in lmerrors:
        plt.loglog(xaxis, lmerrors[f],PLOTSTYLES[f],
                   label = latex_convert(labels[f]), linestyle = '-')
    ax = plt.gca()
    ax.set_xlim(ax.get_xlim()[::-1])
    ax.grid()
    plt.ylabel("L2 error")
    plt.title(title)
    plt.legend(loc=0)
    plt.savefig(pdf, format ='pdf')
    pdf.close()

def save_lmerror(lmerrors,path,start_refine,stop_refine):
    """Create a file for the Lagrange multiplier error data"""
                             
    lmdatafile = open(path + "_lmdata.txt","w")
    lmdatafile.write("\n Refinement levels %s\n"%str(range(start_refine,stop_refine)))
    for f in lmerrors.keys():
        lmdatafile.write("Relative errors of Lagrange Multiplier condition \
                              %s \n %s \n"%(f,str(lmerrors[f])))
    lmdatafile.close()

def save_convergencedata(L2errorsperfunc,path,start_refine,stop_refine):
    """Create a file for the error data"""
                             
    convergencedata = open(path + "_data.txt","w")
    convergencedata.write("\n Refinement levels %s\n"%str(range(start_refine,stop_refine)))
    for k in L2errorsperfunc.keys():
        convergencedata.write("L2 errors of %s \n %s \n"%(k,str(L2errorsperfunc[k])))
    convergencedata.close()

def save_loglogdata(L2errorsperfunc,path,xaxis,functions):
    """Create a file for the slopes of the loglog plots relative to the previous points"""
                             
    loglogdata = open(path + "_loglogdata.txt","w")
    loglogdata.write("Slopes between succesive log log points \n \n")
    conv_rates = {f:[] for f in functions}
    for k in L2errorsperfunc.keys():
        for j in range(len(L2errorsperfunc[k]) - 1):
            rise = ln(L2errorsperfunc[k][j+1]) - ln(L2errorsperfunc[k][j])
            run = ln(xaxis[j+1]) - ln(xaxis[j])
            conv_rates[k].append(rise/run)                                        
        loglogdata.write("%s \n %s \n"%(k, str(conv_rates[k])))
    loglogdata.close()

def plot_timings(timingdata,xaxis,path):
    """
    Plot the runtime of assembly and solve in graph vs. num vertices
    """
    #Plot the various data points                                                               
    plt.figure()
    plt.xlabel("Number of mesh vertices")
    pdf = PdfPages(path + "timings")
    for op in timingdata.keys():
        plt.loglog(xaxis, timingdata[op],'D',
                   label = op, linestyle = '-')
    ax = plt.gca()
##        ax.set_xlim(ax.get_xlim()[::-1])
    ax.grid()
    plt.ylabel("Runtime (s)")
    #plt.title(plottitle)
    plt.legend(loc=0)
    plt.savefig(pdf, format ='pdf')
    pdf.close()

    #estimate slopes
    loglogdata = open(path + "timing_slopes.txt","w")
    loglogdata.write("Slopes between succesive log log timings \n \n")
    conv_rates = {op:[] for op in timingdata.keys()}
    for op in timingdata.keys():
        for j in range(len(timingdata[op]) - 1):
            rise = ln(timingdata[op][j+1]) - ln(timingdata[op][j])
            run = ln(xaxis[j+1]) - ln(xaxis[j])
            conv_rates[op].append(rise/run)                                        
        loglogdata.write("%s \n %s \n"%(op, str(conv_rates[op])))
    loglogdata.close()

def latex_convert(input):
    return "$" + input + "$"

if __name__ == "__main__":
    print "Type of test all,fluid,struc,mesh,fluidstruc,strucmesh,meshfluid, order, \
           bc, refinements start, refinements end"
    
    import sys
    if len(sys.argv) > 7:
        print ("Usage: python %s [all|fluid|struc|mesh|fluidstruc|strucmesh|meshfluid], \
                elem_order, bctype[normal,dirichlet,neumann],start_refine [int], end_refine [int], \
                optional(solve) [t/f]"% sys.argv[0])
        exit()
    test = sys.argv[1]
    elem_order = int(sys.argv[2])
    bctype = sys.argv[3]
    start_refine = int(sys.argv[4])
    end_refine = int(sys.argv[5])
    if len(sys.argv) == 7 and sys.argv[6] == "f":
        solve = False
    else:
        solve = True
    generate_plots(test,elem_order,bctype,start_refine,end_refine,solve)
