% ### Daisy-Chained Laser PLL Phase Tracking Simulation ###
%
% SPDX-FileCopyrightText: 2023 Deutsches Elektronen-Synchrotron DESY, MSK, LbSync <lbsync-expert@desy.de>
% SPDX-License-Identifier: GPL-3.0-or-later
%
% Author: Maximilian Schütte <maximilian.schuette@desy.de>
%
% Control studies for chained VCO PLL phase tracking.
yalmip('clear'); clear;

% CASES
% 1 - Simple 2nd-order plant, no input delay
% 2 - Extra pole at 10kHz to induce finite gain margin
% 3 - 1us input delay to induce finite gain margin
CASE = 1;

drv_bw = 25e3; % 25 kHz driver bandwidth (typical DESY PZT4)
drv_k = 40; % 40 V/V gain (typical PZT4 gain in unipolar mode)
vco_k = 250; % (2*pi * k_omega=6.5 * n_harm=6) ~= 250

% Plant model in mrad @ 1.3GHz
ctrl_scale = 1;
G = zpk([], [0, -2*pi*drv_bw], (2*pi*drv_bw) * drv_k * vco_k);
if CASE == 2
    G = G * zpk([], -2*pi*10e3, 2*pi*10e3);
elseif CASE == 3
    G = G * pade(tf(1, InputDelay=10e-6), 1);
end

Fr = zpk([], 0, 2);
Fd = zpk(-2*pi*1e5, [0, -2*pi*5], 500 / (2*pi*1e5) * (2*pi*5));
Fn = tf(db2mag(-100));
%Fn = zpk([], -2*pi*1e4, 2*pi*1e4*db2mag(-100));

[P, Pinfo] = build_pll_model(G, Fr, Fd, Fn, Rolloff=1, TimeUnit="milliseconds");

P1 = getIOTransfer(P, ["w", "u"], ["r", "y", "e"]);
P2 = getIOTransfer(P, ["r", "w(2)", "w(3)", "u"], ["y", "e"]);

P1 = ss(zpk(P1));
P2 = ss(zpk(P2));

Pchain = chain(P1, P2);
Pchain = lft([1, -1], Pchain, 2, 0);
Pchain = prescale(Pchain);
Pchain.OutputName{1} = 'z';

%% Determine optimal gains to define grid range
h2goal = TuningGoal.Variance("w", "z", 1);
systopts = systuneOptions(RandomStart=20, Display="off");

K1 = tunableGain("K", 1, 1);
Ptmp = ss(zpk(getIOTransfer(P, ["w", "u"], ["z", "e"])));
CL1 = lft(Ptmp, K1);
CL1wy = lft(P1(2:3, :), K1);
CL1ry = lft(P2, K1);
CL1 = systune(CL1, h2goal, [], systopts);
K1opt = CL1.Blocks.K.Gain.Value;

K2 = tunableGain("K", 2, 2);
K2.Gain.Free = eye(2);
CL2 = lft(Pchain, K2);
CL2 = systune(CL2, h2goal, [], systopts);
K2opt = diag(CL2.Blocks.K.Gain.Value);

baseline = norm(lft(Pchain, blkdiag(K1opt, K1opt)));

%% Grid scan
if CASE == 1
    K = linspace(-12, 12, 51);
elseif CASE == 2
    K = linspace(-6, 3, 51);
elseif CASE == 3
    K = linspace(-6, 3, 51);
end

[K1, K2] = meshgrid(K, K);

gamma = zeros(size(K1));
for i = 1:numel(K1)
    k1 = db2mag(K1(i)) * K1opt;
    k2 = db2mag(K2(i)) * K1opt;
    gamma(i) = norm(lft(Pchain, blkdiag(k1, k2)));
end
gamma_rel = mag2db(gamma ./ baseline);

%% PLOTTING
figure();
bode(CL1); hold("on");
bode(setBlockValue(CL1, 'K', K1opt / 1e5));
bode(setBlockValue(CL1, 'K', K1opt * 1e5));

figure();
bode(CL1wy); hold("on");
bode(setBlockValue(CL1wy, 'K', K1opt / 1e2));
bode(setBlockValue(CL1wy, 'K', K1opt * 1e2));

figure();
bode(CL1ry); hold("on");
bode(setBlockValue(CL1ry, 'K', K1opt / 1e2));
bode(setBlockValue(CL1ry, 'K', K1opt * 1e2));

figure();
bode(CL2); hold("on");
bode(setBlockValue(CL2, 'K', blkdiag(K2opt(1) / 1e3, K2opt(2))));
bode(setBlockValue(CL2, 'K', blkdiag(K2opt(1) * 1e3, K2opt(2))));

figure();
bode(CL2); hold("on");
bode(setBlockValue(CL2, 'K', blkdiag(K2opt(1), K2opt(2) / 1e3)));
bode(setBlockValue(CL2, 'K', blkdiag(K2opt(1), K2opt(2) * 1e3)));

%%
figure();
CLx1 = ss(setBlockValue(CL2, 'K', blkdiag(K1opt, K1opt)));
CLx2 = ss(setBlockValue(CL2, 'K', blkdiag(K1opt * 4, K1opt / 1.8)));
bode(CLx1, CLx2);

figure();
bode(CLx1'*CLx1, CLx2'*CLx2);
figure();
bode(systrace(CLx1'*CLx1), systrace(CLx2'*CLx2));

%% Contour
a = min(gamma_rel(:));
b = max(gamma_rel(:));
n = 35;
step = round((b - a) / n, 1, "significant");

figure();
if CASE == 1
    lvls = (floor(-3/step)*step):step:(ceil(6/step)*step);
    cdata = contour(K1, K2, gamma_rel, lvls); hold("on");
    lvls2 = (floor(-3/step)*step):step/4:0;
    lvls2 = setdiff(lvls2, lvls);
    cdata_fine = contour(K1, K2, gamma_rel, lvls2, '--');

    writematrix(cdata', "pll_chain_ex_a.dat", Delimiter='\t');
    writematrix(cdata_fine', "pll_chain_ex_a_fine.dat", Delimiter='\t');
elseif CASE == 2
    step = 1;
    lvls = (floor(a/step)*step):step:(ceil(12/step)*step);
    cdata = contour(K1, K2, gamma_rel, lvls); hold("on");
    step = step / 10;
    lvls2 = (floor(a/step)*step):   step:-3;
    lvls2 = setdiff(lvls2, lvls);
    cdata_fine = contour(K1, K2, gamma_rel, lvls2, '--');

    writematrix(cdata', "pll_chain_ex_b.dat", Delimiter='\t');
    writematrix(cdata_fine', "pll_chain_ex_b_fine.dat", Delimiter='\t');
elseif CASE == 3
    lvls = (floor(a/step)*step):step:(ceil(12/step)*step);
    cdata = contour(K1, K2, gamma_rel, lvls); hold("on");
    step = step / 10;
    lvls2 = (floor(a/step)*step):step/4:0;
    lvls2 = setdiff(lvls2, lvls);
    cdata_fine = contour(K1, K2, gamma_rel, lvls2, '--');

    writematrix(cdata', "pll_chain_ex_c.dat", Delimiter='\t');
    writematrix(cdata_fine', "pll_chain_ex_c_fine.dat", Delimiter='\t');
end

hold("on");
scatter(0, 0, [], "red", '+');
