% ### Robust 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 with uncertain damping using IQC & nonsmooth methods.

%clear;

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

% SCALINGS
ctrl_scale = 1e-2; % to make solution ~1
noise_scale = 1e1; % to make nominal H2-norm ~1

% LMI TOLERANCE
% epsilon = 1e-9;

% CHAIN LENGTH
N = 3; % 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
reso_cf = 1e5; % mechanical resonance frequency
reso_dp = ureal("damp", 0.25, Range=[0.05, 0.5]); % 0.5 mechanical resonance frequency damping
reso_dp.NominalValue = mean(reso_dp.Range);
reso = tf([1, 2*2*pi*reso_cf, (2*pi*reso_cf)^2], [1, 2*reso_dp*(2*pi*reso_cf), (2*pi*reso_cf)^2], ...
          InputName="u", OutputName="uu"); % added later, so parameters are independent

% Plant model in mrad @ 1.3GHz
G = ctrl_scale * 1e3 * zpk([], [0, -2*pi*drv_bw], (2*pi*drv_bw) * drv_k * vco_k);
if CASE == 2
    G = G * zpk([], -2*pi*50e3, 2*pi*50e3);
end

Fr = noise_scale * zpk([], 0, 2);
Fd = noise_scale * zpk(-2*pi*1e5, [0, -2*pi*5], 500 / (2*pi*1e5) * (2*pi*5));
Fn = noise_scale * zpk([], -2*pi*1e6, 2*pi*1e6*db2mag(-100));  % full-block S-procedure doesn't allow feedthrough here.

% Build generalized VCO model
[P, Pinfo] = build_pll_model(G, Fr, Fd, Fn, Rolloff=10); %
reso = chgTimeUnit(reso, P.TimeUnit); % adapt TimeUnit of resonance (which is added later)

Pr = ss(getIOTransfer(P, 'w(1)', ["r", "r"])); Pr.OutputName = ["r", "y"]; % Equiv: Pr = [1; 1] * Pinfo.Fr; Pr.OutputName = ["r", "y"];
P_ = getIOTransfer(P, ["r", "w(2)", "w(3)", "u"], ["y", "e"]);

P.InputName{end} = 'uu';
P = connect(zpk(P), reso, ["w", "u"], P.OutputName);
[M, Delta, BlkStruct] = lftdata(P);
M = nameMissingIO(M, "p", "q");

P_.InputName{end} = 'uu';
P_ = connect(zpk(P_), reso, ["w", "r", "u"], P_.OutputName);
[M_, Delta_, BlkStruct_] = lftdata(P_);
M_ = nameMissingIO(M_, "p", "q");


% Construct repeated VCOs for daisy chaining
M_list = [{Pr}, repmat({M_}, 1, N)];

% Build chain interconnection
Mchain = chain(M_list{:});

% Network performance signal
Mchain = lft([1, -1], Mchain, 2, 0);
Mchain.OutputName{1} = 'z';

Mchain = connect(Mchain, ["p", "w", "u"], ["q", "z", "e"]); % Reorder inputs/outputs

pidx = startsWith(Mchain.InputName, 'p'); np = sum(pidx);
widx = startsWith(Mchain.InputName, 'w'); nw = sum(widx);
uidx = startsWith(Mchain.InputName, 'u'); nu = sum(uidx);
qidx = startsWith(Mchain.OutputName, 'q'); nq = sum(qidx);
zidx = startsWith(Mchain.OutputName, 'z'); nz = sum(zidx);
yidx = startsWith(Mchain.OutputName, 'e'); ny = sum(yidx);
nx = order(Mchain);

Mchain = prescale(Mchain);
Mchain = truncate_coeff(Mchain, 1e-9);
%Mchain = balreal(Mchain);

Delta = repmat({Delta}, [1,N]); Delta = blkdiag(Delta{:});
BlkStruct = repmat(BlkStruct, [1,N]);

% Get extremal uncertainty values and random samples in between
unc_hull_vrtx = arrayfun(@(x) ones(x, 2) * diag([-1, 1]), [BlkStruct.Occurrences], 'UniformOutput', false);
unc_hull_vrtx = combvec(unc_hull_vrtx{:});
%zeta = linspace(-1, 1, 50+1)';
%[zeta_lb, zeta_ub] = meshgrid(zeta);

%% Determine optimal nominal gains and performance and worst-case performance of nominal controller
h2goal = TuningGoal.Variance("w", "z", 1);
systopts = systuneOptions(RandomStart=50, Display="off");

K = tunableGain("K", nu, ny);
K.Gain.Free = eye(nu);
assert(nu == ny, "Assumes decentralized setup!");
CL = lft(lft(zeros(N), Mchain), K);
CL = systune(CL, h2goal, [], systopts);
Kopt = CL.Blocks.K.Gain.Value;

J_worst = 0;
warning("off", "Control:analysis:NormInfinite3");
for i = 1:size(unc_hull_vrtx, 2)
    J = norm(lft(lft(unc_hull_vrtx(:,i), Mchain), Kopt));
    if J > J_worst, J_worst = J; end
end; clear i J;
warning("on", "Control:analysis:NormInfinite3");

fprintf("-- Optimal nominal controler --\n");
fprintf("Optimal nominal controller: %s\n", num2str(diag(Kopt)'));
fprintf("Optimal nominal H2-Norm: %g\n", norm(lft(lft(zeros(N), Mchain), Kopt)));
fprintf("Worst-case H2-Norm: %g\n", J_worst);

%% Nonsmooth IQC Synthesis
opt = systuneOptions(RandomStart=20, Display='final');

% Controller
K = arrayfun(@(i) tunablePID("K" + i, 'P'), 1:N, UniformOutput=false); K = blkdiag(K{:});
MchainCL = lft(Mchain, K);

% Augmented Multiplier
Sigma = tf();
Psi = tf();
phi0 = tf();

%- Robustness
%-- Nx arbitrarily fast time-varying repeated parameter
D = arrayfun(@(i) realp("D" + i, 0), 1:N, UniformOutput=false);
for i = 1:N, D{i}.Minimum = 0; end; clear i;
D = blkdiag(D{:});
W = zeros(nq);

Sigma = append(Sigma, D);
phi0 = append(phi0, W);
Psi = append(Psi, -(D' * D) / 2);

%- Performance
%-- H2
PsiY = 0;
for i = 1:1
    a = realp("a" + i, 0);
    b = realp("b" + i, 0);
    X = realp("X" + i, zeros(nw));
    X.Free = eye(nw); % Try diagonal only, sufficient degrees of freedom
    Y = realp("Y" + i, zeros(nw));
    Y.Free = eye(nw); % same
    % a.Minimum = 0; % technically this should be enforced, but makes random sampling linear instead of log.
    a.Maximum = 2*pi*1e6; % Leads to a in [0, aMax] with log sampling.
    % b.Minimum = 0; % doesn't matter if positive or negative
    % abs(b) should be bounded, but this would make sampling linear instead
    % of log. In any case, sign may be changed by Y.
    b.Maximum = 1e3; % lower bound damping
    X.Minimum = 0;

    A = [-a, b;
         -b, -a];
    A = fKron(eye(nw), A);
    B = [0; 1];
    B = fKron(eye(nw), B);
    C = [Y, X];
    C = C(:, reshape([1:nw; nw+1:2*nw], [], 1)'); % interlace.

    PsiY = PsiY + ss(A, B, C, zeros(nw));
    %tmp2 = X * tf([1, a], [1, 2*a, a*a + b*b]) + Y * tf(b, [1, 2*a, a*a + b*b]);
end; clear X Y a b A B C;
PsiY.InputName = 'k'; PsiY.OutputName = 'l';

Sigma = append(Sigma, eye(nz));
phi0 = append(phi0, zeros(nw, nz));
Psi = append(Psi, -PsiY);


%% Goals
sigma_max_goal = TuningGoal.Gain('v', 'c', 1);
%sigma_max_goal.Focus = [0, 2*pi*1e6];
pos_real_goal = TuningGoal.Passivity('k', 'l');
h2_goal = HermIntegral('k', 'l', 1);


%% Direct Optimization
ns = size(Sigma, 1);
sys = sect(-[phi0 * MchainCL + Psi, zeros(np+nw, ns); ...
             Sigma * MchainCL, -eye(ns) / 2]);
sys.InputName = 'v'; sys.OutputName = 'c';

[sol, fsoft, ghard, info] = systune(append(sys, PsiY), h2_goal, [sigma_max_goal, pos_real_goal], opt);
softhard = [[info.fSoft]', reshape([info.gHard], 2, [])'];
%softhard = [[info.fSoft]', [info.gHard]'];
[softhardsorted, softhardsortedidx] = sortrows(softhard);


%%
sys = replaceBlock(sys, sol.Blocks);
MchainCL = replaceBlock(MchainCL, sol.Blocks);
D = replaceBlock(D, sol.Blocks);
D_ = double(D);
% W = replaceBlock(W, sol.Blocks);
W_ = double(W);
PsiY = replaceBlock(PsiY, sol.Blocks);
PsiY_ = ss(PsiY);
K_ = replaceBlock(K, sol.Blocks);

%% Some statistics
figure();
scatter(softhard(:,2), softhard(:,1), 12, 'filled');
xlim([min(min(softhard(:,2)), 1), min(max(softhard(:,2)), 1.01)]);
set(gca, YScale='log', XScale='log');
xlabel('Hard Constraint Objective');
ylabel('Soft Constraint Objective');
figure;
scatter(1:size(softhard, 1), softhardsorted(:,1), 12, 'filled');
xlabel('Run Number');
ylabel('Soft Constraint Objective');
set(gca, YScale='log');
yyaxis('right');
scatter(1:size(softhard, 1), softhardsorted(:,2), 12, 'filled');
ylabel('Hard Constraint Objective');
ylim([min(min(softhard(:,2)), 1), min(max(softhard(:,2)), 1.01)]);
set(gca, YScale='log');

%% Validate sigma max condition
figure(); sigma(sys);

%% Validate performance multiplier
MchainCL_d = cell(20,1);

for i = 2:length(MchainCL_d)-1
    MchainCL_d{i} = lft(diag(2 * rand(1, N) - 1), MchainCL);
    MchainCL_d{i} = systrace(MchainCL_d{i}' * MchainCL_d{i});
end
MchainCL_d{1} = lft(-1 * eye(N), MchainCL);
MchainCL_d{1} = systrace(MchainCL_d{1}' * MchainCL_d{1});
MchainCL_d{end} = lft(eye(N), MchainCL);
MchainCL_d{end} = systrace(MchainCL_d{end}' * MchainCL_d{end});

% Plot traces
figure; bode(MchainCL_d{:}, systrace(PsiY_ + PsiY_'));
figure; bode(systrace(PsiY_ + PsiY_') - MchainCL_d{1});

% Plot eigenvalues of performance FDI
mult = append(eye(nz), -(PsiY_ + PsiY_'));
test = MchainCL(np+1:end, nq+1:end);
fdisys = [ss(test); eye(nw)]' * mult * [ss(test); eye(nw)];
f = logspace(-2,6,1e4);
fdiresp = freqresp(fdisys, f);
assert(sum(diag(sum(abs(imag(fdiresp)), 3))) < 1e-6);
fdiresp = make_diag_real(fdiresp);
res = cell2mat(arrayfun(@(page) eig(fdiresp(:,:,page)), 1:size(fdiresp,3), 'UniformOutput', false));
res2 = cell2mat(arrayfun(@(page) out2(@() chol(-fdiresp(:,:,page))), 1:size(fdiresp,3), 'UniformOutput', false));
figure;
semilogx(f, sort(real(res))'); yyaxis('right');
semilogx(f, res2 > 0); ylim([-0.5, 1.5]); yyaxis('left');


%% Validate robustness multiplier
% Plot eigenvalues of robust stability FDI
mult = blkdiag(D_'*D_, -D_'*D_);
test = MchainCL(1:np, 1:np);
fdisys = [ss(test); eye(np)]' * mult * [ss(test); eye(np)];
f = logspace(-2,6,1e4);
fdiresp = freqresp(fdisys, f);
%assert(sum(diag(sum(abs(imag(fdiresp)), 3))) < 1e-6);
fdiresp = make_diag_real(fdiresp);
res = cell2mat(arrayfun(@(page) eig(fdiresp(:,:,page)), 1:size(fdiresp,3), 'UniformOutput', false));
res2 = cell2mat(arrayfun(@(page) out2(@() chol(-fdiresp(:,:,page))), 1:size(fdiresp,3), 'UniformOutput', false));
figure;
semilogx(f, sort(real(res))'); yyaxis('right');
semilogx(f, res2 > 0); ylim([-0.5, 1.5]); yyaxis('left');


%% Validate both
mult = append(D_'*D_, eye(nz), -D_'*D_, -PsiY_-PsiY_');
fdisys = [ss(MchainCL); eye(nw+np)]' * mult * [ss(MchainCL); eye(nw+np)];
f = logspace(-2,6,1e4);
fdiresp = freqresp(fdisys, f);
%assert(sum(diag(sum(abs(imag(fdiresp)), 3))) < 1e-6);
fdiresp = make_diag_real(fdiresp);
res = cell2mat(arrayfun(@(page) eig(fdiresp(:,:,page)), 1:size(fdiresp,3), 'UniformOutput', false));
res2 = cell2mat(arrayfun(@(page) out2(@() chol(-fdiresp(:,:,page))), 1:size(fdiresp,3), 'UniformOutput', false));
figure;
semilogx(f, sort(real(res))'); yyaxis('right');
semilogx(f, res2 > 0); ylim([-0.5, 1.5]); yyaxis('left');


return;

%% Validate multiplier KYP
yalmip('clear');
yalmip_opt = sdpsettings('solver', 'mosek', 'verbose', 1, 'debug', 1);
P = [];
Pvar = [];
PP = tf();
F = [];
Fvar = [];

%- Robustness
%-- Arbitrarily fast time-varying repeated parameter
tmp = [D_' * D_, W_';
       W_, -D_' * D_];
P = blkdiag(P, D);
PP = append(PP, blkdiag(eye(nq), eye(nq))); 

Dvar = diag(sdpvar(nq,1));
Wvar = zeros(nq);
Pvar = blkdiag(Pvar, [Dvar, Wvar; Wvar', -Dvar]);
Fvar = Fvar + [Dvar >= 0];


%- Performance

%- Connect
Gkyp = PP * [Gg_(1:nq, 1:np), Gg_(1:nq, np+1:np+nw);
             eye(np), zeros(np, nw);
             Gg_(nq+1:nq+nz, 1:np), Gg_(nq+1:nq+nz, np+1:np+nw);
             zeros(nz, np), eye(nw)];
N = order(Gkyp);

X = sdpvar(N);

U = [eye(N), zeros(N, size(Gkyp,2));
     Gkyp.A, Gkyp.B;
     Gkyp.C, Gkyp.D];
V = [zeros(N), X, zeros(N, size(P,2));
     X, zeros(N), zeros(N, size(P,2));
     zeros(size(P,1), N), zeros(size(P,1), N), P];
F = F + [U'*V*U <= 0];

sol_diag = optimize(F, [], yalmip_opt);
check(F)


%% Analyze closed loop
figure; bode(CL);


%% Cleanup
try
    autoArrangeFigures;
end

%% Utility
function sys = nameMissingIO(sys, in_name, out_name)
    sys.InputName(ismissing(sys.InputName)) = cellstr(compose("%s(%d)", in_name, 1:sum(ismissing(sys.InputName))));
    sys.OutputName(ismissing(sys.OutputName)) = cellstr(compose("%s(%d)", out_name, 1:sum(ismissing(sys.OutputName))));
end