% ### 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
CASE = 2;

% LMI FORM
LMI_FORM = 1;

% LMI TOLERANCE
epsilon = 1e-9;

% CHAIN LENGTH
N = 8; % Chain length

%%
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 = 1e3 * 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);
end

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

[P, Pinfo] = build_pll_model(G, Fr, Fd, Fn, Rolloff=1, TimeUnit="milliseconds");
P1 = getIOTransfer(P, ["w", "u"], ["r", "y", "e"]);
P1 = ss(zpk(P1));
P2 = getIOTransfer(P, ["r", "w(2)", "w(3)", "u"], ["y", "e"]);
P2 = ss(zpk(P2));
P2 = repmat({P2}, 1, N-1);

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

PchainT = Pchain;

if CASE == 0 %% this pole elimination seems illegal
    PchainT = minreal(PchainT);
    PchainT = canon(PchainT);
    PchainT = truncate_coeff(PchainT, 1e-8);
end

%PchainT = balreal(PchainT);
%PchainT = prescale(PchainT);
PchainT = output_state_transform(PchainT, 'e');
PchainT.C(end-ny+1:end,1:end-N) = 0;
PchainT = truncate_coeff(PchainT, 1e-10);
PchainT = prescale(PchainT);

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

A = PchainT.A;
Bu = PchainT.B(:,uidx); Bw = PchainT.B(:,widx);
Cy = PchainT.C(yidx, :); Cz = PchainT.C(zidx, :);
Dzu = PchainT.D(zidx, uidx); Dzw = PchainT.D(zidx, widx);
Dyu = PchainT.D(yidx, uidx); Duw = PchainT.D(yidx, widx);
assert(~any(Dzw, 'all'), "H2-Norm is infinite for systems with feedthrough!");
assert(~any(Duw, 'all'), "Framework assumes measurement noise is bandlimited!");
assert(~any(Dyu, 'all'), "Framework assumes plant has no feedthrough!");

S = [zeros(nu, nx-ny), eye(ny)];
T = S;
R = si_opt_r(T);
assert(isequal(R, R^(nx-1) > 0), 'R can be improved');

%% 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);
CL1 = systune(CL1, h2goal, [], systopts);
K1opt = CL1.Blocks.K.Gain.Value;
%fprintf("Baseline performance: %g\n", norm(lft(Pchain, diag(repelem(K1opt, N)))));

K2 = tunableGain("K", nu, ny);
K2.Gain.Free = eye(nu);
assert(nu == ny, "Assumes decentralized setup!");
CL2 = lft(Pchain, K2);
CL2 = systune(CL2, h2goal, [], systopts);
K2opt = CL2.Blocks.K.Gain.Value;
%fprintf("Optimal H2-Norm: %g\n", norm(lft(Pchain,K2opt)));

%% Solve P SI problem
sol = struct();
h2syn_lmi_opts = h2syn_lmi_options(eliminate_variables=false, sdp_options=sdpsettings('solver', 'mosek', 'debug', 1, 'verbose', 2));
h2syn_lmi_opts.tolerance = epsilon;
[sol.K, sol.CL, sol.gamma, sol.info] = h2syn_si(PchainT, ny, nu, 0, eye(ny), h2syn_lmi_opts);

%fprintf("SI problem status is %d\n", sol_si.sol_diag.problem);
%fprintf("SI controller gains: %s\n", num2str(diag(sol_si.K)'));
%fprintf("SI controller H2-Norm: %g\n", sol_si.J);
%fprintf("SI H2-Norm upper bound: %g\n", sol_si.gamma);

fprintf("Relative improvement: %gdB\n", mag2db(sol.info.J / norm(lft(Pchain, diag(repelem(K1opt, N))))));
fprintf("Relative improvement: %g%%\n", 100 * (sol.info.J / norm(lft(Pchain, diag(repelem(K1opt, N)))) - 1));
fprintf("Optimality gap: %g%%\n", 100 * (sol.info.J / norm(lft(Pchain, K2opt)) - 1));
fprintf("Bound gap: %gdB\n", mag2db(sol.gamma / sol.info.J));

baseline = sol.info.J;

sol_si = sol;

%% Solve PI SI problem
sol = struct();
h2syn_lmi_opts = h2syn_lmi_options(eliminate_variables=false, sdp_options=sdpsettings('solver', 'mosek', 'debug', 1, 'verbose', 2));
h2syn_lmi_opts.tolerance = 0;
[sol.K, sol.CL, sol.gamma, sol.info] = h2syn_si(PchainT, ny, nu, ny, kron(ones(ny), eye(ny)), h2syn_lmi_opts);
sol_si = sol;

%fprintf("SI problem status is %d\n", sol_si.sol_diag.problem);
%fprintf("SI controller gains: %s\n", num2str(diag(sol_si.K)'));
%fprintf("SI controller H2-Norm: %g\n", sol_si.J);
%fprintf("SI H2-Norm upper bound: %g\n", sol_si.gamma);

fprintf("Relative improvement: %gdB\n", mag2db(sol_si.J / norm(lft(Pchain, diag(repelem(K1opt, N))))));
fprintf("Relative improvement: %g%%\n", 100 * (sol_si.J / norm(lft(Pchain, diag(repelem(K1opt, N)))) - 1));
fprintf("Optimality gap: %g%%\n", 100 * (sol_si.J / norm(lft(Pchain, K2opt)) - 1));
fprintf("Bound gap: %gdB\n", mag2db(sol_si.gamma / sol_si.J));

baseline = sol_si.J;

%% SI Grid scan
epsilon = epsilon / 10;

switch CASE
    case 1
        Kgain = linspace(-20, 12, 11);
    case 2
        Kgain = linspace(-6, 3, 11);
end

[K1g, K2g] = meshgrid(Kgain, Kgain);

%- Y needs to be symmetric, no other way known to do that without loop.
%Y = blkdiag(sdpvar(nx-2), sdpvar(), sdpvar());
Y = sdpvar(nx);
for j = 1:size(R, 1)
    for k = 1:size(R, 2)
        if ~R(j,k)
            Y(j,k) = 0;
        end
    end
end; clear j k;

Z = sdpvar(nw); % (only used if Dzu is non-zero)

sdp_opts = sdpsettings('solver', 'mosek', 'debug', 0, 'verbose', 0);
%sdp_opts.mosek.MSK_IPAR_PRESOLVE_USE = 'MSK_OFF';
%sdp_opts.mosek.MSK_IPAR_INTPNT_SCALING = 'MSK_OFF';
%sdp_opts.mosek.MSK_DPAR_INTPNT_CO_TOL_INFEAS = 1e-18;
%sdp_opts.mosek.MSK_DPAR_INTPNT_CO_TOL_PFEAS = 1e-16;
%sdp_opts.mosek.MSK_DPAR_INTPNT_CO_TOL_DFEAS = 1e-16;
%sdp_opts.mosek.MSK_DPAR_INTPNT_CO_TOL_REL_GAP = 1e-18;

gamma = zeros(size(K1g));
prob_stat = zeros(size(K1g));

for i = 1:numel(K1g)
    k1g = db2mag(K1g(i));
    k2g = db2mag(K2g(i));
    switch CASE
        case 1
            Kout = sol_si.K * blkdiag(k1g, k2g);
        case 2
            Kout = blkdiag(K1opt, K1opt) * blkdiag(k1g, k2g);
    end
    Kstat = Kout * Cy;
    M = Kstat * Y;

    switch LMI_FORM
        case 1
            cnst = [Y >= epsilon * eye(nx);
                    (A + Bu*Kstat)*Y + Y*(A + Bu*Kstat)' + Bw*Bw' <= -epsilon * eye(nx)];
            obj = trace((Cz + Dzu*Kstat) * Y * (Cz + Dzu*Kstat)');
        case 2 % appears to make no numerical difference compared to 1
            cnst = [Y >= epsilon * eye(nx);
                    (A*Y + Bu*Kstat*Y) + (A*Y + Bu*Kstat*Y)' + Bw*Bw' <= -epsilon * eye(nx)];
            obj = trace((Cz + Dzu*Kstat) * Y * (Cz + Dzu*Kstat)');
        case 3 % appears to make no numerical difference compared to 1
            cnst = [Y >= epsilon * eye(nx);
                    [(A*Y + Bu*M) + (A*Y + Bu*M)', Bw;
                     Bw', -eye(nw)] <= -epsilon * eye(nx+nw)];
            obj = trace((Cz + Dzu*Kstat) * Y * (Cz + Dzu*Kstat)');
        case 4 % appears to make no numerical difference compared to 1
            cnst = [Y >= epsilon * eye(nx);
                    A*Y + Y*A' + Bu*M + M'*Bu' + Bw*Bw' <= -epsilon * eye(nx);
                    [Z, M;
                     M', Y] >= 0];
            obj = trace(Cz*Y*Cz' + Cz*M'*Dzu' + Dzu*M*Cz' + Dzu*Z*Dzu');
        case 5 % appears to make no numerical difference compared to 1
            cnst = [Y >= epsilon * eye(nx);
                    [A*Y + Y*A' + Bu*M + M'*Bu', Bw;
                     Bw', -eye(nw)] <= -epsilon * eye(nx+nw);
                    [Z, M;
                     M', Y] >= 0];
            obj = trace(Cz*Y*Cz' + Cz*M'*Dzu' + Dzu*M*Cz' + Dzu*Z*Dzu');
        case 6 % Careful here, Y is actually Y inverse in SI sense.
            cnst = [Y >= epsilon * eye(nx);
                    (A + Bu*Kstat)' * Y + Y * (A + Bu*Kstat) + Cz'*Cz <= -epsilon * eye(nx)];
            obj = trace(Bw'*Y*Bw);
        case 7 % Visible difference to 1
            cnst = [Y >= epsilon * eye(nx);
                    [Y*A' + Y*Kstat'*Bu' + A*Y + Bu*Kstat*Y, Y*Cz';
                     Cz*Y, -eye(nz)] <= -epsilon * eye(nx+nz);
                    [Y, Bw;
                     Bw', Z] >= 0];
            obj = trace(Z);
    end
    sol_diag = optimize(cnst, obj, sdp_opts);

    prob_stat(i) = sol_diag.problem;
    gamma(i) = sqrt(value(obj));
end
gamma(prob_stat == 1) = nan;

gamma_rel = mag2db(gamma ./ baseline);


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

figure();
switch CASE
    case 1
        lvls = (floor(a/step)*step):step:(ceil(b/step)*step);
        cdata = contour(K1g, K2g, gamma_rel, lvls); hold("on");
        lvls2 = lvls(1):step/10:lvls(2);
        lvls2 = setdiff(lvls2, lvls);
        cdata_fine = contour(K1g, K2g, gamma_rel, lvls2, '--');
        
        writematrix(cdata', "pll_chain_si_ex_a.dat", Delimiter='\t');
        writematrix(cdata_fine', "pll_chain_si_ex_a_fine.dat", Delimiter='\t');
    case 2
        lvls = (floor(a/step)*step):step:(ceil(b/step)*step);
        cdata = contour(K1g, K2g, gamma_rel, lvls); hold("on");
        step = step / 10;
        lvls2 = lvls(1):step:lvls(2);
        lvls2 = setdiff(lvls2, lvls);
        cdata_fine = contour(K1g, K2g, gamma_rel, lvls2, '--');
        
        writematrix(cdata', "pll_chain_si_ex_b.dat", Delimiter='\t');
        writematrix(cdata_fine', "pll_chain_si_ex_b_fine.dat", Delimiter='\t');
end

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