function [x,y,t,u,A,B,b] = HeatEqCD
% Heat diffusion in a slab

% Warning: this code assumes a case dependet os.

% Setup, geometry and time interval
L = 1.0;
W = 1.0;
T = 100;

% Number of intervals in space and time
m = 20;
n = 20;
s = 160;

% Heat conduction coefficient
Kx = 0.001;
Ky = Kx;

% Theta scheme (the parameter theta) theta=1 is fully implicit, theta=0 is
% fully explicit, and theta=1/2 is the Crank Nicolson scheme
theta = 1;

% Allocate arrays
u = zeros((m+1)*(n+1),s+1);
A = zeros((m+1)*(n+1),(m+1)*(n+1));
B = zeros((m+1)*(n+1),(m+1)*(n+1));
I = eye((m+1)*(n+1),(m+1)*(n+1));
b = zeros((m+1)*(n+1),1);

% Create mesh
dx = L/m;
dy = W/n;
dt = T/s;
t = (0:s)*dt;
x = (0:m)*dx;
y = (0:n)*dy;

nux = Kx/dx^2;
nuy = Ky/dy^2;

% The explicit scheme is stable provided that (nux+nuy)*dt<=1/2

% Initial condition
for j=2:n
    for i=2:m
        k = (j-1)*(m+1) + i;
        u(k,1) = u0(x(i),y(j));
    end
end

% Matrix construction
% Interior points, discretization of dt*Laplacian
for j=2:n
    for i=2:m
        k = (j-1)*(m+1) + i;
        l = k;

        A(k,l-(m+1)) = nuy*dt;
        A(k,l-1) = nux*dt;
        A(k,l) = -(2*nux*dt + 2*nuy*dt);
        A(k,l+1) = nux*dt;
        A(k,l+(m+1)) = nuy*dt;
    end
end

% B is I + (1 - theta)*dt*(-Laplacian), and A is I - theta*dt*(-Laplacian)
B = I + (1 - theta)*A;
A = I - theta*A;

% Boundary conditions
j=1;
i=1;
k = (j-1)*(m+1) + i;
l = k;

A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
A(k,l+(m+1)) = -alpha(x(i),y(j))*theta*(2*nuy*dt);
A(k,l+1) = -alpha(x(i),y(j))*theta*(2*nux*dt);

for i=2:m
    k = (j-1)*(m+1) + i;
    l = k;


    A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
    A(k,l+(m+1)) = -alpha(x(i),y(j))*theta*(2*nuy*dt);
    A(k,l-1) = -alpha(x(i),y(j))*theta*nux*dt;
    A(k,l+1) = -alpha(x(i),y(j))*theta*nux*dt;
end

i=m+1;
k = (j-1)*(m+1) + i;
l = k;

A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
A(k,l+(m+1)) = -alpha(x(i),y(j))*theta*(2*nuy*dt);
A(k,l-1) = -alpha(x(i),y(j))*theta*(2*nux*dt);

for j=2:n
    i = 1;
    k = (j-1)*(m+1) + i;
    l = k;

    A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
    A(k,l+1) = -alpha(x(i),y(j))*theta*(2*nux*dt);
    A(k,l-(m+1)) = -alpha(x(i),y(j))*theta*nuy*dt;
    A(k,l+(m+1)) = -alpha(x(i),y(j))*theta*nuy*dt;

    i=m+1;
    k = (j-1)*(m+1) + i;
    l = k;
    A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
    A(k,l-1) = -alpha(x(i),y(j))*theta*(2*nux*dt);
    A(k,l-(m+1)) = -alpha(x(i),y(j))*theta*nuy*dt;
    A(k,l+(m+1)) = -alpha(x(i),y(j))*theta*nuy*dt;
end

j=n+1;
i=1;
k = (j-1)*(m+1) + i;
l = k;

A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
A(k,l-(m+1)) = -alpha(x(i),y(j))*theta*(2*nuy*dt);
A(k,l+1) = -alpha(x(i),y(j))*theta*(2*nux*dt);

for i=2:m
    k = (j-1)*(m+1) + i;
    l = k;

    A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
    A(k,l-(m+1)) = -alpha(x(i),y(j))*theta*(2*nuy*dt);
    A(k,l-1) = -alpha(x(i),y(j))*theta*nux*dt;
    A(k,l+1) = -alpha(x(i),y(j))*theta*nux*dt;
end

i=m+1;
k = (j-1)*(m+1) + i;
l = k;

A(k,l) = alpha(x(i),y(j))*(1 + theta*(2*nux*dt + 2*nuy*dt)) + beta(x(i),y(j));
A(k,l-(m+1)) = -alpha(x(i),y(j))*theta*(2*nuy*dt);
A(k,l-1) = -alpha(x(i),y(j))*theta*(2*nux*dt);

% Time stepping loop
for r=2:s+1

    b = B*u(:,r-1);

    % Right hand side
    for j=2:n
        for i=2:m
            k = (j-1)*(m+1) + i;

            b(k) = b(k) + dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1)));
        end
    end

    % Boundary conditions
    j=1;
    i=1;
    k = (j-1)*(m+1) + i;

    b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nuy*dt)*u(k+(m+1),r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nux*dt)*u(k+1,r-1) ...
        + alpha(x(i),y(j))*(dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
        + alpha(x(i),y(j))*2*dt/dx*sqrt(2)*(theta*Kgradu(x(i),y(j),t(r)) ...
        + (1 - theta)*Kgradu(x(i),y(j),t(r)))) ...
        + beta(x(i),y(j))*g(x(i),y(j),t(r));

    for i=2:m
        k = (j-1)*(m+1) + i;

        b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*(2*nuy*dt)*u(k+(m+1),r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nux*dt*u(k-1,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nux*dt*u(k+1,r-1) ...
            + alpha(x(i),y(j))*(dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
            + alpha(x(i),y(j))*2*dt/dx*(theta*Kgradu(x(i),y(j),t(r)) ...
            + (1 - theta)*Kgradu(x(i),y(j),t(r)))) ...
            + beta(x(i),y(j))*g(x(i),y(j),t(r));
    end

    i=m+1;
    k = (j-1)*(m+1) + i;

    b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nuy*dt)*u(k+(m+1),r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nux*dt)*u(k+1,r-1) ...
        + alpha(x(i),y(j))*(dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
        + alpha(x(i),y(j))*2*dt/dx*sqrt(2)*(theta*Kgradu(x(i),y(j),t(r)) ...
        + (1 - theta)*Kgradu(x(i),y(j),t(r)))) ...
        + beta(x(i),y(j))*g(x(i),y(j),t(r));

    for j=2:n
        i = 1;
        k = (j-1)*(m+1) + i;

        b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*(2*nux*dt)*u(k+1,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nuy*dt*u(k-(m+1),r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nuy*dt*u(k+(m+1),r-1) ...
            + dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
            + alpha(x(i),y(j))*2*dt/dx*(theta*Kgradu(x(i),y(j),t(r)) ...
            + (1 - theta)*Kgradu(x(i),y(j),t(r))) ...
            + beta(x(i),y(j))*g(x(i),y(j),t(r));

        i=m+1;
        k = (j-1)*(m+1) + i;
        l = k;

        b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*(2*nux*dt)*u(k-1,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nuy*dt*u(k-(m+1),r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nuy*dt*u(k+(m+1),r-1) ...
            + dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
            + alpha(x(i),y(j))*2*dt/dx*(theta*Kgradu(x(i),y(j),t(r)) ...
            + (1 - theta)*Kgradu(x(i),y(j),t(r))) ...
            + beta(x(i),y(j))*g(x(i),y(j),t(r));
    end

    j=n+1;
    i=1;
    k = (j-1)*(m+1) + i;

    b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nuy*dt)*u(k-(m+1),r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nux*dt)*u(k+1,r-1) ...
        + dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
        + alpha(x(i),y(j))*2*dt/dx*sqrt(2)*(theta*Kgradu(x(i),y(j),t(r)) ...
        + (1 - theta)*Kgradu(x(i),y(j),t(r))) ...
        + beta(x(i),y(j))*g(x(i),y(j),t(r));

    for i=2:m
        k = (j-1)*(m+1) + i;

        b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*(2*nuy*dt)*u(k-(m+1),r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nux*dt*u(k-1,r-1) ...
            + alpha(x(i),y(j))*(1 - theta)*nux*dt*u(k+1,r-1) ...
            + dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
            + alpha(x(i),y(j))*2*dt/dx*(theta*Kgradu(x(i),y(j),t(r)) ...
            + (1 - theta)*Kgradu(x(i),y(j),t(r))) ...
            + beta(x(i),y(j))*g(x(i),y(j),t(r));
    end

    i=m+1;
    k = (j-1)*(m+1) + i;

    b(k) = alpha(x(i),y(j))*(1 - (1 - theta)*(2*nux*dt + 2*nuy*dt))*u(k,r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nuy*dt)*u(k-(m+1),r-1) ...
        + alpha(x(i),y(j))*(1 - theta)*(2*nux*dt)*u(k-1,r-1) ...
        + dt*(theta*f(x(i),y(j),t(r)) + (1 - theta)*f(x(i),y(j),t(r-1))) ...
        + alpha(x(i),y(j))*2*dt/dx*sqrt(2)*(theta*Kgradu(x(i),y(j),t(r)) ...
        + (1 - theta)*Kgradu(x(i),y(j),t(r))) ...
        + beta(x(i),y(j))*g(x(i),y(j),t(r));

    % Below is where we solve a system of algebraic at each time step
    u(:,r) = A\b;

end

% Visualization

maxu = max(max(u));
minu = min(min(u));

for r=1:s+1
    surf(x,y,reshape(u(:,r),m+1,n+1));
    axis([0 1 0 1 minu maxu]);
    %     view([0 90]);
    shading interp;
    Frames(r) = getframe;
end

%Extra functions: right hand side, initial condition, and boundary conditions
function f = f(x,y,t)
f = 0;

function u0 = u0(x,y)
u0 = x*(x-1)*y*(y-1);

function alpha = alpha(x,y)
alpha = 1;

function beta = beta(x,y)
beta = 0;

function g = g(x,y,t)
g = 0;

function Kgradu = Kgradu(x,y,t)
Kgradu = 0;