close all; clear; clc;

%% Load dataset
load('resonator_alignment_Ns_60x60.mat')

%% Fit cavity waist
plt = false;
gaussian = @(x,c,w,h) exp( -4.*log(2).*((x-c)./w).^2 ) .* h;
m = @(x,p) gaussian(x,0,p(1),p(2)); % Model the cavity waist as a Gaussian

options = optimset('MaxFunEvals',1e4,'MaxIter',1e4,'Display','none');

Ws = zeros(numel(g1s),numel(g2s));

if plt
    figure(1); clf;
    colormap(inferno);
    a = subplot(1,2,1);
    ih = imagesc(g1s,g2s,Ws);
    grid on; box on; axis square;
    a.YDir = 'normal';
    xlabel('g1'); ylabel('g2');
    title('Cavity beam waist')

    subplot(1,2,2);
    lh = plot(ps,zeros(size(ps)),ps,zeros(size(ps)));
    grid on; box on; axis square;
    xlabel('position'); ylabel('Norm. reflection count');
    ylim([0 1]);
end

tic
for g1i = 1:numel(g1s)
    for g2i = 1:numel(g2s)
        N = squeeze(Ns(g1i,g2i,:));
        N = N-1; N = N-min(N); N = N./max(N);
        f = @(p) mean((N-m(ps',p)).^2);
        pf = fminsearch(f,[0.1 1],options);
        pf(1) = max(min(abs(pf(1)),1),0); % Coerce to range (0,1)
        Ws(g1i,g2i) = pf(1);
        
        if plt
            ih.CData = Ws;
            lh(1).YData = N;
            lh(2).YData = m(ps,pf);
            drawnow
        end
    end
    if ~mod(g1i,round(numel(g1s)/20)), fprintf('%.f%% done\n',g1i/numel(g1s)*100); end
end
toc

%% Plot
figure(2); clf;
a = axes; hold on;
colormap(a,inferno);
imagesc(g1s,g2s,Ws');
a = gca; a.YDir = 'normal';
axis square;
plot(g1s(g1s<0),1./g1s(g1s<0),'w--','linewidth',1);
plot(g1s(g1s>0),1./g1s(g1s>0),'w--','linewidth',1);
plot([0 0],[-2 2],'w--','linewidth',1); plot([-2 2],[0 0],'w--','linewidth',1);
a = gca; a.YDir = 'normal';
grid on; box on; axis equal; xlim([-2 2]); ylim([-2 2]);
xlabel('g1'); ylabel('g2');