close all; clear; clc;

%% Constants
% Atomic units
hbar = 1; % Atomic units
m = 1; % Electron mass
omega = 1; % Hartree per hbar (per radian)
L = 15; % Domain size
xx = linspace(-L,L,1e3); dx = xx(2)-xx(1);
V = @(x) m/2*omega.^2*x.^2;

%% Calculate wavefunctions
% Calculate the wavefunctions for first 50 states
% of the quantum harmonic oscillator

e = 0; de = 0.1; % Energy guess and energy step
tol = 1e-6; % Wavefunction tolerance
Ns = 50; % Number of states to calculate
E = zeros(1,Ns); Psi = zeros(length(xx),Ns); % Preallocate
%load('U_and_E_harmonic_N50.mat'); % < - Saved E and Psi
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('%.0f done\n',j);
    
    %psi2 = log10(abs(psi)+1./xx);
    xx2 = xx(xx>0);
    psi2 = medfilt1(log10(abs(psi(xx>0))))+(10*(1-xx2/L))';
    f = @(x) interp1(xx2,psi2,x);
    %options = optimset('MaxFunEvals',1e9);
    x0 = fminsearch(f,xx(end));
    figure(1); clf; plot(xx,f(xx),x0,f(x0),'o'); grid on; title(sprintf('j = %.0f',j)); drawnow;
    psi(xx>x0)=0;
    psi = psi./sqrt(sum(psi.^2).*dx);
    
    Psi(:,j) = psi; E(j) = e;
    de = 0.1; e = e+de; % Reset step size
end
%save('U_and_E_harmonic_N50.mat','E','Psi');

%% Plot Franck-Condon overlap
figure(1); clf;
a = axes; hold on;
for e = unique(sort([0.1 1:1:2]))
FC = zeros(Ns,Ns);
U2 = interp1(xx+6,Psi,xx,'linear',0);
for s1i = 1:Ns
    for s2i = 1:(Ns+1-s1i)
        FC(s2i,s1i) = abs(sum(Psi(:,s1i).*U2(:,s2i+s1i-1)).*dx);
    end
end
En = exp(-(0:Ns-1)/e); En = En'./sum(En);
plot(1:Ns,sgolayfilt(FC*En,1,3),'.-');
end
a.YScale = 'log'; xlim([2 35]); grid on; box on;
xlabel('\DeltaG / hbar\omega');
ylabel('Franck-Condon overlap');

%% Plot wavefunctions
figure(2); clf; hold on;
plot(xx,V(xx),'k-');
cm = viridis(Ns);
for j = 1:Ns
    patch(xx,Psi(:,j)+E(j),cm(j,:),'edgecolor','none','facealpha',0.2);
    plot(xx,Psi(:,j)+E(j),'color',cm(j,:));
end
grid on; box on; axis square;
xlim([-1 1]*L); ylim([0 50]);
xlabel('Generalised motion coordinate (GMC)'); ylabel('Energy / hbar\omega')