%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% test of idealised point spread function generation
% see Adrian & Westerweel (2003), page 96, eq 3.18, 3.19

addpath([pwd '/../../../misc']);

%% input
n_r      = 200;
n_z      = 200;
n_k      = 50;
f_number = 8;
lambda   = 0.532;
Lz       = 500;
Lr       = 25;

%% grid
R        = linspace(-Lr, Lr, n_r);
Z        = linspace(-Lz, Lz, n_z);

phi      = 2*pi*Z/lambda / (2*f_number)^2;
psi      = 2*pi*R/lambda / (2*f_number);

[phi_mat, psi_mat] = ndgrid(phi, psi);

%% Lommel functions
% equations 8.19
Un       = zeros(n_z, n_r, 3);
Vn       = zeros(n_z, n_r, 3);
phi_psi  = phi_mat./psi_mat;
psi_phi  = psi_mat./phi_mat;

for n = 0 : 2
   for k = 0 : n_k-1
      Un(:,:,n+1)    = Un(:,:,n+1) ...
                     + (-1)^k * besselj(n+2*k, psi_mat) .* (phi_mat./psi_mat).^(n+2*k);
      Vn(:,:,n+1)    = Vn(:,:,n+1) ...
                     + (-1)^k * besselj(n+2*k, psi_mat) .* (psi_mat./phi_mat).^(n+2*k);
   end
end

%% point spread function
% equation 3.18
I1       = (Un(:,:,2).^2 + Un(:,:,3).^2) .* (2./phi_mat).^2;
I2       = (1 + Vn(:,:,1).^2 + Vn(:,:,2).^2 ...
         - 2*Vn(:,:,1).*cos(0.5*(phi_mat + psi_mat.^2./phi_mat)) ...
         - 2*Vn(:,:,2).*sin(0.5*(phi_mat + psi_mat.^2./phi_mat))) .* (2./phi_mat).^2;

%% plot PSF
figure(1);
clf;
subplot(2, 1, 1);
imagesc(Z, R, I1');
axis equal tight xy;
xlabel('$Z-Z_0$ [$\mu$m]', 'interpreter','latex','fontsize',24);
ylabel('$R$ [$\mu$m]', 'interpreter','latex','fontsize',24);
colorbar;
set(gca,'clim',[0 0.2]);
colormap(1- gray);

subplot(2, 1, 2);
imagesc(Z, R, I2');
axis equal tight xy;
xlabel('$Z-Z_0$ [$\mu$m]', 'interpreter','latex','fontsize',24);
ylabel('$R$ [$\mu$m]', 'interpreter','latex','fontsize',24);
colorbar;
set(gca,'clim',[0 1]);

%% slices at different focal positions
figure(2);
clf;
cmap=brewermap(64,'Spectral')*0.85;
for n = 1 : 20 : n_z/2
   color    = interpcmap(cmap, Z(n), [Z(1) 0]);
   plot(R, I1(n,:), '-', 'color', color);
   hold on;
end
cb=colorbar;
colormap(cmap);
set(gca,'clim',[Z(1) 0]);
xlabel('$R$ [$\mu$m]', 'interpreter','latex','fontsize',24);
ylabel('$I/I_0$', 'interpreter','latex','fontsize',24);
ylabel(cb, '$Z-Z_0$ [$\mu$m]', 'interpreter','latex','fontsize',24);
set(gca,'ticklabelinterpreter','latex','fontsize',24);