clear all;
close all;
%% initial parameters
% This is a green path
lambda=532e-9;
nFS=1.4607; %fused silica refractive index at 532nm.

% distances and positions relative to the launcher. See alog 59094.
%15+138+42: collimator to HM1
%119: HM1 to HL1
%147: HM1 to HL2
%584: HM1 to HM2
% HM2 to HL3.
%16: HM3 to HL3.
zHM1= (15+138+42)*1e-3;
zHM2=zHM1+584e-3;
zHM3=zHM2+55e-3;
zAM3=zHM3+131e-3;

zHL1=zHM1+119e-3;
zHL2=zHM1+147e-3;
zHL3=zHM2+39e-3;

% ROC and focal length of lenses
RHL1=150e-3;
RHL2=150e-3;
RHL3=100e-3;
fHL1=RHL1/(nFS-1);
fHL2=RHL2/(nFS-1);
fHL3=RHL3/(nFS-1);

% ideal waist location 
zTargetWaist=432e-3+zAM3;
wTargetWaist=104.25e-6;

%% Measurement and fit, Jun 11
SHRailMark= (61- [61, 51, 41, 31, 25, 11, 8, 36, 46, 43])*1e-2
%[110,210,260 ,310, 360 ,410,460 ,510 ,610 ]*1e-3; %scan head rail position
%SH_BS=49e-3; % scan head to the DCPD-scanhead BS when the head was at 610mm, BS was moved a bit today.
SH_BS=110e-3; % scan head to the DCPD-scanhead BS when the head was at 610mm, BS was moved a bit today.
BS_flipper=89e-3; % BS to flipper. This changed too.
flipper_AM3= 100e-3; %flipper to A:M3 didn't change.
zNS=zAM3+flipper_AM3+BS_flipper+SH_BS+SHRailMark;
wP= [493, 339, 276,  464, 620, 929, 1024, 357, 274, 268]*1e-6/2; %using 13.5
wY=[476, 320,  279, 488, 638, 978,  1057, 370, 269, 266]*1e-6/2; %using 13.5
wD4SP=wP*1.013; %I only recorded the 1/e2!!! Just doing a dodgy numerical conversion..
wD4SY=wY*1.013;


beamSize=@(w0, z0, z) sqrt(w0^2+((z-z0)*lambda/pi/w0).^2);
fitfun=@(beta, z) beamSize(beta(1), beta(2), z);

betaP=nlinfit(zNS, wD4SP, fitfun, [2e-4, 1.5]);
betaY=nlinfit(zNS, wD4SY, fitfun, [2e-4, 1.5]);
[w0P, z0P, M2P]=fitMsquared(zNS, wD4SP, lambda);
[w0Y, z0Y, M2Y]=fitMsquared(zNS, wD4SY, lambda);

%% Quick and dirty add 1064 path!
SHRailMark_1064 = (61-[61, 51, 41, 31])*1e-2;
wY_1064 = [613, 834, 1124,  1402]*1e-6/2;
wP_1064 = [614, 845,  1150, 1430]*1e-6/2;
lambda2 = 1064e-9;
beamSize_1064=@(w0, z0, z) sqrt(w0^2+((z-z0)*lambda2/pi/w0).^2);
fitfun_1064=@(beta, z) beamSize_1064(beta(1), beta(2), z);

zNS_1064=zAM3+flipper_AM3+BS_flipper+SH_BS+SHRailMark_1064;

betaP_1064=nlinfit(zNS_1064, wP_1064, fitfun_1064, [2e-4, 1.5]);
betaY_1064=nlinfit(zNS_1064, wY_1064, fitfun_1064, [2e-4, 1.5]);
[w0P_1064, z0P_1064, M2P_1064]=fitMsquared(zNS_1064, wP_1064, lambda2);
[w0Y_1064, z0Y_1064, M2Y_1064]=fitMsquared(zNS_1064, wY_1064, lambda2);

%%

figure
zz=0:0.02:2;
plot(zz, beamSize(betaP(1), betaP(2), zz), 'r--', zz, beamSize(betaY(1), betaY(2), zz), 'b--')
hold on
plot(zz, w0P*sqrt(1+M2P^2*(lambda/pi/w0P^2)^2*(zz-z0P).^2), 'm', zz, w0Y*sqrt(1+M2Y^2*(lambda/pi/w0Y^2)^2*(zz-z0Y).^2), 'g');
plot(zNS, wD4SP, 'rx', zNS, wD4SY, 'bo')
legend('P Gauss fit', 'Y Gauss fit', ['P fit M^2\=', num2str(M2P)], ['Y fit M^2\=', num2str(M2Y)], 'P D4S', 'Y D4S');
grid on
xlim([1, 1.8])
ylim('auto')

%% plot the 1064 path
figure
zz=0:0.02:2;
plot(zz, beamSize_1064(betaP_1064(1), betaP_1064(2), zz), 'r--', zz, beamSize_1064(betaY_1064(1), betaY_1064(2), zz), 'b--')
hold on
plot(zz, w0P_1064*sqrt(1+M2P_1064^2*(lambda2/pi/w0P_1064^2)^2*(zz-z0P_1064).^2), 'm', zz, w0Y_1064*sqrt(1+M2Y_1064^2*(lambda2/pi/w0Y_1064^2)^2*(zz-z0Y_1064).^2), 'g');
plot(zNS_1064, wP_1064, 'rx', zNS_1064, wY_1064, 'bo')
legend('P Gauss fit', 'Y Gauss fit', ['P fit M^2\=', num2str(M2P)], ['Y fit M^2\=', num2str(M2Y)], 'P 1/e2', 'Y 1/e2');
grid on
xlim([1, 1.8])
ylim('auto')
title('Beam size of A path 1064 beam')

%% Seed the beamPath using the fit
PPath=beamPath();
PPath.addComponent(component.lens(fHL1, zHL1, 'H:L1')); 
PPath.addComponent(component.lens(fHL2, zHL2, 'H:L2')); 
PPath.addComponent(component.lens(fHL3, zHL3, 'H:L3')); 
PPath.addComponent(component.flatMirror(zHM3, 'H:M3'));
PPath.addComponent(component.flatMirror(zAM3, 'A:M3'));
% Using the measured parameter as seed is problematic because it's downstream of lenses
% and moving lenses won't do anything downstream of lenses.
% We'd temporarily use measured parameter as seed, then propagate that back
% upstream, and use that as the new seed.
PPath.seedWaist(betaP(1), betaP(2), lambda);
newPQ=PPath.qPropagate(0); % zero is the collimator position
PPath.seedz=0;
PPath.seedq=newPQ;
% Repeat this for YAW.
YPath=PPath.duplicate();
YPath.seedWaist(betaY(1), betaY(2), lambda);
newYQ=YPath.qPropagate(0); % zero is the collimator position
YPath.seedz=0;
YPath.seedq=newYQ;

%% Set the target waist
PPath.targetWaist(wTargetWaist, zTargetWaist, lambda);
YPath.targetWaist(wTargetWaist, zTargetWaist, lambda);

figure
PPath.plotSummary(zz);
subplot(211);
hold on
plot(zNS, wD4SP, 'x');
plot(PPath.targetz, 0, '*');
text(PPath.targetz+2e-2, 0, ['target waist ', num2str(PPath.targetq.waistSize*1e6), 'um']);
text(PPath.targetz+2e-2, -1.2e-4, [' at ', num2str(PPath.targetz), 'm']);
text(PPath.targetz+2e-2, -2.4e-4, ['Target Overlap: ', num2str(PPath.targetOverlap)])
title('Measured and fit for P')

figure
YPath.plotSummary(zz);
hold on
plot(zNS, wD4SY, 'o');
plot(YPath.targetz, 0, '*');
text(YPath.targetz+2e-2, 0, ['target waist ', num2str(YPath.targetq.waistSize*1e6), 'um']);
text(YPath.targetz+2e-2, -1.2e-4, [' at ', num2str(YPath.targetz), 'm']);
text(YPath.targetz+2e-2, -2.4e-4, ['Target Overlap: ', num2str(YPath.targetOverlap)])
title('Measured and fit for Y')

%% Optimize by moving L1 and L2 relative to current position by +-20cm.
newPPath=PPath.optimizePath('H:L1', [-0.2, 0.2], 'H:L2', [-0.2, 0.2], '-r');
newYPath=YPath.optimizePath('H:L1', [-0.2, 0.2], 'H:L2', [-0.2, 0.2], '-r');

figure
newPPath.plotSummary(zz);
title('P path if we move L1 and L2');
subplot 211
hold on
plot(newPPath.targetz, 0, 'x');
text(newPPath.targetz+2e-2, 0, ['target waist ', num2str(newPPath.targetq.waistSize*1e6), 'um']);
text(newPPath.targetz+2e-2, -1.2e-4, [' at ', num2str(newPPath.targetz), 'm']);
text(newPPath.targetz+2e-2, -2.4e-4, ['Target overlap: ', num2str(newPPath.targetOverlap)]);
figure
newYPath.plotSummary(zz);
title('Y path if we move L1 and L2');
subplot 211
hold on
plot(newYPath.targetz, 0, 'x');
text(newYPath.targetz+2e-2, 0, ['target waist ', num2str(newYPath.targetq.waistSize*1e6), 'um']);
text(newYPath.targetz+2e-2, -1.2e-4, [' at ', num2str(newYPath.targetz), 'm']);
text(newYPath.targetz+2e-2, -2.4e-4, ['Target overlap: ', num2str(newYPath.targetOverlap)]);

disp(['P overlap before and after optimizing: ', num2str(PPath.targetOverlap), ', ', num2str(newPPath.targetOverlap)]);
disp(['P: new L1 pos = old L1 pos + (', num2str((newPPath.component('H:L1').z-PPath.component('H:L1').z)*1000), ') mm']);
disp(['P: new L2 pos = old L2 pos + (', num2str((newPPath.component('H:L2').z-PPath.component('H:L2').z)*1000), ') mm']);
disp(['P: new L1-L2 distance = ', num2str((newPPath.component('H:L2').z-newPPath.component('H:L1').z)*1000), 'mm'])
disp(['Y overlap before and after optimizing: ', num2str(YPath.targetOverlap), ', ', num2str(newYPath.targetOverlap)]);
disp(['Y: new L1 pos = old L1 pos + (', num2str((newYPath.component('H:L1').z-YPath.component('H:L1').z)*1000), ') mm']);
disp(['Y: new L2 pos = old L2 pos + (', num2str((newYPath.component('H:L2').z-YPath.component('H:L2').z)*1000), ') mm']);
disp(['Y: new L1-L2 distance = ', num2str((newYPath.component('H:L2').z-newYPath.component('H:L1').z)*1000), 'mm'])


function [w0_, z0_, msquared_]=fitMsquared(z_, w_, lambda_)
    % w_^2= w0^2+M^4*theta^2*(z-z0)^2,
    % theta=Gaussian divergence angle = lambda/pi/w0
    % Therefore, if w_^2=A*z^2+B*z+C,
    %   A=M^4*theta^2, B=-2*z0*A, C=A*z0^2+w0^2
    % i.e.
    % z0=-B/A/2
    % w0^2=C-A*z0^2
    % theta = lambda/pi/w0
    % M^4=A/theta^2
    P_=polyfit(z_, w_.^2, 2);
    z0_=-P_(2)/P_(1)/2;
    w0_=sqrt(P_(3)-P_(1)*z0_.^2);
    theta_=lambda_/pi/w0_;
    msquared_=sqrt(P_(1)/theta_^2);
    return;
end

