#!/usr/bin/env python

"""This script generates plots from data found in a results directory"""

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

# First added:  2012-04-08
# Last changed: 2012-05-02

import sys
import pylab as pl

def read_ascii(dirname, filename):
    try:
        output = open(dirname + "/" + filename).read()
    except:
        return None
    return [[float(word) for word in line.split(" ") if len(word) > 0]
            for line in output.split("\n") if \
                (len(line) > 0 and not "level" in line)]

def extract_reference_value():
    return float(open("reference_value.txt").read())

def extract_functional(dirname):
    output = read_ascii(dirname, "goal_functional_final.txt")
    M_T = [line[1] for line in output]
    M_I = [line[2] for line in output]
    return M_T, M_I

def extract_errors(dirname):
    output = read_ascii(dirname, "error_estimates.txt")
    if output is None: return [None]*11
    E_0   = [line[1]  for line in output]
    E_0_F = [line[2]  for line in output]
    E_0_S = [line[3]  for line in output]
    E_0_M = [line[4]  for line in output]
    E     = [line[5]  for line in output]
    E_h   = [line[6]  for line in output]
    E_k   = [line[7]  for line in output]
    E_c   = [line[8]  for line in output]
    E_c_F = [line[9]  for line in output]
    E_c_S = [line[10] for line in output]
    E_c_M = [line[11] for line in output]
    return E_0, E_0_F, E_0_S, E_0_M, E, E_h, E_k, E_c, E_c_F, E_c_S, E_c_M

def extract_num_dofs(dirname):
    output = read_ascii(dirname, "num_dofs.txt")
    N = [line[2] for line in output]
    return N

def extract_timesteps(dirname):
    output = read_ascii(dirname, "timesteps.txt")
    if output is None: return None, None, None
    maxlevel = int(max(line[0] for line in output))
    t = []
    k = []
    R = []
    for i in range(maxlevel + 1):
        t.append([line[1] for line in output if int(line[0]) == i])
        k.append([line[2] for line in output if int(line[0]) == i])
        R.append([line[3] for line in output if int(line[0]) == i])
    return t, k, R

# Check command-line arguments
if len(sys.argv) != 2:
    print "Usage: plot <results directory>"
    exit(1)

# Extract directory
dirname = sys.argv[1]
print "Plotting results found in directory", dirname

# Extract data
M_0 = extract_reference_value()
M_T, M_I = extract_functional(dirname)
E_0, E_0_F, E_0_S, E_0_M, E, E_h, E_k, E_c, E_c_F, E_c_S, E_c_M = extract_errors(dirname)
N = extract_num_dofs(dirname)
t, k, Rk = extract_timesteps(dirname)

# Compute error
e = [abs(M_I[i] - M_0) for i in range(len(M_I))]

# Plot goal functionals
pl.figure()

pl.subplot(2, 1, 1); pl.plot(M_I, '-o')
pl.plot(range(len(M_I)), [M_0]*len(M_I), '--')
pl.grid(True)
pl.title("Goal functional")

pl.subplot(2, 1, 2); pl.semilogy(e, '-o')
pl.grid(True)
pl.title("Error")
pl.xlabel("Refinement level")

# Plot errors
if E_0 is not None:

    # Plot
    #eplot = pl.loglog
    eplot = pl.plot
    pl.figure()

    pl.subplot(4, 1, 1)
    eplot(N[:len(E)], E, '-o')
    eplot(N[:len(E_0)], E_0, '-^')
    eplot(N[:len(e)], e, '-s')
    pl.legend(["$E$", "$E_0$", "$e$"])
    pl.title("Errors and error estimates")
    pl.grid(True)

    pl.subplot(4, 1, 2)
    eplot(N[:len(E_h)], E_h, 'b-o')
    eplot(N[:len(E_k)], E_k, 'b-^')
    eplot(N[:len(E_c)], E_c, 'b-s')
    pl.legend(["$E_h$", "$E_k$", "$E_c$"])
    pl.grid(True)

    pl.subplot(4, 1, 3)
    eplot(N[:len(E_c)], E_c, 'b-o')
    eplot(N[:len(E_c_F)], E_c_F, 'b-^')
    eplot(N[:len(E_c_S)], E_c_S, 'b-s')
    eplot(N[:len(E_c_M)], E_c_M, 'b-*')
    pl.legend(["$E_0$", "$E_0^F$", "$E_0^S$", "$E_0^M$"])
    pl.grid(True)
    pl.xlabel("Number of degrees of freedom")

    pl.subplot(4, 1, 4)
    eplot(N[:len(E_0)], E_0, 'g-o')
    eplot(N[:len(E_0_F)], E_0_F, 'g-^')
    eplot(N[:len(E_0_S)], E_0_S, 'g-s')
    eplot(N[:len(E_0_M)], E_0_M, 'g-*')
    pl.legend(["$E_c$", "$E_c^F$", "$E_c^S$", "$E_c^M$"])
    pl.grid(True)
    pl.xlabel("Number of degrees of freedom")

# Plot efficiency indices
if E_0 is not None:

    # Compute error and efficiency indices
    I_0 = [E_0[i] / e[i] for i in range(min(len(e), len(E_0)))]
    I = [E[i] / e[i] for i in range(min(len(e), len(E)))]

    # Plot
    pl.figure()
    pl.plot(N[:len(I)], I, '-o')
    pl.plot(N[:len(I_0)], I_0, '-^')
    pl.legend(["$I$", "$I_0$"])
    pl.title("Efficiency indices")
    pl.grid(True)
    pl.xlabel("Number of degrees of freedom")

# Plot time steps and residuals
if t is not None:
    pl.figure()
    pl.subplot(2, 1, 1)
    for i in range(len(t)):
        pl.plot(t[i], k[i])
    pl.title("Time steps $k = k(t)$")
    pl.grid(True)
    pl.legend(["%d" % i for i in range(len(t))])
    pl.subplot(2, 1, 2)
    for i in range(len(t)):
        pl.plot(t[i], Rk[i])
    pl.title("Time residuals $R_k = R_k(t)$")
    pl.xlabel("$t$")
    pl.grid(True)
    pl.legend(["%d" % i for i in range(len(t))])

pl.show()
