clc; clear all;

N = 30; % Size of the grid. Assume N > 3

L = 8; % Size of physical space (L*L)
h = L/N; % Size of steps

% Construct d/dx and d/dy matrix which will have size N^2 * N^2
a = [1; 0; -1; zeros(N - 3,1)];
A = zeros(N,N);
for i = 1:N
A(:, i) = circshift(a, i - 2);
end
A(1,N) = 0; A(N, 1) = 0;
E = eye(N);

% Do the tensor product

dx = kron(A, E)./(2*h);
dy = kron(E, A)./(2*h);

% Construct the potential matrix

v = @(x,y) 10*(sinh(x)^2 + sinh(y)^2)/(cosh(x) + cosh(y))^2;
V = zeros(N*N,N*N);
for i = 1:N
for j = 1:N
pos = j + N*(i - 1);
V(pos, pos) = v(-L/2 + h*(i-1), -L/2 + h*(j-1));
end
end

% Now construct the Hamiltonian matrix

H = -dx^2 - dy^2 + V;

[Psi, D] = eig(H);
eval = diag(D);

[useless, perm] = sort(eval);

eval = eval(perm); Psi = Psi(:,perm);

%% Plot a choice of Psi

fprintf('Preparing the plot... \n');
k = 0;
PsiFn = zeros(N,N);

for j = 1:N
for i = 1:N
PsiFn(i,j) = Psi(i + N*(j - 1), k + 1);
end
end
[y,z] = meshgrid(-L/2:h:L/2 - h);
% Ploting |Psi|^2
figure;
surf(y,z,PsiFn);
%surf(y,z,abs(PsiFn),'EdgeColor','none','LineStyle','none','FaceLighting','phong')
;