clear vars;
colormap(parula(5));

a = 0.04;
h_x = 0.6;
h_t = 0.1;
w = 0.85;

len = 50;
p = a / (h_t^2);
q = 1 / (h_x^2);

y_g = 10;
x_g = 10;

g = @(x) sin(w*x);
idx = @(y, x) (y-1)*len + x;
[X,Y] = meshgrid(0: h_x: (len-1)*h_x);

f_prev = zeros(len*len, 1);
f_curr = zeros(len*len, 1);

% Calculate A matrix
A = zeros(len*len, len*len);

border_indexes = [];

for i = 1:len
    for j = 1:len
        if all(i > 1 & i < len & j > 1 & j < len)
            c = idx(i,   j);
            n = idx(i+1, j);
            e = idx(i,   j+1);
            s = idx(i-1, j);
            w = idx(i,   j-1);

            A(c, c) = p + 4*q;
            A(c, n) = -q;
            A(c, e) = -q;
            A(c, s) = -q;
            A(c, w) = -q;
        else
            c = idx(i,   j);
            border_indexes = [border_indexes c];
            A(c, c) = 1;
        end
    end
end

%   Point G
idx_g = idx(y_g, x_g);
A(idx_g, :) = zeros(len*len, 1);
A(idx_g, idx_g) = 1;

% Main loop
for i = 1:360
    b = p * (2*f_curr - f_prev);
    b(idx_g) =  g((i+1)*h_t);
    for j = 1:length(border_indexes)
        b(border_indexes(j)) = 0;
    end
    f_next = gmres(A, b, 1, 0.01);

    mesh(X, Y, reshape(f_next, [len, len]));
    caxis([-1 1]);
    axis([0 (len-1)*h_x 0 (len-1)*h_x -1 1]);
    view(90,0);

    pause(0.05);

    f_prev = f_curr;
    f_curr = f_next;

end