z                 = sort((26:-1:1))*1e-3;
d                 = sort([246, 240, 235, 230, 226, 222, 218, 215, 213, 210, 206, 204, 201, 199, 197, 195, 194, 193, 192, 191, 190, 191, 190, 189, 189, 189])*1e-6;
Wavelength        = 532e-9;
TrimIndex         = 7;
Seed              = [100e-6 -5e-2];

[w0, dw0, z0, dz0]  = extractWaist(z, d, Wavelength, Seed, TrimIndex);

%% 

z                 = z0;
f                 = 500e-3;
Mr                = abs(f/(z-f));
r                 = zr/(z-f);
M                 = Mr/sqrt(1+r^2);
zp_to_f           = M^2*(z-f)*10^6; % distance of waist to focal point in um
w0p               = M*w0*1e3;       % waist after objective in um

%%

z                 = z0+dz0;
f                 = 50e-3;
Mr                = abs(f/(z-f));
r                 = zr/(z-f);
M                 = Mr/sqrt(1+r^2);
zp_to_fmax        = M^2*(z-f)*10^6; % distance of waist ot focal point in um
w0pmax            = M*w0*1e3;       % waist after objective in um
  
z                 = z0-dz0;
f                 = 50e-3;
Mr                = abs(f/(z-f));
r                 = zr/(z-f);
M                 = Mr/sqrt(1+r^2);
zp_to_fmin        = M^2*(z-f)*10^6; % distance of waist ot focal point in um
w0pmin            = M*w0*1e3; % waist after objective in um

function [w0, dw0, z0, dz0]  = extractWaist(z, d, lb, seed, trim_idx)
    
    w                 = d/2;
    
    fitmodel_waist    = @(w0,z0,x)w0*sqrt(1+(x-z0).^2/(pi.*w0.^2/lb).^2);
    fitmodel_waist_far= @(w0,z0,x)(x-z0)*lb/pi./w0;
    
    lower_bounds      = [0, -inf]; % Setting lower bound for w0 as 0 to prevent negative waist sizes
    upper_bounds      = [inf, inf];
    z_trimmed         = z(end-trim_idx:end);
    w_trimmed         = w(end-trim_idx:end);
    fitwaistfar       = fit(z_trimmed', w_trimmed', fitmodel_waist_far, 'StartPoint', seed, 'Lower', lower_bounds, 'Upper', upper_bounds); % fit first as if waist far away
    
    seed              = [fitwaistfar.w0 fitwaistfar.z0];
    fitwaist          = fit(z', w', fitmodel_waist, 'StartPoint', seed, 'Lower', lower_bounds, 'Upper', upper_bounds);
    
    ci                = confint(fitwaist,1-exp(-1)); %63% uncertainty confidence interval
    ws                = ci(:,1)*1e3; % in mm
    w0                = mean(ws);
    dw0               = diff(ws)/2;
    z0s               = ci(:,2)*1e3; % in mm
    z0                = mean(z0s);
    dz0               = diff(z0s)/2;
    
    % Plot
    zplot             = linspace(-0.03, 0.03, 1001);
    waistplot         = feval(fitwaist,zplot);
    waistplotfar      = feval(fitwaistfar, zplot);
     
    figure(1); 
    clf; 
    set(gcf,'Position',[100 100 950 750])
    set(gca,'FontSize',16,'Box','On','Linewidth',2); 
    hold all
    plot(z*1e3, w*1e6, 'ok', 'Linewidth',2, 'DisplayName', 'Measured Data')
    plot(zplot*1e3, waistplot*1e6,'r-','Linewidth',2, 'DisplayName', 'Fit at near distances')
    plot(zplot*1e3, waistplotfar*1e6,'b--','Linewidth',2,  'DisplayName', 'Fit at far distances') 
    ylim([min(w)*1e6, max(w)*1e6])
    grid on
    xlabel('Distance (mm)');
    ylabel('Beam Radius (mum)')
    legend('Location', 'NorthWest')
end