clear; close all;  % not needed, except for debugging

%% Build model
%- Load model parameter set
datapth = "~/Data";

lbsync = {load("models/lbsync.mat")};
lbsync = [lbsync, load(fullfile(datapth, "2024-07-01 XFEL Characterization", "data/lbsync.mat"))];
lbsync = [lbsync, load(fullfile(datapth, "2024-10-10 XFEL Characterization", "data/lbsync.mat"))];
lbsync = [lbsync, load(fullfile(datapth, "2024-12-11 XFEL Characterization", "data/lbsync.mat"))];
lbsync = [lbsync, load(fullfile(datapth, "2024-12-12 XFEL Characterization", "data/lbsync.mat"))];
lbsync = [lbsync, load("models/lbsync_single_link_pz_res.mat")];
lbsync = merge_structs(lbsync{:});

%%
%- Build model
time_unit = "milliseconds";
scaling = 15;

mo = chgTimeUnit(10^(scaling - lbsync.xfel.mo.scaling) * lbsync.xfel.mo.Fd, time_unit);
mo.InputName = "w"; mo.Outputname = "y";
mo.OutputUnit = "femtoseconds";
mo.UserData = lbsync.xfel.mo.Fd.UserData;

mlo = build_laser_model(lbsync.xfel.laser.mlo1, Scaling=scaling, TimeUnit=time_unit);
slo = build_laser_model(lbsync.xfel.laser.slo1, Scaling=scaling, TimeUnit=time_unit);
lonk = build_link_model(lbsync.xfel.link.slo1, Scaling=scaling, TimeUnit=time_unit);
relonk = build_link_model(lbsync.xfel.link.relonk, Scaling=scaling, TimeUnit=time_unit);

%%
f = lbsync.xfel.mo.Fd.UserData.raw.data{1}.Frequency;
f = f(1 <= f & f <= 1e6);
w = 2*pi * f * 1e-3;
smoothed = false;
frds = struct();

frds.mo = load_frd(mo, smoothed, 1e6);
frds.mlo = load_frd(mlo.Fd, smoothed, 1e3);
frds.lonk = load_frd(lonk.Fd, smoothed, 500);
frds.slo = load_frd(slo.Fd, smoothed, 1e3);
frds.relonk = load_frd(relonk.Fd, smoothed, 500);

figure();
bodemag(frds.mo, frds.mlo, frds.lonk, frds.slo, frds.relonk, {2*pi, 2e6*pi});
legend("MO", "MLO", "LONK", "SLO", "ReLONK");

frds = structfun(@(frd) interp(frd, f), frds, UniformOutput=false);
frds = structfun(@(frd) chgTimeUnit(frd, time_unit), frds, UniformOutput=false);

%% MLO Lock
mlo.CLcol_y = chain2(mo, mlo.CLcol);
mlo.CLcol_e = getIOTransfer(mlo.CLcol_y, 'w', 'e');
mlo.CLfrd = mlo.CL * blkdiag(frds.mo, frds.mlo, mlo.Fn);
mlo.CLfrd.InputName = "w";
mlo.CLfrd_e = getIOTransfer(mlo.CLfrd, 'w', 'e');

[rcol, rfrd] = prepare_responses(mlo, f, "MLO");
ref = load(fullfile(datapth, "2024-12-12 XFEL Round-Trip Experiment/data/", "xfel_mlo1_resid.mat"));
ref = ref.noise_tr;

plot_lock(rcol, rfrd, ref, "MLO");
export_responses(rcol, rfrd, "mlo", 1e-3);

%% SLO Link Lock
lonk.CLcol_y = chain2(replaceBlock(mlo.CLcol_y, 'e', 1), lonk.CLcol);
lonk.CLcol_e = getIOTransfer(lonk.CLcol_y, 'w', 'e');
lonk.CLfrd = lonk.CL * blkdiag(replaceBlock(mlo.CLfrd, 'e', 1), frds.lonk, lonk.Fn);
lonk.CLfrd.InputName = "w";
lonk.CLfrd_e = getIOTransfer(lonk.CLfrd, 'w', 'e');

[rcol, rfrd] = prepare_responses(lonk, f, "LONK");
ref = load(fullfile(datapth, "2024-12-12 XFEL Round-Trip Experiment/data/", "xfel_lonk_resid.mat"));
ref = ref.noise_tr;

plot_lock(rcol, rfrd, ref, "LONK");
export_responses(rcol, rfrd, "lonk", 1e-3);

%% SLO Lock
slo.CLcol_y = chain2(replaceBlock(lonk.CLcol_y, 'e', 1), slo.CLcol);
slo.CLcol_e = getIOTransfer(slo.CLcol_y, 'w', 'e');
slo.CLfrd = slo.CL * blkdiag(replaceBlock(lonk.CLfrd, 'e', 1), frds.slo, slo.Fn);
slo.CLfrd.InputName = "w";
slo.CLfrd_e = getIOTransfer(slo.CLfrd, 'w', 'e');

[rcol, rfrd] = prepare_responses(slo, f, "SLO");
ref = load(fullfile(datapth, "2024-12-12 XFEL Round-Trip Experiment/data/", "xfel_slo1_resid.mat"));
ref = ref.noise_tr;

plot_lock(rcol, rfrd, ref, "SLO");
export_responses(rcol, rfrd, "slo", 1e-3);

%% Return Link Lock
relonk.CLcol_y = chain2(replaceBlock(slo.CLcol_y, 'e', 1), relonk.CLcol);
relonk.CLcol_e = getIOTransfer(relonk.CLcol_y, 'w', 'e');
relonk.CLfrd = relonk.CL * blkdiag(replaceBlock(slo.CLfrd, 'e', 1), frds.relonk, relonk.Fn);
relonk.CLfrd.InputName = "w";
relonk.CLfrd_e = getIOTransfer(relonk.CLfrd, 'w', 'e');

[rcol, rfrd] = prepare_responses(relonk, f, "ReLONK");
ref = load(fullfile(datapth, "2024-12-12 XFEL Round-Trip Experiment/data/", "xfel_relonk_resid.mat"));
ref = ref.noise_tr;

plot_lock(rcol, rfrd, ref, "ReLONK");
export_responses(rcol, rfrd, "relonk", 1e-3);

%% Round-Trip Measurement
rtrip_det = ss([], [], [], [1, db2mag(-75)], InputName={'r', 'w'}, OutputName="y", TimeUnit="milliseconds");
CLcol_y = chain2(mo, mlo.CLcol, lonk.CLcol, slo.CLcol, relonk.CLcol, rtrip_det, AnalysisPoint={'r(2)'});
CLcol_y = replaceBlock(CLcol_y, 'e', 1);
CLcol_y.InputName = "w";
CLcol_e = getIOTransfer(CLcol_y, 'w', {'r(2)', 'y'});
CLcol_e = [1, -1] * CLcol_e;
CLfrd = chain2(mlo.CL, lonk.CL, slo.CL, relonk.CL, rtrip_det, AnalysisPoint={'r(2)'});
CLfrd = CLfrd * blkdiag(frds.mo, frds.mlo, frds.lonk, frds.slo, frds.relonk, mlo.Fn, lonk.Fn, slo.Fn, relonk.Fn, 1);
CLfrd = replaceBlock(CLfrd, 'e', 1);
CLfrd.InputName = "w";
CLfrd_e = getIOTransfer(CLfrd, 'w', {'r(2)', 'y'});
CLfrd_e = [1, -1] * CLfrd_e;

rcol = abs(squeeze(freqresp(CLcol_e, w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "ReLONK", "ReLONKDetection", "RoundTripDetection", "TimingJitter"]);
rfrd = abs(squeeze(frd(CLfrd_e).ResponseData))';
rfrd = array2table([f, mag2db(rfrd), pow2db(sum(rfrd.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "LONK", "SLO", "ReLONK", "MLODetection", "LONKDetection", "SLODetection", "ReLONKDetection", "RoundTripDetection", "TimingJitter"]);
rfrd = movevars(rfrd, rcol.Properties.VariableNames);

ref = load(fullfile(datapth, "2024-12-12 XFEL Round-Trip Experiment/data/", "xfel_roundtrip_resid.mat"));
ref = ref.noise_tr;

% Compute NoFloorIntegratedJitter
rcol = compute_integrated_jitters(rcol);
rfrd = compute_integrated_jitters(rfrd);

plot_lock(rcol, rfrd, ref, "Round-Trip");
plot_lock2(rfrd, ref, "Round-Trip");
legend("MO", "MLO", "MLO MZM", "LONK", "LONK OXC", "SLO", "SLO OXC", "ReLONK", "ReLONK OXC", "Round-Trip Detection", "Total");


export_responses(rcol, rfrd, "roundtrip", 1e-3);

%%
try
    TileFigures([], [], [], 1);
end

%% Helper functions
function fr = load_frd(sys, smoothed, floor_start)
    trc = sys.UserData.raw.data{smoothed+1};
    fr = frd(db2mag(trc.TimingJitter), trc.Frequency, InputName="w", OutputName="d", FrequencyUnit="Hz");
    fresp = abs(squeeze(freqresp(sys, trc.Frequency, "Hz")));
    fr.ResponseData(:,:,trc.Frequency >= floor_start) = fresp(trc.Frequency >= floor_start);
end

function tr = compute_integrated_jitters(tr)
    tr.IntegratedTimingJitter = sqrt(cumtrapz(tr.Frequency, db2pow(tr.TimingJitter)));
    tr.ReverseIntegratedTimingJitter = sqrt(-flip(cumtrapz(flip(tr.Frequency), flip(db2pow(tr.TimingJitter)))));
    mask = tr.Frequency < 1e5;
    tr.ReverseIntegratedTimingJitterNoFloor = nan(height(tr), 1);
    tr.ReverseIntegratedTimingJitterNoFloor(mask) = sqrt(-flip(cumtrapz(flip(tr.Frequency(mask)), flip(db2pow(tr.TimingJitter(mask))))));
end

function [rcol, rfrd] = prepare_responses(sys, f, name)
    rcol = abs(squeeze(freqresp(sys.CLcol_e, f, "Hz")))';
    rcol = [rssq(rcol(:,1:end-2), 2), rcol(:,end-1:end)]; % Sum up all contributions to reference
    rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], VariableNames=["Frequency", "Reference", name, "Detection", "TimingJitter"]);
    rcol = compute_integrated_jitters(rcol);
    rfrd = abs(squeeze(frd(sys.CLfrd_e).ResponseData))';
    rfrd = [rssq(rfrd(:,1:end-2), 2), rfrd(:,end-1:end)];
    rfrd = array2table([f, mag2db(rfrd), pow2db(sum(rfrd.^2, 2))], VariableNames=["Frequency", "Reference", name, "Detection", "TimingJitter"]);
    rfrd = compute_integrated_jitters(rfrd);
end

function plot_lock(rcol, rfrd, ref, name)
    fig = figure(); tl = tiledlayout(fig, 2, 1);
    ax = nexttile(tl);
    ax.XScale="log"; hold(ax, "on");
    plot(ax, rfrd, "Frequency", contains(rfrd.Properties.VariableNames, "Integrated"));
    h = plot(ax, rcol, "Frequency", contains(rcol.Properties.VariableNames, "Integrated"));
    for i = 1:numel(h), h(i).SeriesIndex = i; h(i).LineStyle = "--"; end
    plot(ax, ref, "Frequency", contains(ref.Properties.VariableNames, "Integrated"), Color="black", DisplayName="Measured");

    ax = nexttile(tl);
    ax.XScale="log"; colororder(ax, "gem12"); hold(ax, "on");
    plot(ax, rcol, "Frequency", 2:(width(rcol)-3));
    h = plot(ax, rfrd, "Frequency", 2:(width(rfrd)-3));
    for i = 1:numel(h), h(i).SeriesIndex = i; end
    plot(ax, ref.Frequency, ref.TimingJitter, "k-", DisplayName="Measured");
    legend(ax, "Reference", name, "Detection", "Total");
    ylim(ax, [-120, inf]);
end

function plot_lock2(rfrd, ref, name)
    fig = figure(); tl = tiledlayout(fig, 3, 1);
    ax = nexttile(tl); ax.FontSize = 12;
    ax.XScale="log"; hold(ax, "on");
    ylabel(ax, "Integrated Jitter (fs)");
    plot(ax, rfrd, "Frequency", "ReverseIntegratedTimingJitterNoFloor");
    plot(ax, ref, "Frequency", "ReverseIntegratedTimingJitterNoFloor", Color="black", DisplayName="Measured");

    ax = nexttile(tl, [2,1]); ax.FontSize = 12;
    ax.XScale="log"; colororder(ax, "gem12"); hold(ax, "on");
    plot(ax, rfrd, "Frequency", 2:(width(rfrd)-3));
    plot(ax, ref.Frequency, ref.TimingJitter, "k-", DisplayName="Measured");
    legend(ax, "Reference", name, "Detection", "Total");
    ylim(ax, [-120, inf]);
    ylabel(ax, "Timing Jitter Spectral Density (dBfs^2/Hz)");
    xlabel(ax, "Fourier Frequency (Hz)");

    fig.Units = "centimeters";
    fig.InnerPosition(3:4) = [20,20];
end

function export_responses(rcol, rfrd, name, tol)
    arguments
        rcol table
        rfrd table
        name (1,1) string
        tol (1,1) double = 1e-3
    end

    rcol = reducepoly_table(rcol, tol);
    rfrd = reducepoly_table(rfrd, tol);
    writetable(rfrd, "xfel_" + name + ".dat", Delimiter='\t');
    writetable(rcol, "xfel_" + name + "_fit.dat", Delimiter='\t');
    %writetable(ref, "xfel_" + name + "_meas.dat", Delimiter='\t');
end
