close all; clear; clc;

%% Setup
% Atomic units
hbar = 1; % Atomic units
m = 1; % Electron mass
omega = 1; % Hartree per hbar (per radian)
L = 5; % Domain size
xx = linspace(-L,L,1e2); dx = xx(2)-xx(1);

V = @(x) m/2*omega.^2*x.^2;

%% Show the behaviour of the edge over the course of the iteration
cm = plasma(35);
figure(1); clf; hold on;
colormap(cm); cb = colorbar; cb.Label.String = 'Iteration';
plot(xx,V(xx),'k'); 
grid on; box on; axis square;
ylim([0 6]);
e = 1.6; de = 0.15; % Energy guess and energy step
tol = 1e-3; % Wavefunction tolerance
Ns = 5; % Number of states to calculate
niter = 0;
U = zeros(numel(xx),1);

while true
    dpsidx = @(x,psi) [0 1;2*m/hbar.^2.*(V(x)-e) 0]*psi;
    psi0 = [0 sqrt(2*m/hbar.^2.*(V(-L)-e))];
    [~,psi] = ode23(dpsidx,xx,psi0);
    psi = real(psi(:,1)); psi = psi./sqrt(sum(psi.^2).*dx);
    niter = niter + 1;
    if abs(psi(end))<tol || de<eps, break; end
    if psi(end) > 0, e = e+de;
    else, de = de/2; e = e-de; end
    plot(xx,psi+e,'color',cm(niter,:),'Linewidth',1);
    drawnow; pause(0.05);
    U(:,niter) = psi;
end

%% Plot
figure(2); clf;
colormap(viridis);
surf(1:(niter-1),xx,U-U(:,end));
zlim([-1 1]*1); caxis([-1 1]*1);
grid on; box on;
xlim([1,niter-1]); ylim([-1 1]*L);
xlabel('Iteration'); ylabel('x'); zlabel('\Delta')
