m=50;

%set up Laplace equation
h = 1.0/(m+1);
I = speye(m,m);
E = sparse(2:m,1:m-1,1,m,m);
D = 2*I-E-E';
A = (kron(D,I)+kron(I,D))/h^2;
b = zeros(m*m,1); %zero right-hand-side


xtrue = A\b; %true solution

%for Jacobi method
diagonal = diag(diag(A));
R = A - diagonal;

%for Gauss-Seidel
L = tril(A);
U = A-L;

%for SOR
sorL = tril(A,-1);
%TODO clean up and make the notation consistent :-)

omega = 2 / (1+sin(pi*h)); %optimal choice
omega2 = 1.5; %some other random choice.

xgs = 100*rand(m*m,1); %initial guess with random values
xj = xgs;
xsor = xgs;
xsor2 = xgs;

stepcount = 25;
errnormj = zeros(stepcount,1);
errnormgs = zeros(stepcount,1);
errnormsor = zeros(stepcount,1);
errnormsor2 = zeros(stepcount,1);
for i = 1:stepcount
    
    %store the error for each step
    errnormj(i) = norm(xj-xtrue);
    errnormgs(i) = norm(xgs-xtrue);
    errnormsor(i) = norm(xsor-xtrue);
    errnormsor2(i) = norm(xsor2-xtrue);
    
    %perform solves
    xj = diagonal \ (b  - R*xj); % Jacobi
    
    xgs = L\(b - U*xgs); %GS
    
    %SOR
    xsor = (diagonal + omega * sorL) \ (omega*b - (omega*U + (omega-1)*diagonal) * xsor);
    xsor2 = (diagonal + omega2 * sorL) \ (omega2*b - (omega2*U + (omega2-1)*diagonal) * xsor2);
    
    %show plots
    figure(1);
    subplot(2,2,1);
    height = reshape(xj-xtrue,m,m);
    bar3(height);
    zlim([0, 100]);
    title('Jacobi');
    
    subplot(2,2,2);
    height = reshape(xgs-xtrue,m,m);
    bar3(height);
    zlim([0, 100]);
    title('Gauss-Seidel');
    
    subplot(2,2,3);
    height = reshape(xsor-xtrue,m,m);
    bar3(height);
    zlim([0, 100]);
    title('SOR-optimal');
    
    subplot(2,2,4);
    height = reshape(xsor2-xtrue,m,m);
    bar3(height);
    zlim([0, 100]);
    title('SOR-other');
    
    pause(0.1);
end

figure(2);
clf;
plot(errnormj);
hold on;
plot(errnormgs);
plot(errnormsor);
plot(errnormsor2);
legend('Jacobi','Gauss-Seidel','SOR-opt','SOR-other');
ylim([0,2600]);
title('Soln error vs. iteration count');
