function [x,y,t,u,A,B,b] = HeatEqOS
% 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;

% The explicit scheme is stable provided that (nux+nuy)*dt<=1/2
nux = Kx/dx^2;
nuy = Ky/dy^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) = -3*nuy*dy/2*alpha(x(i),y(j))/(-sqrt(2)) ...
    - 3*nux*dx/2*alpha(x(i),y(j))/(-sqrt(2)) + beta(x(i),y(j));
A(k,l+(m+1)) = 4*nuy*dy/2*alpha(x(i),y(j))/(-sqrt(2));
A(k,l+2*(m+1)) = -nuy*dy/2*alpha(x(i),y(j))/(-sqrt(2));
A(k,l+1) = 4*nux*dx/2*alpha(x(i),y(j))/(-sqrt(2));
A(k,l+2) = -nux*dx/2*alpha(x(i),y(j))/(-sqrt(2));

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

    A(k,l) = -3*nuy*dy/2*alpha(x(i),y(j))*(-1) + beta(x(i),y(j));
    A(k,l+(m+1)) = 4*nuy*dy/2*alpha(x(i),y(j))*(-1);
    A(k,l+2*(m+1)) = -nuy*dy/2*alpha(x(i),y(j))*(-1);
end

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

A(k,l) = -3*nuy*dy/2*alpha(x(i),y(j))/(-sqrt(2)) ...
    + 3*nux*dx/2*alpha(x(i),y(j))/(sqrt(2)) + beta(x(i),y(j));
A(k,l+(m+1)) = 4*nuy*dy/2*alpha(x(i),y(j))/(-sqrt(2));
A(k,l+2*(m+1)) = -nuy*dy/2*alpha(x(i),y(j))/(-sqrt(2));
A(k,l-1) = -4*nux*dx/2*alpha(x(i),y(j))/(sqrt(2));
A(k,l-2) = nux*dx/2*alpha(x(i),y(j))/(sqrt(2));

for j=2:n
    i = 1;
    k = (j-1)*(m+1) + i;
    l = k;
    A(k,l) = - 3*nux*dx/2*alpha(x(i),y(j))*(-1) + beta(x(i),y(j));
    A(k,l+1) = 4*nux*dx/2*alpha(x(i),y(j))*(-1);
    A(k,l+2) = -nux*dx/2*alpha(x(i),y(j))*(-1);

    i=m+1;
    k = (j-1)*(m+1) + i;
    l = k;
    A(k,l) = 3*nux*dx/2*alpha(x(i),y(j)) + beta(x(i),y(j));
    A(k,l-1) = -4*nux*dx/2*alpha(x(i),y(j));
    A(k,l-2) = nux*dx/2*alpha(x(i),y(j));
end

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

A(k,l) = 3*nuy*dy/2*alpha(x(i),y(j))/(sqrt(2)) ...
    - 3*nux*dx/2*alpha(x(i),y(j))/(-sqrt(2)) + beta(x(i),y(j));
A(k,l-(m+1)) = -4*nuy*dy/2*alpha(x(i),y(j))/(sqrt(2));
A(k,l-2*(m+1)) = nuy*dy/2*alpha(x(i),y(j))/(sqrt(2));
A(k,l+1) = 4*nux*dx/2*alpha(x(i),y(j))/(-sqrt(2));
A(k,l+2) = -nux*dx/2*alpha(x(i),y(j))/(-sqrt(2));

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

    A(k,l) = 3*nuy*dy/2*alpha(x(i),y(j)) + beta(x(i),y(j));
    A(k,l-(m+1)) = -4*nuy*dy/2*alpha(x(i),y(j));
    A(k,l-2*(m+1)) = nuy*dy/2*alpha(x(i),y(j));
end

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

A(k,l) = 3*nuy*dy/2*alpha(x(i),y(j))/(sqrt(2)) ...
    + 3*nux*dx/2*alpha(x(i),y(j))/(sqrt(2)) + beta(x(i),y(j));
A(k,l-(m+1)) = -4*nuy*dy/2*alpha(x(i),y(j))/(sqrt(2));
A(k,l-2*(m+1)) = nuy*dy/2*alpha(x(i),y(j))/(sqrt(2));
A(k,l-1) = -4*nux*dx/2*alpha(x(i),y(j))/(sqrt(2));
A(k,l-2) = nux*dx/2*alpha(x(i),y(j))/(sqrt(2));

% 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))*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))*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))*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))*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;

        b(k) = alpha(x(i),y(j))*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))*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))*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))*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;