% Parameters
Nx       = 256;
Ny       = 256;
Lx       = 10; 
Ly       = 10;

x        = linspace(-Lx/2, Lx/2, Nx);
dx       = x(2) - x(1);
Kmax     = pi*Nx/Lx;
kx       = linspace(-Kmax,Kmax,Nx+1);
kx       = kx(1:end-1); 
dkx      = kx(2)-kx(1);
kx       = fftshift(kx);

y        = linspace(-Ly/2, Ly/2, Ny);
dy       = y(2) - y(1);
Kmax     = pi*Ny/Ly;
ky       = linspace(-Kmax,Kmax,Ny+1);
ky       = ky(1:end-1); 
dky      = ky(2)-ky(1);
ky       = fftshift(ky);

[X, Y]   = ndgrid(x, y);
[KX, KY] = ndgrid(kx,ky);

% Operators
% Kinetic energy
KEop     = 0.5 * (KX.^2+KY.^2);
HKin     = @(w) ifft2(KEop .* fft2(w));

% Trap Potential
% Harmonic potential
V        = 0.5 * (X.^2 + Y.^2);
HV       = @(w) V.*w;

H        = @(w) HKin(w) + HV(w);

% Initial guess: random Gaussian
psi      = exp(-0.5 * (X.^2 + Y.^2)) .* (1 + 0.1*randn(size(X)));
psi      = psi / sqrt(sum(sum(abs(psi).^2)) * dx * dy);  % Normalize

% Gradient descent parameters (Heavy Ball)
alpha    = 1e-6;
beta     = 0.95;
epsilon  = 1e-8;
max_iter = 1000;

psi_old  = psi;
energy   = zeros(1, max_iter);

% Live plotting setup
figure(1);

for idx = 1:max_iter
    J           = H(psi); % Compute gradient: -Laplace + V

    % Chemical potential
    mu          = real(sum(conj(J(:)) .* psi(:))) / sum(abs(psi(:)).^2);

    % Residual and energy
    diff        = J - mu * psi;
    res         = sum(abs(diff(:)).^2) * dx * dy;
    energy(idx) = real(sum(conj(psi(:)) .* J(:)) * dx * dy);

    if res < epsilon
        fprintf('Converged at iteration %d\n', idx);
        break;
    end

    % Heavy-ball update
    psi_new     = (1 + beta)*psi - alpha * J - beta * psi_old;
    psi_old     = psi;

    % Normalize
    psi         = psi_new / sqrt(sum(sum(abs(psi_new).^2)) * dx * dy);

    % Plot the evolution of |psi|^2 every 50 iterations
    if mod(idx, 10) == 0
        imagesc(x, y, abs(psi).^2);
        title(sprintf('Iteration %d', idx));
        axis equal tight;
        colorbar;
        drawnow;
        pause(0.25)
    end
end

% Trim unused energy slots
energy          = energy(1:idx);

% Analytical ground state
psi_analytic    = (1/pi)^(1/2) * exp(-(X.^2 + Y.^2)/2);
psi_analytic    = psi_analytic / sqrt(sum(sum(abs(psi_analytic).^2)) * dx * dy);

% L2 error
L2_error        = sqrt(sum(sum(abs(psi - psi_analytic).^2)) * dx * dy);
fprintf('L2 error compared to analytical solution: %.5e\n', L2_error);

% Plot results
figure(2);
subplot(2,3,1); imagesc(x, y, abs(psi).^2); title('$|\psi|^2$ Numerical', 'Interpreter', 'latex', 'FontSize', 14); axis equal tight; colorbar;
subplot(2,3,2); imagesc(x, y, abs(psi_analytic).^2); title('$|\psi|^2$ Analytical', 'Interpreter', 'latex', 'FontSize', 14); axis equal tight; colorbar;
subplot(2,3,3); imagesc(x, y, abs(psi).^2 - abs(psi_analytic).^2); title('Difference $|\psi|^2$', 'Interpreter', 'latex', 'FontSize', 14); axis equal tight; colorbar;
subplot(2,3,[4 5 6]); plot(1:idx, energy, 'b-', 'LineWidth', 1.5); hold on;
yline(1, 'r--', 'Analytical E=1'); xlabel('Iteration'); ylabel('Energy');
title('Energy Convergence'); grid on; legend('Numerical', 'Analytical');