% Usage: [u] = pset3prob3(x, u, T, b, mu)
%
% Solves heat equation given initial u and x, running for a time T, with
% parameters b and mu = dt/dx^2.
%
% Return value is u(x,T).

function [u] = pset3prob3(x, u, T, b, mu)
    % h = plot(x, u);
    % axis([min(x) max(x) min(u), max(u)])
    % set(h,'EraseMode', 'xor');

    dx = (max(x) - min(x)) / (length(u) - 1);
    dt = dx^2 * mu;
    t = 0;

    % force b to be a vector b(x).
    if (length(b) == 1)
      b = b * ones(size(x));
    end

    % matrix for Crank-Nicolson:
    A = -mu * diag(b) + 0.5 * mu * (diag(b(1:end-1),1) + diag(b(2:end),-1));
    A(1,length(x)) = 0.5 * mu * b(1);
    A(length(x),1) = 0.5 * mu * b(end);
    
    while (t < T)
            t = t + dt;
            u = ((eye(length(x)) - A) \ ((eye(length(x)) + A) * u.')).';

	    % set(h, 'XData', x, 'YData', u);
	    % drawnow;
	    % plot(x, u);
    end