"""
1D wave equation with u=0 at the boundary.
Simplest possible implementation.

The key function is::

  u, x, t, cpu = (I, V, f, c, L, dt, C, T, user_action)

which solves the wave equation u_tt = c**2*u_xx on (0,L) with u=0
on x=0,L, for t in (0,T].  Initial conditions: u=I(x), u_t=V(x).

T is the stop time for the simulation.
dt is the desired time step.
C is the Courant number (=c*dt/dx), which specifies dx.
f(x,t) is a function for the source term (can be 0 or None).
I and V are functions of x.

user_action is a function of (u, x, t, n) where the calling
code can add visualization, error computations, etc.
"""

import numpy as np

def solver(I, V, f, c, L, dt, C, T, user_action=None):
    """Solve u_tt=c^2*u_xx + f on (0,L)x(0,T]."""
    Nt = int(round(T/dt))
    t = np.linspace(0, Nt*dt, Nt+1)   # Mesh points in time
    dx = dt*c/float(C)
    Nx = int(round(L/dx))
    x = np.linspace(0, L, Nx+1)       # Mesh points in space
    C2 = C**2                      # Help variable in the scheme
    if f is None or f == 0 :
        f = lambda x, t: 0
    if V is None or V == 0:
        V = lambda x: 0

    u   = np.zeros(Nx+1)   # Solution array at new time level
    u_1 = np.zeros(Nx+1)   # Solution at 1 time level back
    u_2 = np.zeros(Nx+1)   # Solution at 2 time levels back

    import time;  t0 = time.time()  # for measuring CPU time

    # Load initial condition into u_1
    for i in range(0,Nx+1):
        u_1[i] = I(x[i])

    if user_action is not None:
        user_action(u_1, x, t, 0)

    # Special formula for first time step
    n = 0
    for i in range(1, Nx):
        u[i] = u_1[i] + dt*V(x[i]) + \
               0.5*C2*(u_1[i-1] - 2*u_1[i] + u_1[i+1]) + \
               0.5*dt**2*f(x[i], t[n])
    u[0] = 0;  u[Nx] = 0

    if user_action is not None:
        user_action(u, x, t, 1)

    # Switch variables before next step
    u_2[:] = u_1;  u_1[:] = u

    for n in range(1, Nt):
        # Update all inner points at time t[n+1]
        for i in range(1, Nx):
            u[i] = - u_2[i] + 2*u_1[i] + \
                     C2*(u_1[i-1] - 2*u_1[i] + u_1[i+1]) + \
                     dt**2*f(x[i], t[n])

        # Insert boundary conditions
        u[0] = 0;  u[Nx] = 0
        if user_action is not None:
            if user_action(u, x, t, n+1):
                break

        # Switch variables before next step
        u_2[:] = u_1;  u_1[:] = u

    cpu_time = time.time() - t0
    return u, x, t, cpu_time


def viz(I, V, f, c, L, Nx, C, T,
        ymax,                      # y axis: [-ymax, ymax]
        u_exact,                   # u_exact(x, t)
        animate='u and u_exact',   # or 'error'
        movie_filename='movie',
        ):
    """Run solver and visualize u at each time level."""
    import matplotlib.pyplot as plt
    import time, glob, os

    class Plot:
        def __init__(self, ymax, frame_name='frame'):
            self.max_error = []   # hold max amplitude errors
            self.max_error_t = [] # time points corresponding to max_error
            self.frame_name = frame_name
            self.ymax = ymax

        def __call__(self, u, x, t, n):
            """user_action function for solver."""
            if animate == 'u and u_exact':
                plt.plot(x, u, 'r-')
                plt.plot(x, u_exact(x, t[n]), 'b--')
                plt.xlabel('x')
                plt.ylabel('u')
                plt.axis([0, L, -self.ymax, self.ymax])
                plt.title('t=%f' % t[n])
                plt.show()
            else:
                error = u_exact(x, t[n]) - u
                local_max_error = np.abs(error).max()
                # self.max_error holds the increasing amplitude error
                if self.max_error == [] or \
                       local_max_error > max(self.max_error):
                    self.max_error.append(local_max_error)
                    self.max_error_t.append(t[n])
                # Use user's ymax until the error exceeds that value.
                # This gives a growing max value of the yaxis (but
                # never shrinking)
                self.ymax = max(self.ymax, max(self.max_error))
                plt.plot(x, error, 'r-')
                plt.xlabel('x')
                plt.ylabel('error')
                plt.axis([0, L, -self.ymax, self.ymax])
                plt.title('t=%f' % t[n])
                plt.show()
            plt.savefig('%s_%04d.png' % (self.frame_name, n))

    # Clean up old movie frames
    for filename in glob.glob('frame_*.png'):
        os.remove(filename)

    plot = Plot(ymax)
    u, x, t, cpu = solver(I, V, f, c, L, Nx, C, T, plot)

    # # Make plot of max error versus time
    # plt.figure()
    # plt.plot(plot.max_error_t, plot.max_error)
    # plt.xlabel('time'); plt.ylabel('max abs(error)')
    # plt.savefig('error.png')
    # plt.savefig('error.pdf')

    # # Make .flv movie file
    # fps = 4  # Frames per second
    # codec2ext = dict(flv='flv') #, libx64='mp4', libvpx='webm', libtheora='ogg')

    # filespec = 'frame_%04d.png'
    # movie_program = 'avconv'  # or 'ffmpeg'
    # for codec in codec2ext:
    #     ext = codec2ext[codec]
    #     cmd = '%(movie_program)s -r %(fps)d -i %(filespec)s '\
    #           '-vcodec %(codec)s %(movie_filename)s.%(ext)s' % vars()
    #     os.system(cmd)

    for filename in glob.glob('frame_*.png'):
        os.remove(filename)

        
def simulations():
    from numpy import sin, cos, pi
    L = 12  # length of domain
    m = 8   # 2L/m is the wave length or period in space (2*pi/k, k=pi*m/L)
    c = 2   # wave velocity
    A = 1   # amplitude
    Nx = 80
    C = 0.8
    dt = C*(L/Nx)/c
    P = 2*pi/(pi*m*c/L)  # 1 period in time
    T = 6*P

    def u_exact(x, t):
        return A*sin(pi*m*x/L)*cos(pi*m*c*t/L)

    def I(x):
        return u_exact(x, 0)

    V = 0
    f = 0

    viz(I, V, f, c, L, dt, C, 2*P,
        0.1, u_exact,
        animate='error',
        movie_filename='error')

    # Very long simulation to demonstrate different curves
    viz(I, V, f, c, L, dt, C, T,
        1.2*A, u_exact,
        animate='u and u_exact',
        movie_filename='solution')

simulations()