% Usage: [u2,ts] = animwave(x, u, T, a, lambda, scheme)
%
% Animate solution of one-way wave equation u_t + a u_x = 0, with
% periodic boundary conditions, given initial u(x), using various
% finite-difference schemes.
%
% x = uniformly spaced x-axis values (row array, length N)
% u = initial condition u(x) (row array, length N)
% T = simulation time
% a = velocity (number, *or* row array of length N for inhomogeneous case)
% lambda = dt/dx
% scheme = time-stepping scheme (optional).
%
% results: u2 = L2 norm |u| vs. time corresponding to array ts.

function [u2,ts] = animwave(x, u, T, a, lambda, scheme)
    if (nargin < 6)
        scheme = 'Forward-Time Backward-Space';
    end

    h = plot(x, u);
    axis([min(x) max(x) -2*max(abs(u)) 2*max(abs(u))]);
    set(h,'EraseMode', 'xor');
    dx = (max(x) - min(x)) / (length(u) - 1);
    dt = dx * lambda;
    t = 0;
    ts = [t]; u2 = [sqrt(u * u' * dx)];
   
    if (strcmp(scheme, 'Crank-Nicolson'))
        if (length(a) == 1)
            a = a * ones(size(x));
        end
        % matrix for Crank-Nicolson:
        A = diag(ones(size(x))) + 0.25 * lambda * (diag(a(1:end-1),1) - diag(a(2:end),-1));
        A(1,length(x)) = -0.25 * lambda * a(1);
        A(length(x),1) = 0.25 * lambda * a(end);
    end 
    
    v = u + 0.5 * lambda * a .* ([u(2:end),u(1)] - [u(end),u(1:end-1)]);; % for leap-frog
    
    while (t < T)
            t = t + dt;
            if (strcmp(scheme, 'Lax-Friedrichs'))
                u = 0.5 * ((1 - a*lambda) .* [u(2:end),u(1)] + (1 + a*lambda) .* [u(end),u(1:end-1)]);
            elseif (strcmp(scheme, 'Lax-Wendroff'))
                u = (1 - (a*lambda).^2) .* u - 0.5 * a*lambda .* ((1 - a*lambda) .* [u(2:end),u(1)] - (1 + a*lambda) .* [u(end),u(1:end-1)]);
            elseif (strcmp(scheme, 'Leap-Frog'))
                w = v - lambda * a .* ([u(2:end),u(1)] - [u(end),u(1:end-1)]);
                v = u; u = w;
            elseif (strcmp(scheme, 'Crank-Nicolson'))
                u = (A \ (u - 0.25 * lambda * a .* ([u(2:end),u(1)] - [u(end),u(1:end-1)])).').';
            else
                u = u - a .* lambda .* (u - [u(end),u(1:end-1)]);
            end
            set(h, 'XData', x, 'YData', u);
            drawnow;
            ts = [ts, t]; u2 = [u2, sqrt(u * u' * dx)];
    end
    uf = u;