% ### 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 sparsity invariance methods.

yalmip('clear');
%clear;

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

% LMI FORM
LMI_FORM = 1;


% 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, TimeUnit="milliseconds");
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 = canon(Mchain); % ensures integrator poles are exactly at zero
Mchain = prescale(Mchain);

MchainT = Mchain; % make copy for output-state transform

%MchainT = balreal(MchainT);
%MchainT = prescale(MchainT);
%MchainT = xperm(MchainT, [1:6,8:10,12:nx,11,7]);
[MchainT,Q] = output_state_transform(MchainT, 'e');
%MchainT.C(zidx | yidx,1:end-N) = 0;
%MchainT = truncate_coeff(MchainT, 1e-10);
MchainT = prescale(MchainT);

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=20, 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);

%% Determine SI optimal robust gains and performance bound and nominal performance and realistic worst-case performance estimate
[sol.K, sol.gamma, ...
 sol.J, sol.sol_diag] = h2syn_si_lpv(MchainT, Delta, BlkStruct, eye(ny), ...
                                     epsilon, solver='mosek', verbose=1, debug=1);
sol_si_lpv = sol;

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), sol_si_lpv.K));
    if J > J_worst, J_worst = J; end
end; clear i J;
warning("off", "Control:analysis:NormInfinite3");

fprintf("-- SI robust controller --\n");
fprintf("Controller gains: %s\n", num2str(diag(sol_si_lpv.K)'));
fprintf("Guaranteed H2-Norm: %g\n", sol_si_lpv.gamma);
fprintf("Nominal H2-Norm: %g\n", norm(setBlockValue(CL, 'K', sol_si_lpv.K)));
fprintf("Worst-case H2-Norm: %g\n", J_worst);

%% Scan optimal gains and performance for known uniform damping
%zeta = linspace(-1, 1, 50+1)';
Kopt_z = zeros(ny, ny, length(zeta));
Jz = zeros(size(zeta));
for i = 1:numel(zeta)
    CL = lft(lft(zeta(i) * eye(N), Mchain), K);
    CL = systune(CL, h2goal, [], systopts);
    Kopt_z(:,:,i) = CL.Blocks.K.Gain.Value;
    Jz(i) = norm(CL);
end; clear i;

%% Scan SI optimal gains and performance for known uniform damping
%zeta = linspace(-1, 1, 50+1);
for i = 1:numel(zeta)
    tmp = lft(zeta(i) * eye(N), Mchain);
    [tmp,Q] = output_state_transform(tmp, 'e');
    tmp = truncate_coeff(tmp, 1e-10);
    tmp = prescale(tmp);
    [sol.K, sol.gamma, sol.J, sol.sol_diag] = h2syn_si(tmp, eye(ny), 1e-9, verbose=0, debug=0);
    sol_si(i) = sol; %#ok<SAGROW>
end; clear i;
sol_si = sol_si';

figure(); hold("on");
plot(zeta, [Jz, [sol_si.J]']);
yyaxis("right");
plot(zeta, [sol_si.gamma]');

%% Scan SI robust gains and performance for varying lower/upper damping bound
nx = order(MchainT);
A = MchainT.A;
Bp = MchainT.B(:, pidx); Bu = MchainT.B(:, uidx); Bw = MchainT.B(:, widx);
Cq = MchainT.C(qidx, :); Cz = MchainT.C(zidx, :); Cy = MchainT.C(yidx, :);
Dqp = MchainT.D(qidx, pidx); Dqw = MchainT.D(qidx, widx); Dqu = MchainT.D(qidx, uidx);
Dzp = MchainT.D(zidx, pidx); Dzw = MchainT.D(zidx, widx); Dzu = MchainT.D(zidx, uidx);
Dyp = MchainT.D(yidx, pidx); Dyw = MchainT.D(yidx, widx); Dyu = MchainT.D(yidx, uidx);

% Verify
assert(~any(Dzw(:)), "H2-performance not finite with feedthrough from w -> z");
assert(~any(Dyp(:)), "Feedthrough in p -> e not allowed, find different LFR");
assert(~any(Dyw(:)), "Feedthrough in w -> e not allowed, add artifical pole");
assert(~any(Dyu(:)), "Feedthrough in u -> e not allowed, add artifical pole");

% Optimization Variables
%- Sparsity Invariance Structure Matrices
S = [zeros(nu, nx-ny), eye(N)];
T = S;
R = si_opt_r(T);
assert(isequal(R, R^(nx-1) > 0), "R can be improved");

%- Y is not symmetric, so we can use this more elegant, loop-less code.
[Ti, Tj] = find(T);
M = sparse(Ti, Tj, sdpvar(sum(T(:)), 1), nu, nx);
clear Ti Tj;

%- X needs to be symmetric, no other way known to do that without loop.
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;

%- H2-performance slack variable
Z = sdpvar(nz);

% Dualized "Stability" Multiplier
W = sdpvar(np); % TODO: Scherer, 1999 check if off-diagonal blocks of W are zero
U = sdpvar(np, nq, 'full');
V = sdpvar(nq);

% Dualized "Performance" Multiplier
Wp = sdpvar(np); % TODO: Scherer, 1999 check if off-diagonal blocks of W are zero
Up = sdpvar(np, nq, 'full');
Vp = sdpvar(nq);

% Solver Options
opt = sdpsettings('solver', 'mosek', 'verbose', 1, 'debug', 1);

% Objective
obj = trace(Z);

% Constraints
U_ = [-(A*Y + Bu*M)', -(Cq*Y + Dqu*M)';
      eye(nx), zeros(nx, nq);
      -Bp', -Dqp';
      zeros(nq, nx), eye(nq);
      -Bw', -Dqw'];
V_ = blkdiag(fOblkdiag(eye(nx)), ...
             [W, U; ...
              U', V], ...
             -eye(nw));

cnst = [Y >= epsilon * eye(size(Y));
        Z >= epsilon * eye(size(Z));
        U_' * V_ * U_ >= epsilon * eye(size(U_,2));
        [Y, -(Cq*Y + Dqu*M)', -(Cz*Y + Dzu*M)';
         -(Cq*Y + Dqu*M), Dqp*Wp*Dqp' - Dqp*Up - (Dqp*Up)' + Vp, Dqp*Wp*Dzp' - (Dzp*Up)';
         -(Cz*Y + Dzu*M), Dzp*Wp*Dqp' - Dzp*Up, Dzp*Wp*Dzp' + Z] >= epsilon * eye(nx+nq+nz);
        V >= epsilon * eye(size(V));
        Vp >= epsilon * eye(size(Vp))
       ];

M_ = [W, U;
      U', V];
Mp_ = [Wp, Up;
       Up', Vp];

%- Finite dimensional robustness verification
idx = [0, cumsum([BlkStruct.Occurrences])];
for i = 1:length(BlkStruct)
    cnst = [cnst;
            V(idx(i)+1:idx(i+1),idx(i)+1:idx(i+1)) >= 0 * epsilon * eye(BlkStruct(i).Occurrences);
            Vp(idx(i)+1:idx(i+1),idx(i)+1:idx(i+1)) >= 0 * epsilon * eye(BlkStruct(i).Occurrences)];
end; clear i idx;

gamma = zeros(size(zeta_lb));
for i = 1:numel(zeta_lb)
    if zeta_lb(i) > zeta_ub(i)
        gamma(i) = nan;
        continue;
    end

    % Generate convex hull of uncertainty set
    tmpl = arrayfun(@(x) ones(x, 2) * diag([zeta_lb(i), zeta_ub(i)]), [BlkStruct.Occurrences], 'UniformOutput', false);
    tmpl = combvec(tmpl{:});
    cnst_hull = [];
    for j = 1:size(tmpl,2)
        U_ = [eye(nq); -diag(tmpl(:,j))'];

        cnst_hull = [cnst_hull;
                     U_' * M_ * U_ <= -epsilon * eye(nq)]; % possibly non-strict
        cnst_hull = [cnst_hull;
                     U_' * Mp_ * U_ <= -epsilon * eye(nq)]; % possibly non-strict
    end; clear j;

    % Solve!
    sol_diag = optimize([cnst; cnst_hull], obj, opt);

    [primal_residual, dual_residual] = check(cnst);
    if sol_diag.problem == 0 && (any(primal_residual < -epsilon) || any(dual_residual < -epsilon))
        warning("Constraint violated beyond definitness margin!");
        sol_diag.problem = 8;
    end

    try
        K = value(M) / value(Y); % Recover controller.
        K = K(end-ny+1:end, end-ny+1:end);
        K = K ./ diag(Cy(:, end-ny+1:end))';
    catch
        warning("Unable to recover controller, check solver status.");
        K = [];
    end

    gamma(i) = sqrt(value(obj));
    if(~isreal(gamma(i))), gamma(i) = Inf; end

    %- Nominal Performance
    try
        P = lyap((A+Bu*K*Cy), Bw * Bw');
        Jnom = sqrt(trace((Cz+Dzu*K*Cy) * P * (Cz+Dzu*K*Cy)'));

        if(~isreal(Jnom)), Jnom = Inf; end
    catch
        Jnom = Inf;
    end
end; clear i;

figure(); contour(zeta_lb, zeta_ub, mag2db(gamma), 50);
axis("equal");


%% PLOTTING
clrs = colororder();
SHOW_POLYLINE_COMPRESSION = true;

benchmark = min([sol_si.gamma]);

%%
fig = figure();
ax = axes();

cdata = contour(normalized2actual(reso_dp, zeta), normalized2actual(reso_dp, zeta), mag2db(gamma/benchmark), 25);
writematrix(cdata', "pll_chain_rob_si_ex_lb_ub_grid.dat", Delimiter='\t');


%%
fig = figure();
ax = axes();
hold("on");

pll_chain_rob_si_ex = table(normalized2actual(reso_dp, zeta), mag2db([sol_si.gamma] / benchmark)', ...
                         mag2db(gamma(:,5) / benchmark), mag2db(gamma(:,15) / benchmark), mag2db(gamma(:,25) / benchmark), mag2db(gamma(:,40) / benchmark), ...
                         mag2db(gamma(10,:) / benchmark)', mag2db(gamma(25,:) / benchmark)', mag2db(gamma(35,:) / benchmark)', mag2db(gamma(45,:) / benchmark)', ...
                         VariableNames=["zeta", "gamma", ...
                         "gamma_z_lb_1", "gamma_z_lb_2", "gamma_z_lb_3", "gamma_z_lb_4", ...
                         "gamma_z_ub_1", "gamma_z_ub_2", "gamma_z_ub_3", "gamma_z_ub_4"]);
writetable(pll_chain_rob_si_ex, "pll_chain_nom_rob_si_ex.dat", Delimiter="\t");

t = pll_chain_rob_si_ex;
plot(t.zeta, t{:,["gamma", ...
                  "gamma_z_lb_1", "gamma_z_lb_2", "gamma_z_lb_3", "gamma_z_lb_4", ...
                 "gamma_z_ub_1", "gamma_z_ub_2", "gamma_z_ub_3", "gamma_z_ub_4"]});
clear t;

return;
%%


%% Code dump
% Add uncertain resonance with individual uncertainties
% M_list = cell(N, 1);
% Delta_list = cell(N, 1);
% BlkStruct_list = cell(N, 1);
% for i = 1:length(P_list)
%     tmp = P_list{i} * blkdiag(eye(size(P_list{i}, 2)-1), reso);
%     tmp.InputName = P_list{i}.InputName;
%     tmp.Blocks.damp.name = "damp" + i;
%     [tmpM, Delta_list{i}, BlkStruct_list{i}] = lftdata(tmp);
%     tmpM.InputName(ismissing(tmpM.InputName)) = cellstr(compose("p(%d)", 1:sum(ismissing(tmpM.InputName))));
%     tmpM.OutputName(ismissing(tmpM.OutputName)) = cellstr(compose("q(%d)", 1:sum(ismissing(tmpM.OutputName))));
%     M_list{i} = tmpM;
% end

%% Scan SI robust gains and performance for selected lower/upper damping bounds
zeta_lb = [-0.8, -0.6, -0.4];
gamma_lb = zeros(3,10);
for i = 1:numel(zeta_lb)
    zeta_ub = linspace(zeta_lb(i), 1, 10+1);
    for j = 1:numel(zeta_ub)
        % Generate convex hull of uncertainty set
        tmpl = arrayfun(@(x) ones(x, 2) * diag([zeta_lb(i), zeta_ub(j)]), [BlkStruct.Occurrences], 'UniformOutput', false);
        tmpl = combvec(tmpl{:});
        cnst_hull = [];
        for k = 1:size(tmpl,2)
            U_ = [eye(nq); -diag(tmpl(:,k))'];
    
            cnst_hull = [cnst_hull;
                         U_' * M_ * U_ <= -epsilon * eye(nq)]; % possibly non-strict
            cnst_hull = [cnst_hull;
                         U_' * Mp_ * U_ <= -epsilon * eye(nq)]; % possibly non-strict
        end; clear k;
    
        % Solve!
        sol_diag = optimize([cnst; cnst_hull], obj, opt);
    
        [primal_residual, dual_residual] = check(cnst);
        if sol_diag.problem == 0 && (any(primal_residual < -epsilon) || any(dual_residual < -epsilon))
            warning("Constraint violated beyond definitness margin!");
            sol_diag.problem = 8;
        end
    
        gamma(i,j) = sqrt(value(obj));
        if(~isreal(gamma(i,j))), gamma(i,j) = Inf; end
    end; clear j;
end; clear i;

zeta_ub = [0.4, 0.6, 0.8];
gamma_ub = zeros(3,10);
for i = 1:numel(zeta_ub)
    zeta_lb = linspace(-1, zeta_ub(i), 10+1);
    for j = 1:numel(zeta_ub)
        % Generate convex hull of uncertainty set
        tmpl = arrayfun(@(x) ones(x, 2) * diag([zeta_lb(j), zeta_ub(i)]), [BlkStruct.Occurrences], 'UniformOutput', false);
        tmpl = combvec(tmpl{:});
        cnst_hull = [];
        for k = 1:size(tmpl,2)
            U_ = [eye(nq); -diag(tmpl(:,k))'];
    
            cnst_hull = [cnst_hull;
                         U_' * M_ * U_ <= -epsilon * eye(nq)]; % possibly non-strict
            cnst_hull = [cnst_hull;
                         U_' * Mp_ * U_ <= -epsilon * eye(nq)]; % possibly non-strict
        end; clear k;
    
        % Solve!
        sol_diag = optimize([cnst; cnst_hull], obj, opt);
    
        [primal_residual, dual_residual] = check(cnst);
        if sol_diag.problem == 0 && (any(primal_residual < -epsilon) || any(dual_residual < -epsilon))
            warning("Constraint violated beyond definitness margin!");
            sol_diag.problem = 8;
        end
    
        gamma(i,j) = sqrt(value(obj));
        if(~isreal(gamma(i,j))), gamma(i,j) = Inf; end
    end; clear j;
end; clear i;

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