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

V = @(x) (x<0).*(x).^2 + ...
    (x>=0).*(4*(2*x).^2./(2+(2*x)).^2);
%V = @(x) m/2*omega.^2*x.^2;

%%
e = 0; de = 0.1; % Energy guess and energy step
tol = 1e-6; % Wavefunction tolerance
Ns = 10; % Number of states to calculate
options = optimset('MaxFunEvals',1e9);
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);
    
    % Zero-off the exponential runaway
    psi2 = log10( psi.^2 );
    f = @(x) interp1(xx(2:end),psi2(2:end),x);
    x0 = fminsearch(f,xx(end),options);
    psi(xx>x0)=0;
    psi = psi./sqrt(sum(psi.^2).*dx);
    
    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 = viridis(Ns);
plot(xx,ones(size(xx)).*3.8,'k--');
for j = 1:Ns
    patch(xx,0.5*U(:,j)+E(j),cm(j,:),'edgecolor','none','facealpha',0.2);
    plot(xx,0.5*U(:,j)+E(j),'color',cm(j,:),'linewidth',1);
end
grid on; box on; axis square;
xlim([-L 20]); ylim([0 6]);
xlabel('x'); ylabel('Potential (V) and wavefunctions');