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{:});
load('tmp.mat'); % solution of optimal synthesis.

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

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

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;

%% Heuristic Tunings
mlo = lbsync.xfel.laser.mlo1;
slo = lbsync.xfel.laser.slo1;
lonk = lbsync.xfel.link.slo1;
relonk = lbsync.xfel.link.relonk;

mlo.C = pid(2, 0.5e-3 * 1.3e9/16/48 * 3) / mlo.k_det / 2;
lonk.C = pid(0.03, 2e-3 * 5e6/5) / lonk.k_det;
slo.C = pid(5, 50e-3 * 1.3e9/16/48 * 3) / slo.k_det / 2;
relonk.C = pid(0.044, 4e-3 * 5e6/5) / relonk.k_det;

mlo = build_laser_model(mlo, Scaling=scaling, TimeUnit=time_unit);
slo = build_laser_model(slo, Scaling=scaling, TimeUnit=time_unit);
lonk = build_link_model(lonk, Scaling=scaling, TimeUnit=time_unit);
relonk = build_link_model(relonk, Scaling=scaling, TimeUnit=time_unit);

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;

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"]);
rcol_heur = compute_integrated_jitters(rcol);

%% Round-Trip Measurement
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);

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;

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"]);
rcol_bo = compute_integrated_jitters(rcol);

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

%% Optimal Synthesis
rcol = abs(squeeze(freqresp(sol.CL, w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "ReLONK", "ReLONKDetection", "RoundTripDetection", "TimingJitter"]);
rcol_syn = compute_integrated_jitters(rcol);

%%
clrs = colororder();
fig = figure(); ax = axes(fig, XScale="log", FontSize=12);
hold(ax, "on");
plot(ax, rcol_heur, "Frequency", "TimingJitter", LineWidth=1.5);
plot(ax, rcol_bo, "Frequency", "TimingJitter", LineWidth=1.5);
plot(ax, rcol_syn, "Frequency", "TimingJitter", LineWidth=1.5);
plot(ax, ref, "Frequency", "TimingJitter", Color="k");
ylabel(ax, "Timing Jitter Spectral Density (dBfs^2 / Hz)");
yyaxis(ax, "right");
ax.LineStyleOrder = "--";
ax.ColorOrder = clrs;
plot(ax, rcol_heur, "Frequency", "IntegratedTimingJitter", LineWidth=1.5);
plot(ax, rcol_bo, "Frequency", "IntegratedTimingJitter", LineWidth=1.5);
plot(ax, rcol_syn, "Frequency", "IntegratedTimingJitter", LineWidth=1.5);
plot(ax, ref, "Frequency", "IntegratedTimingJitter", Color="k");
ax.YColor = 'k';
legend(ax, ["Operator Tuning (sim)", "BayOpt Tuning (sim)", "Opt. Synthesis (sim)", "BayOpt Tuning (meas)"], Location="east");
xlabel(ax, "Fourier Frequency (Hz)");
ylabel(ax, "Integrated Timing Jitter (fs)");

%% Helper functions
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