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

%%% CONFIGURATION %%% --------------------------------------------------------------------------------------------------
%% Set user-defined config parameters
time_unit  = "milliseconds";
scaling    = 15;             % model unit scale, 15 = femtosecond, etc.
pade_order = 2;
bandlimit  = inf;
min_decay  = 1;

% Other stuff used throughout the script
h2_goal = TuningGoal.Variance("w", "z", 1);


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

try
    load("models/lbsync_final.mat");
catch
    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{:});

    input_scaling = 1e-3;
    lbsync.xfel.laser.mlo1.input_scaling     = input_scaling;
    lbsync.xfel.laser.slo1.input_scaling     = input_scaling;
    lbsync.xfel.laser.sase3sxp.input_scaling = input_scaling;
    lbsync.xfel.link.slo1.input_scaling      = input_scaling;
    lbsync.xfel.link.relonk.input_scaling    = input_scaling;
    lbsync.xfel.link.sase3sxp.input_scaling  = input_scaling;
    clear input_scaling;
end

%% Build model
mo      = build_oscillator_model(lbsync.xfel.mo,        Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=1e6);
mlo     = build_laser_model(lbsync.xfel.laser.mlo1,     Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);
slo     = build_laser_model(lbsync.xfel.laser.slo1,     Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);
ppl     = build_laser_model(lbsync.xfel.laser.sase3sxp, Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);
lonk    = build_link_model(lbsync.xfel.link.slo1,       Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);
ppllink = build_link_model(lbsync.xfel.link.sase3sxp,   Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);

ppl_chain = chain2(mo.Pcol, mlo.Pcol, lonk.Pcol, slo.Pcol, ppllink.Pcol, ppl.Pcol, AnalysisPoint={'r(1)'});
ppl_chain = getIOTransfer(ppl_chain, {'w', 'u'}, {'r(1)', 'y', 'e'});
ppl_chain = connect(ppl_chain, sumblk("z = r(1) - y(1)"), {'w', 'u'}, {'z', 'e'});
ppl_chain = ss(ppl_chain);

sys = ppl_chain;
N = 5;

uidx = startsWith(sys.InputName, 'u'); nu = sum(uidx);
widx = startsWith(sys.InputName, 'w'); nw = sum(widx);
yidx = startsWith(sys.OutputName, 'e'); ny = sum(yidx);
zidx = startsWith(sys.OutputName, 'z'); nz = sum(zidx);
nx = order(sys);

sys_pade = pade(sys, pade_order);
nx_pade = order(sys_pade);

sys = prescale(sys);
sys_pade = prescale(sys_pade);

%% Synthesize Benchmark
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=25, Display="final");
ctrl = arrayfun(@(i) tunablePID("ctrl" + i, "PI"), 1:N, UniformOutput=false);
ctrl = blkdiag(ctrl{:});
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);

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

%%
rcol = abs(squeeze(freqresp(lft(sys, sol.K), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
fig = figure(); tiledlayout(3, 1); fig.Units = "centimeters"; fig.InnerPosition(3:4) = [20, 20];
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12"); ax1.FontSize = 12;
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12"); ax2.FontSize = 12;

plot(ax1, rcol, "Frequency", "IntegratedTimingJitter");
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

ylim([-80, inf]);

ylabel(ax2, "Timing Jitter Spectral Density (dBfs^2/Hz)");
ylabel(ax1, "Integrated Timing Jitter (fs)");
xlabel(ax1, "Fourier Frequency (Hz)");
legend(rcol.Properties.VariableNames(2:end));


%% Export
writetable(reducepoly_table(rcol, 1e-3), "xfel_lbsync_ppl_jitter_before.dat", delimiter="\t");

%hold(ax2, "on");
%for i = 2:(width(rcol) - 4)
%    tmp = rmmissing(rcol(:,[1,i]));
%    plot(ax2, tmp{:,1}, tmp{:,2}, LineStyle=":", Color="black");
%end

%% PID control / PI + Lead-Lag
%% 100 kHz Bw
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=25, Display="final");
ctrl = arrayfun(@(i) tunableSS("ctrl" + i, 2, 1, 1), 1:N, UniformOutput=false);
for i = 1:N, ctrl{i}.A.Value(1,1) = 0; ctrl{i}.A.Free = diag([0, 1]); ctrl{i}.A.Minimum = -2*pi*100e3; end
ctrl = blkdiag(ctrl{:});
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);

%%
rcol = abs(squeeze(freqresp(lft(sys, sol.K), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
fig = figure(); tiledlayout(3, 1); fig.Units = "centimeters"; fig.InnerPosition(3:4) = [20, 20];
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12"); ax1.FontSize = 12;
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12"); ax2.FontSize = 12;

plot(ax1, rcol, "Frequency", "IntegratedTimingJitter");
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

ylim([-80, inf]);

ylabel(ax2, "Timing Jitter Spectral Density (dBfs^2/Hz)");
ylabel(ax1, "Integrated Timing Jitter (fs)");
xlabel(ax1, "Fourier Frequency (Hz)");
legend(rcol.Properties.VariableNames(2:end));

%% Less PPL Loop Dly
new_ppl = lbsync.xfel.laser.sase3sxp; new_ppl.loop_dly = 2e-6; new_ppl.drv_lp_fc = 150e3;
new_ppl = build_laser_model(new_ppl, Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);

ppl_chain = chain2(mo.Pcol, mlo.Pcol, lonk.Pcol, slo.Pcol, ppllink.Pcol, new_ppl.Pcol, AnalysisPoint={'r(1)'});
ppl_chain = getIOTransfer(ppl_chain, {'w', 'u'}, {'r(1)', 'y', 'e'});
ppl_chain = connect(ppl_chain, sumblk("z = r(1) - y(1)"), {'w', 'u'}, {'z', 'e'});
ppl_chain = ss(ppl_chain);

sys = ppl_chain;

sys_pade = pade(sys, pade_order);
nx_pade = order(sys_pade);

sys = prescale(sys);
sys_pade = prescale(sys_pade);

%% 100 kHz Bw
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=25, Display="final");
ctrl = arrayfun(@(i) tunableSS("ctrl" + i, 2, 1, 1), 1:N, UniformOutput=false);
for i = 1:N, ctrl{i}.A.Value(1,1) = 0; ctrl{i}.A.Free = diag([0, 1]); ctrl{i}.A.Minimum = -2*pi*100e3; end
ctrl = blkdiag(ctrl{:});
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);

%%
rcol = abs(squeeze(freqresp(lft(sys, sol.K), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
fig = figure(); tiledlayout(3, 1); fig.Units = "centimeters"; fig.InnerPosition(3:4) = [20, 20];
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12"); ax1.FontSize = 12;
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12"); ax2.FontSize = 12;

plot(ax1, rcol, "Frequency", "IntegratedTimingJitter");
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

ylim([-80, inf]);

ylabel(ax2, "Timing Jitter Spectral Density (dBfs^2/Hz)");
ylabel(ax1, "Integrated Timing Jitter (fs)");
xlabel(ax1, "Fourier Frequency (Hz)");
legend(rcol.Properties.VariableNames(2:end));


%% Better MO
new_mo = lbsync.xfel.mo;
new_mo.Fd = chgTimeUnit(tf(db2mag(45), [1, 0]) + db2mag(-72), "milliseconds");
new_mo.Fd.InputName = "Fd"; new_mo.OutputName = "d";
new_mo = build_oscillator_model(new_mo, Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=1e6);

ppl_chain = chain2(new_mo.Pcol, mlo.Pcol, lonk.Pcol, slo.Pcol, ppllink.Pcol, new_ppl.Pcol, AnalysisPoint={'r(1)'});
ppl_chain = getIOTransfer(ppl_chain, {'w', 'u'}, {'r(1)', 'y', 'e'});
ppl_chain = connect(ppl_chain, sumblk("z = r(1) - y(1)"), {'w', 'u'}, {'z', 'e'});
ppl_chain = ss(ppl_chain);

sys = ppl_chain;

sys_pade = pade(sys, pade_order);
nx_pade = order(sys_pade);

sys = prescale(sys);
sys_pade = prescale(sys_pade);

%% 100 kHz Bw
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=25, Display="final");
ctrl = arrayfun(@(i) tunableSS("ctrl" + i, 2, 1, 1), 1:N, UniformOutput=false);
for i = 1:N, ctrl{i}.A.Value(1,1) = 0; ctrl{i}.A.Free = diag([0, 1]); ctrl{i}.A.Minimum = -2*pi*100e3; end
ctrl = blkdiag(ctrl{:});
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);

%%
rcol = abs(squeeze(freqresp(lft(sys, sol.K), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
fig = figure(); tiledlayout(3, 1); fig.Units = "centimeters"; fig.InnerPosition(3:4) = [20, 20];
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12"); ax1.FontSize = 12;
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12"); ax2.FontSize = 12;

plot(ax1, rcol, "Frequency", "IntegratedTimingJitter");
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

ylim([-80, inf]);

ylabel(ax2, "Timing Jitter Spectral Density (dBfs^2/Hz)");
ylabel(ax1, "Integrated Timing Jitter (fs)");
xlabel(ax1, "Fourier Frequency (Hz)");
legend(rcol.Properties.VariableNames(2:end));

%% Better MZM
new_mlo = lbsync.xfel.laser.mlo1; new_mlo.Fn.Numerator = 1e-3;
new_mlo = build_laser_model(new_mlo, Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);

ppl_chain = chain2(new_mo.Pcol, new_mlo.Pcol, lonk.Pcol, slo.Pcol, ppllink.Pcol, new_ppl.Pcol, AnalysisPoint={'r(1)'});
ppl_chain = getIOTransfer(ppl_chain, {'w', 'u'}, {'r(1)', 'y', 'e'});
ppl_chain = connect(ppl_chain, sumblk("z = r(1) - y(1)"), {'w', 'u'}, {'z', 'e'});
ppl_chain = ss(ppl_chain);

sys = ppl_chain;

sys_pade = pade(sys, pade_order);
nx_pade = order(sys_pade);

sys = prescale(sys);
sys_pade = prescale(sys_pade);

%% 100 kHz Bw
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=25, Display="final");
ctrl = arrayfun(@(i) tunableSS("ctrl" + i, 2, 1, 1), 1:N, UniformOutput=false);
for i = 1:N, ctrl{i}.A.Value(1,1) = 0; ctrl{i}.A.Free = diag([0, 1]); ctrl{i}.A.Minimum = -2*pi*100e3; end
ctrl = blkdiag(ctrl{:});
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);

%%
rcol = abs(squeeze(freqresp(lft(sys, sol.K), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
fig = figure(); tiledlayout(3, 1); fig.Units = "centimeters"; fig.InnerPosition(3:4) = [20, 20];
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12"); ax1.FontSize = 12;
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12"); ax2.FontSize = 12;

plot(ax1, rcol, "Frequency", "IntegratedTimingJitter");
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

ylim([-80, inf]);

ylabel(ax2, "Timing Jitter Spectral Density (dBfs^2/Hz)");
ylabel(ax1, "Integrated Timing Jitter (fs)");
xlabel(ax1, "Fourier Frequency (Hz)");
legend(rcol.Properties.VariableNames(2:end));


%% No Link PZ Resonance
new_lonk = lbsync.xfel.link.slo1; new_lonk.G_pz = tf(1, TimeUnit="milliseconds"); %new_lonk.drv_lp_fc = 100e3;
new_ppllink = lbsync.xfel.link.sase3sxp; new_ppllink.G_pz = tf(1, TimeUnit="milliseconds"); %new_ppllink.drv_lp_fc = 100e3;
new_lonk    = build_link_model(new_lonk, Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);
new_ppllink = build_link_model(new_ppllink, Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);

ppl_chain = chain2(new_mo.Pcol, new_mlo.Pcol, new_lonk.Pcol, slo.Pcol, new_ppllink.Pcol, new_ppl.Pcol, AnalysisPoint={'r(1)'});
ppl_chain = getIOTransfer(ppl_chain, {'w', 'u'}, {'r(1)', 'y', 'e'});
ppl_chain = connect(ppl_chain, sumblk("z = r(1) - y(1)"), {'w', 'u'}, {'z', 'e'});
ppl_chain = ss(ppl_chain);

sys = ppl_chain;

sys_pade = pade(sys, pade_order);
nx_pade = order(sys_pade);

sys = prescale(sys);
sys_pade = prescale(sys_pade);

%% 100 kHz Bw
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=25, Display="final");
ctrl = arrayfun(@(i) tunableSS("ctrl" + i, 2, 1, 1), 1:N, UniformOutput=false);
for i = 1:N, ctrl{i}.A.Value(1,1) = 0; ctrl{i}.A.Free = diag([0, 1]); ctrl{i}.A.Minimum = -2*pi*100e3; end
ctrl = blkdiag(ctrl{:});
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);

%%
rcol = abs(squeeze(freqresp(lft(sys, sol.K), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
fig = figure(); tiledlayout(3, 1); fig.Units = "centimeters"; fig.InnerPosition(3:4) = [20, 20];
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12"); ax1.FontSize = 12;
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12"); ax2.FontSize = 12;

plot(ax1, rcol, "Frequency", "IntegratedTimingJitter");
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

ylim([-80, inf]);

ylabel(ax2, "Timing Jitter Spectral Density (dBfs^2/Hz)");
ylabel(ax1, "Integrated Timing Jitter (fs)");
xlabel(ax1, "Fourier Frequency (Hz)");
legend(rcol.Properties.VariableNames(2:end));


%% Interconnected controllers
%% 100 kHz Bw
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=50, Display="final");

%ctrl = [tunableSS("ctrl1", 2, 1, 1), 0, tunableSS("ctrl2", 2, 1, 1), 0, tunableSS("ctrl3", 2, 1, 1);
%     0, tunableSS("ctrl10", 2, 1, 1), 0, 0, 0;
%     tunableSS("ctrl4", 2, 1, 1), 0, tunableSS("ctrl5", 2, 1, 1), 0, tunableSS("ctrl6", 2, 1, 1);
%     0, 0, 0, tunableSS("ctrl11", 2, 1, 1), 0;
%     tunableSS("ctrl7", 2, 1, 1), 0, tunableSS("ctrl8", 2, 1, 1), 0, tunableSS("ctrl9", 2, 1, 1);];
%for i = 1:11, ctrl.Blocks.("ctrl"+i).A.Value(1,1) = 0; ctrl.Blocks.("ctrl"+i).A.Free = diag([0, 1]); ctrl.Blocks.("ctrl"+i).A.Minimum = -2*pi*100e3; end

ctrl1 = tunableSS("ctrl1", 6, 1, 5); ctrl1.B.Free(:,[2,4]) = 0; ctrl1.B.Value(:,[2,4]) = 0;
ctrl2 = tunableSS("ctrl2", 2, 1, 5); ctrl2.B.Free(:,[1,3,4,5]) = 0; ctrl2.B.Value(:,[1,3,4,5]) = 0; 
ctrl3 = tunableSS("ctrl3", 6, 1, 5); ctrl3.B.Free(:,[2,4]) = 0; ctrl3.B.Value(:,[2,4]) = 0;
ctrl4 = tunableSS("ctrl4", 2, 1, 5); ctrl4.B.Free(:,[1,2,3,5]) = 0; ctrl4.B.Value(:,[1,2,3,5]) = 0; 
ctrl5 = tunableSS("ctrl5", 6, 1, 5); ctrl5.B.Free(:,[2,4]) = 0; ctrl5.B.Value(:,[2,4]) = 0;

ctrl1 = ctrl1 * blkdiag(1, 1, pade(tf(1,InputDelay=20e-6),1), 1, pade(tf(1,InputDelay=20e-6),1));
ctrl3 = ctrl3 * blkdiag(pade(tf(1,InputDelay=20e-6),1), 1, 1, 1, 1);
ctrl5 = ctrl5 * blkdiag(pade(tf(1,InputDelay=20e-6),1), 1, 1, 1, 1);


ctrl = [ctrl1;ctrl2;ctrl3;ctrl4;ctrl5];
for i = [2,4], ctrl.Blocks.("ctrl"+i).A.Value(1,1) = 0; ctrl.Blocks.("ctrl"+i).A.Free = diag([0, 1]); ctrl.Blocks.("ctrl"+i).A.Minimum = -2*pi*100e3; end

%ctrl = [tunableSS("ctrl1", 2, 1, 1), 0, tunableSS("ctrl2", 2, 1, 1), 0, tunableSS("ctrl3", 2, 1, 1);
%     0, tunableSS("ctrl10", 2, 1, 1), 0, 0, 0;
%     tunableSS("ctrl4", 2, 1, 1), 0, tunableSS("ctrl5", 2, 1, 1), 0, tunableSS("ctrl6", 2, 1, 1);
%     0, 0, 0, tunableSS("ctrl11", 2, 1, 1), 0;
%     tunableSS("ctrl7", 2, 1, 1), 0, tunableSS("ctrl8", 2, 1, 1), 0, tunableSS("ctrl9", 2, 1, 1);];
%for i = 1:11, ctrl.Blocks.("ctrl"+i).A.Value(1,1) = 0; ctrl.Blocks.("ctrl"+i).A.Free = diag([0, 1]); ctrl.Blocks.("ctrl"+i).A.Minimum = -2*pi*100e3; end

%ctrl = blkdiag(tunableSS("ctrl1", 8, 2, 2), tunableSS("ctrl2", 8, 2, 2), tunableSS("ctrl3", 2, 1, 1));
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [TuningGoal.ControllerPoles('ctrl1',0,1/sqrt(2),2*pi*100); TuningGoal.ControllerPoles('ctrl3',0,1/sqrt(2),2*pi*100); TuningGoal.ControllerPoles('ctrl5',0,1/sqrt(2),2*pi*100);], ns_opts);
%[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);

%%
rcol = abs(squeeze(freqresp(lft(sys, replaceBlock(ctrl, sol.info(2).Blocks)), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
fig = figure(); tiledlayout(3, 1); fig.Units = "centimeters"; fig.InnerPosition(3:4) = [20, 20];
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12"); ax1.FontSize = 12;
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12"); ax2.FontSize = 12;

plot(ax1, rcol, "Frequency", "IntegratedTimingJitter");
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

ylim([-80, inf]);

ylabel(ax2, "Timing Jitter Spectral Density (dBfs^2/Hz)");
ylabel(ax1, "Integrated Timing Jitter (fs)");
xlabel(ax1, "Fourier Frequency (Hz)");
legend(rcol.Properties.VariableNames(2:end));

%% Export
writetable(reducepoly_table(rcol, 1e-3), "xfel_lbsync_ppl_jitter_after.dat", delimiter="\t");

% hold(ax2, "on");
% for i = 2:(width(rcol) - 4)
%     tmp = rmmissing(rcol(:,[1,i]));
%     plot(ax2, tmp{:,1}, tmp{:,2}, LineStyle=":", Color="black");
% end

%% Short links
new_lonk = lbsync.xfel.link.slo1; new_lonk.G_pz = tf(1, TimeUnit="milliseconds"); new_lonk.link_len = 200;
new_lonk = build_link_model(new_lonk, Scaling=scaling, TimeUnit=time_unit, FilterMinDecay=min_decay, Bandlimit=bandlimit);

ppl_chain6 = chain2(new_mo.Pcol, new_mlo.Pcol, new_lonk.Pcol, slo.Pcol, new_ppllink.Pcol, new_ppl.Pcol, AnalysisPoint={'r(1)'});
ppl_chain6 = getIOTransfer(ppl_chain6, {'w', 'u'}, {'r(1)', 'y', 'e'});
ppl_chain6 = connect(ppl_chain6, sumblk("z = r(1) - y(1)"), {'w', 'u'}, {'z', 'e'});
ppl_chain6 = ss(ppl_chain6);

sys = ppl_chain6;

sys_pade = pade(sys, pade_order);
nx_pade = order(sys_pade);

sys = prescale(sys);
sys_pade = prescale(sys_pade);


%% 100 kHz Bw
sol = struct();
ns_opts = systuneOptions(MaxIter=1e3, RandomStart=25, Display="final");
ctrl = arrayfun(@(i) tunableSS("ctrl" + i, 2, 1, 1), 1:N, UniformOutput=false);
for i = 1:N, ctrl{i}.A.Value(1,1) = 0; ctrl{i}.A.Free = diag([0, 1]); ctrl{i}.A.Minimum = -2*pi*100e3; end
ctrl = blkdiag(ctrl{:});
ctrl = chgTimeUnit(ctrl, time_unit);
CL = lft(sys_pade, ctrl);

[sol.CL, sol.gamma, ~, sol.info] = systune(CL, h2_goal, [], ns_opts);
sol.K = replaceBlock(ctrl, sol.CL.Blocks);

print_jitter(sys, sol.K, sol.gamma);


%%
rcol = abs(squeeze(freqresp(lft(sys, sol.K), w)))';
rcol = array2table([f, mag2db(rcol), pow2db(sum(rcol.^2, 2))], ...
                   VariableNames=["Frequency", "MO", "MLO", "MLODetection", "LONK", "LONKDetection", "SLO", "SLODetection", "PPLLink", "PPLLinkDetection", "PPL", "PPLDetection", "TimingJitter"]);
rcol = compute_integrated_jitters(rcol);

%% Show Jitter
figure(); tiledlayout(3, 1);
ax1 = nexttile(); ax1.XScale = "log"; hold("on"); colororder("gem12");
ax2 = nexttile([2,1]); ax2.XScale = "log"; hold("on"); colororder("gem12");

plot(ax1, rcol, "Frequency", contains(rcol.Properties.VariableNames, "Integrated"));
plot(ax2, rcol, "Frequency", 2:(width(rcol)-4));
plot(ax2, rcol, "Frequency", "TimingJitter", Color="black", LineStyle="--");

%%% Functions %%% ------------------------------------------------------------------------------------------------------

function print_jitter(sys, ctrl, gamma, pade_order, range)
    arguments
        sys lti
        ctrl lti
        gamma (1,1) double = nan
        pade_order (1,1) {mustBeInteger,mustBePositive,mustBeFinite} = 1
        range (2,1) double {mustBeNonnegative,mustBeFinite} = [1e0, 1e6]
    end

    cl = lft(sys, ctrl);
    cl_pade = pade(cl, pade_order);

    if isfinite(gamma)
        fprintf("Guaranteed performance: %.3g\n", gamma * sqrt(1e3) / sqrt(2));
    end
    if isstable(cl_pade)
        fprintf("Padé-approximated performance: %.3g\n", norm(cl_pade) * sqrt(1e3) / sqrt(2));
    else
        fprintf("Warning: System not stable!")
    end
    fprintf("Integrated jitter 1Hz-1MHz: %.3g\n", integrate_freqresp(cl_pade, range));
    fprintf("Integrated jitter 1Hz-1MHz: %.3g\n", integrate_freqresp(cl, range));
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)];
    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)
    figure(); tiledlayout(2, 1);
    ax = nexttile();
    ax.XScale="log"; hold("on");
    plot(rfrd, "Frequency", contains(rfrd.Properties.VariableNames, "Integrated"));
    h = plot(rcol, "Frequency", contains(rcol.Properties.VariableNames, "Integrated"));
    for i = 1:numel(h), h(i).SeriesIndex = i; h(i).LineStyle = "--"; end
    plot(ref, "Frequency", contains(ref.Properties.VariableNames, "Integrated"), Color="black", DisplayName="Measured");

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