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
L2 = 10;
xx = linspace(-L,L,1e3); dx = xx(2)-xx(1);
xx2 = linspace(-L2,L2,1e3); dx2 = xx2(2)-xx2(1);

V = @(x) (x<-2).*(x+2).^2 + ...
    (x>=-2 & x<0).*(1.8-0.2*sin(x*2*pi/2*3.5)) + ...
    (x>=0).*1.5.*(4*x.^2./(2+x).^2);
V = @(x) m/2*omega.^2*x.^2;

%% Solve
e = 0; de = 0.1; % Energy guess and energy step
tol = 1e-6; % Wavefunction tolerance
Ns = 5; % Number of states to calculate
U = zeros(length(xx),Ns); % Preallocate
for j = 1:Ns
    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);
        if abs(psi(end))<tol || de<eps, break; end
        if psi(end)*(-1)^(j+1) > 0, e = e+de;
        else, de = de/2; e = e-de; end
    end
    fprintf('%d/%d done\n',j,Ns);
    U(:,j) = psi; E(j) = e;
    de = 0.1; e = e+de; % Reset step size
end

%% Plot
figure(1); clf; hold on;
plot(xx,V(xx),'k-');
cm = lines;
for j = 1:Ns
    patch([0 xx 0],[0 U(:,j)' 0]+E(j),cm(j,:),'edgecolor','none','facealpha',0.15);
    plot(xx,U(:,j)+E(j),'color',cm(j,:));
end
grid on; box on;
xlim([-1 1]*L); ylim([0 6]);
xlabel('x'); ylabel('Potential (V)');
axis square;
