% function [u,t] = two_d_heat(L,t_f,h,k,u0)
% 
% this function computes solutions of the heat equation in 2-d on a square grid
% with periodic boundary conditions.  It uses Euler timestepping.  It needs
% inputs L (square side length), t_f (final time), h (space step), k (time step)
% and u0 (initial data)
%
function [u,t] = two_d_heat(L,t_f,h,k,u0)

% number of subintervals
N = L/h;

% I am doing everything to be consistent with the notes.  (0,0) is the
% left-bottom index, (N,0) is the right bottom, (0,N) is the left top
% and (N,N) is the right top.
%
% If we're at (i,j) then this corresponds to the I = i+1 + j*(N+1) th
% entry of the vector X.
%
% define the matrix for the points away from the boundaries at 0 and N
for i=1:N-1
    for j=1:N-1
        % five point stencil.  Have -4 at (i,j)
        I = 1+i+j*(N+1);
        A(I,I) = -4;
        % Have +1 at (i,j+1) *above* (i,j)
        i1 = i;
        j1 = j+1;
        I1 = 1+i1+j1*(N+1);
        A(I,I1) = 1;
        % Have +1 at (i-1,j) *left* of (i,j)
        i1 = i-1;
        j1 = j;
        I1 = 1+i1+j1*(N+1);
        A(I,I1) = 1;
        % Have +1 at (i+1,j) *right* of (i,j)
        i1 = i+1;
        j1 = j;
        I1 = 1+i1+j1*(N+1);
        A(I,I1) = 1;
        % have +1 at (i,j-1) *below* (i,j)
        i1 = i;
        j1 = j-1;
        I1 = 1+i1+j1*(N+1);
        A(I,I1) = 1;
    end
end
% boundary points on the left side: i=0.  Need to fix the left point.
i=0;
for j=1:N-1
    % five point stencil.  Have -4 at (i,j)
    I = 1+i+j*(N+1);
    A(I,I) = -4;
    % Have +1 at (i,j+1) *above* (i,j)
    i1 = i;
    j1 = j+1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (-1,j) *left* of (0,j).  Periodic BC --> 0==N, -1==N-1
    i1 = N-1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (i+1,j) *right* of (i,j)
    i1 = i+1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % have +1 at (i,j-1) *below* (i,j)
    i1 = i;
    j1 = j-1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
end
% boundary points on the right side: i=N
i=N;
for j=1:N-1
    % five point stencil.  Have -4 at (i,j)
    I = 1+i+j*(N+1);
    A(I,I) = -4;
    % Have +1 at (i,j+1) *above* (i,j)
    i1 = i;
    j1 = j+1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (i-1,j) *left* of (i,j)
    i1 = i-1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (N+1,j) *right* of (N,j) periodic bc --> N==0, N+1==1
    i1 = 1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % have +1 at (i,j-1) *below* (i,j)
    i1 = i;
    j1 = j-1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
end
% boundary points on the bottom: j=0
j=0;
for i=1:N-1
    % five point stencil.  Have -4 at (i,j)
    I = 1+i+j*(N+1);
    A(I,I) = -4;
    % Have +1 at (i,j+1) *above* (i,j)
    i1 = i;
    j1 = j+1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (i-1,j) *left* of (i,j)
    i1 = i-1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (i+1,j) *right* of (i,j)
    i1 = i+1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % have +1 at (i,-1) *below* (i,0)  period BC --> 0==N, -1==N-1
    i1 = i;
    j1 = N-1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
end
% boundary points on the tip: j=N
j=N;
for i=1:N-1
    % five point stencil.  Have -4 at (i,j)
    I = 1+i+j*(N+1);
    A(I,I) = -4;
    % Have +1 at (i,N+1) *above* (i,N) N == 0, N+1 == 1
    i1 = i;
    j1 = 1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (i-1,j) *left* of (i,j)
    i1 = i-1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % Have +1 at (i+1,j) *right* of (i,j)
    i1 = i+1;
    j1 = j;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
    % have +1 at (i,j-1) *below* (i,j)
    i1 = i;
    j1 = j-1;
    I1 = 1+i1+j1*(N+1);
    A(I,I1) = 1;
end
% bottom-left corner: i=0 & j=0
i=0;
j=0;
% five point stencil.  Have -4 at (i,j)
I = 1+i+j*(N+1);
A(I,I) = -4;
% Have +1 at (i,j+1) *above* (i,j)
i1 = i;
j1 = j+1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (-1,0) *left* of (0,0)  0==N, -1 ==N-1
i1 = N-1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (i+1,j) *right* of (i,j)
i1 = i+1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% have +1 at (0,-1) *below* (0,0) -1==N-1
i1 = i;
j1 = N-1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% bottom-right corner: i=N,j=0
i=N;
j=0;
% five point stencil.  Have -4 at (i,j)
I = 1+i+j*(N+1);
A(I,I) = -4;
% Have +1 at (i,j+1) *above* (i,j)
i1 = i;
j1 = j+1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (i-1,j) *left* of (i,j)
i1 = i-1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (N+1,0) *right* of (N,0)  N+1==1
i1 = N+1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% have +1 at (N,-1) *below* (N,0) -1 == N-1
i1 = i;
j1 = N-1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% top-left corner: i=0,j=N
i=0;
j=N;
% five point stencil.  Have -4 at (i,j)
I = 1+i+j*(N+1);
A(I,I) = -4;
% Have +1 at (0,N+1) *above* (0,N)  N+1 == 1
i1 = i;
j1 = 1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (-1,N) *left* of (0,N)  -1 == N-1
i1 = N-1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (i+1,j) *right* of (i,j)
i1 = i+1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% have +1 at (i,j-1) *below* (i,j)
i1 = i;
j1 = j-1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% top-right corner: i=N,j=N
i=N;
j=N;
% five point stencil.  Have -4 at (i,j)
I = 1+i+j*(N+1);
A(I,I) = -4;
% Have +1 at (N,N+1) *above* (N,N) N+1 == 1
i1 = i;
j1 = 1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (i-1,j) *left* of (i,j)
i1 = i-1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% Have +1 at (N+1,N) *right* of (N,N)  N+1 == 1
i1 = 1;
j1 = j;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;
% have +1 at (i,j-1) *below* (i,j)
i1 = i;
j1 = j-1;
I1 = 1+i1+j1*(N+1);
A(I,I1) = 1;

M = t_f/k;
D = 1;
lambda = D*k/h^2
% fill in the vector X as in the notes: bottom row first, reading left
% to right, followed by second-to-last row, reading left to right and so on
X = reshape(flipud(u0)',(N+1)^2,1);
u(:,:,1) = u0;
t(1) = 0;
for j=1:M
    t(j+1)=j*k;
    X = X + lambda*A*X;
    u(:,:,j+1) = flipud(reshape(X,N+1,N+1)');
end

% to see what's going on with all this reshaping and flipping... uncomment
% the following and execute it.
% for i=1:4
%     for j=1:4 
%         A(i,j) = j+(i-1)*4; 
%     end
% end
% A
% X = reshape(flipud(A)',16,1)
% B = flipud(reshape(X,4,4)')