__author__ = "Marie E. Rognes (meg@simula.no)"
__copyright__ = "Copyright (C) 2012 Marie Rognes"
__license__  = "GNU LGPL version 3 or any later version"

"""
This module contains functions for solving the system of equations:

  F'(t) = 1 - S(t),
  S'(t) = F(t) - 1

with the initial conditions

  F(0) = F^0
  S(0) = S^0

using explicit Euler, implicit Euler and the Crank-Nicolson method,
and in addition some tools for plotting solutions using
pylab/matplotlib.
"""

# Import plotting package:
import pylab

# Import numerical linear algebra tools from numpy
import numpy
import numpy.linalg

def plot_solutions_versus_time(times, solutions, method, dt):
    """
    Plot solutions versus times
    """

    assert(len(times) == len(solutions)), \
        "Mismatching length of times and solutions"
    pylab.figure()
    pylab.plot(times, [x[0] for x in solutions], '.-', label="F")
    pylab.plot(times, [x[1] for x in solutions], 'x-', label="S")
    pylab.grid(True)
    pylab.legend()
    pylab.title("Plot of (F, S) versus time t with %s" % method)
    pylab.xlabel("t")
    tag = "_".join(method.split(" "))
    pylab.savefig("solutions_versus_time_%s_dt%g.pdf" % (tag, dt))

def phase_plot(solutions, method, dt):
    """
    Plot phase plot of solutions
    """
    pylab.figure()
    pylab.plot([x[0] for x in solutions], [x[1] for x in solutions],
               '.-')
    pylab.grid(True)
    pylab.title("Phase plot of S versus F with %s" % method)
    pylab.xlabel("F")
    pylab.ylabel("S")
    tag = "_".join(method.split(" "))
    pylab.savefig("phase_plot_%s_dt%g.pdf" % (tag, dt))

def phase_plot3D(times, solutions, method, dt):
    """
    Plot phase plot of solutions and time using 3D plotting tools
    """
    try:
        import matplotlib.pyplot as plt
        from mpl_toolkits.mplot3d import Axes3D
    except:
        print "mpl_toolkits.mplot3d or matplotlib.pyplot not available."
        return

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot([x[0] for x in solutions], [x[1] for x in solutions], times)
    plt.title("Phase plot of F, S and time with %s" % method)
    plt.xlabel("F")
    plt.ylabel("S")
    tag = "_".join(method.split(" "))
    pylab.savefig("phase_plot_3D%s_dt%g.pdf" % (tag, dt))

def explicit_euler(dt, T, ic):
    """
    Solve the system of equations using the explicit Euler method
    """

    # Check that we get two initial conditions
    assert(len(ic) == 2)

    # Store solutions in list of tuples: one tuple for each time
    times = [0, ]
    solutions = [ic, ]

    # Extract initial conditions
    (F0, S0) = ic

    # Start time-loop: let F_, S_ be the previous solutions and F, S
    # be the current solutions
    t = dt
    (F_, S_) = ic
    while (t <= T):

        # Define the new solutions from the old solutions
        F = F_ + dt*(1 - S_)
        S = S_ + dt*(F_ - 1)

        # Store the new solutions
        solutions += [(F, S)]
        times += [t]

        # Prepare for next iteration by updating the previous values
        # to hold the just-computed values
        (F_, S_) = (F, S)
        t += dt

    return times, solutions

def implicit_euler(dt, T, ic):
    """
    Solution scheme:
     F_n + dt S_n  = F_{n-1} + dt
    - dt F_n + S_n = S_{n-1} - dt
    """

    # Check that we get two initial conditions
    assert(len(ic) == 2)

    # Store solutions in list of tuples: one tuple for each time
    times = [0, ]
    solutions = [ic, ]

    # Extract initial conditions
    (F0, S0) = ic

    # Start time-loop: let F_, S_ be the previous solutions and F, S
    # be the current solutions
    t = dt
    (F_, S_) = ic

    # Precompute matrix and inverse (is constant in time loop so
    # that's ok)
    A = numpy.matrix([[1, dt], [-dt, 1]])
    invA = numpy.linalg.inv(A)
    while (t <= T):

        # Compute current right-hand side
        b = numpy.matrix([F_ + dt, S_ - dt]).transpose()

        # Solve system by multiplying by precomputed inverse (ok for
        # small systems like this one)
        y = invA*b

        # Extract components (numpy returns 2 x 1 matrix)
        F = y[(0, 0)]
        S = y[(1, 0)]

        # Store the new solutions
        solutions += [(F, S)]
        times += [t]

        # Prepare for next iteration by updating the previous values
        # to hold the just-computed values
        (F_, S_) = (F, S)
        t += dt

    return times, solutions

def implicit_euler(dt, T, ic):
    """
    Solution scheme:
     F_n + dt S_n  = F_{n-1} + dt
    - dt F_n + S_n = S_{n-1} - dt
    """

    # Check that we get two initial conditions
    assert(len(ic) == 2)

    # Store solutions in list of tuples: one tuple for each time
    times = [0, ]
    solutions = [ic, ]

    # Extract initial conditions
    (F0, S0) = ic

    # Start time-loop: let F_, S_ be the previous solutions and F, S
    # be the current solutions
    t = dt
    (F_, S_) = ic

    # Precompute matrix and inverse (is constant in time loop so
    # that's ok)
    A = numpy.matrix([[1, dt], [-dt, 1]])
    invA = numpy.linalg.inv(A)
    while (t <= T):

        # Compute current right-hand side
        b = numpy.matrix([F_ + dt, S_ - dt]).transpose()

        # Solve system by multiplying by precomputed inverse (ok for
        # small systems like this one)
        y = invA*b

        # Extract components (numpy returns 2 x 1 matrix)
        F = y[(0, 0)]
        S = y[(1, 0)]

        # Store the new solutions
        solutions += [(F, S)]
        times += [t]

        # Prepare for next iteration by updating the previous values
        # to hold the just-computed values
        (F_, S_) = (F, S)
        t += dt

    return times, solutions

def crank_nicolson(dt, T, ic):

    # Check that we get two initial conditions
    assert(len(ic) == 2)

    # Store solutions in list of tuples: one tuple for each time
    times = [0, ]
    solutions = [ic, ]

    # Extract separate initial conditions components
    (F0, S0) = ic

    # Start time-loop: let F_, S_ be the previous solutions and F, S
    # be the current solutions
    t = dt
    (F_, S_) = ic

    # Precompute matrix and inverse (is constant in time loop so
    # that's ok)
    A = numpy.matrix([[1, 0.5*dt], [-0.5*dt, 1]])
    invA = numpy.linalg.inv(A)
    while (t <= T):

        # Compute current right-hand side (need to represent as 1 x 2
        # matrix for numpy's sake)
        b = numpy.matrix([F_ - 0.5*dt*S_ + dt, S_ + 0.5*dt*F_ - dt]).transpose()

        # Solve system by multiplying by precomputed inverse (ok for
        # small systems, like this one)
        y = invA*b

        # Extract components (numpy returns 2 x 1 matrix)
        F = y[(0, 0)]
        S = y[(1, 0)]

        # Store the new solutions
        solutions += [(F, S)]
        times += [t]

        # Prepare for next iteration by updating the previous values
        # to hold the just-computed values
        (F_, S_) = (F, S)
        t += dt

    return times, solutions

def compute_energy(F, S):
    return (F - 1)**2 + (S - 1)**2

def main(method, ic, T, dt, show_plots=True):

    # Solve
    if method == "explicit Euler":
        times, solutions = explicit_euler(dt, T, ic)
    elif method == "implicit Euler":
        times, solutions = implicit_euler(dt, T, ic)
    elif method == "Crank-Nicolson":
        times, solutions = crank_nicolson(dt, T, ic)
    else:
        print "Unknown solution method: %s" % method
        exit()

    # Compute r_N at final N
    r_N = compute_energy(solutions[-1][0], solutions[-1][1])

    # Finish here if we do not want to plot
    if not show_plots:
        return r_N

    # Plot solutions versus time
    plot_solutions_versus_time(times, solutions, method, dt)

    # Plot phase plot of solutions
    phase_plot(solutions, method, dt)

    # Plot 3D phase plot of solutions in time
    phase_plot3D(times, solutions, method, dt)

    return r_N


if __name__ == "__main__":

    ic = (0.9, 0.1)
    T = 10.0

    # Use this to generate plots
    ns = (1, 10)
    maybe = True

    # Use this to look at errors
    #ns = (1, 2, 4, 8, 16, 32, 64, 128)
    #maybe = False

    # Define time steps
    dts = [0.1/n for n in ns]

    # Uncomment at will
    #method = "explicit Euler"
    #method = "implicit Euler"
    method = "Crank-Nicolson"

    exact = compute_energy(ic[0], ic[1])
    values = []
    for dt in dts:
        r_N = main(method, ic, T, dt, show_plots=maybe)
        values += [r_N]

    # Generate LaTeX table using package booktabs for prettier tables:
    print "\\begin{center}"
    print "\\begin{tabular}{ccccc}"
    print "\\toprule"
    print "$\\Delta t$ & $r$ & $r_N$ & $|r - r_N|$ & $\\frac{|r - r_N|}{\Delta t}$ \\\\"
    print "\\midrule"
    for i in range(len(dts)):
        print "%10.4g &" % dts[i],
        print "%10.4g &" % exact,
        print "%10.4g &" % values[i],
        print "%10.4g &" % abs(exact - values[i]),
        print "%10.4g" % (abs(exact - values[i])/dts[i]),
        print "\\\\"
    print "\\bottomrule"
    print "\\end{tabular}"
    print "\\end{center}"

    # Uncomment this line if you actually want to look at the plots
    #pylab.show()