clear; close all;

% Load parameters for system
load('lbsync.mat');

sys = build_model(lbsync.lab);

G = sys.G;
C = pid(10, 0.1);

Fr = sys.Fr;
%F_r = chgTimeUnit(sys.F_r, 'milliseconds');
Fr.P{1}(1) = 0; % Add integral behaviour
Fr = ss(Fr);
%Fr.D = 0; % Explicitly strip feed-through for H2-Norm computation

Fd = sys.Fd;
%F_d = chgTimeUnit(sys.F_d, 'milliseconds');
Fd.P{1}(1) = 0; % Add integral behaviour
Fd = ss(Fd);
%Fd.D = 0; % Explicitly strip feed-through for H2-Norm computation

G.u = 'u';
G.y = 'y';
Fr.u = 'w(1)';
Fr.y = 'r';
Fd.u = 'w(2)';
Fd.y = 'd';
C.u = 'e';
C.y = 'u';

G_ss_gen = connect(G, Fr, Fd, sumblk('e = r - y - d'), {'u', 'w'}, 'e');
G_ss_gen = balreal(G_ss_gen);

kp = logspace(-2, 2, 100);
ki = linspace(0, 1e6, 20);
f = logspace(0, 6, 1e4);

H_r_pwr = squeeze(abs(freqresp(Fr, f, 'Hz')).^2);
H_d_pwr = squeeze(abs(freqresp(Fd, f, 'Hz')).^2);

[Kp, Ki] = meshgrid(kp, ki);
jitter = 0 * Kp;
h2norm = 0 * Kp;

%%
start_tic = tic;
last_tic = start_tic;
for i=1:numel(Kp)
    if toc(last_tic) > 10
        it_per_sec = i / toc(start_tic);
        fprintf('Now testing Kp = %g, Ki = %g. Average %g iterations per second. Estimated %g seconds till completion.\n', ...
            Kp(i), Ki(i), it_per_sec, (numel(Kp) - i) / it_per_sec);
        last_tic = tic;
    end

    C.Kp = Kp(i) / sys.k_phi;
    C.Ki = Ki(i) / sys.k_phi;

    if ~isstable(feedback(C * G, 1))
        jitter(i) = NaN;
        h2norm(i) = NaN;
        continue;
    end

    S = 1 / (1 + C * G);
    % T = C * G * S;
    S_pwr = squeeze(abs(freqresp(S, f, 'Hz')).^2);
    % T_pwr = squeeze(abs(freqresp(T, f, 'Hz')).^2);

    y = S_pwr .* (H_d_pwr + H_r_pwr);
    jitter(i) = sqrt(trapz(f, y));

    %cl = connect(G_ss_gen, C, 'w', 'e');
    %h2norm(i) = norm(cl);
end

surf(kp, ki, jitter, 'FaceColor', 'interp');
hold on;
surf(kp, ki, h2norm);
xlabel('P-Gain');
ylabel('I-Gain');
zlabel('Integrated Jitter 1Hz - 1MHz');
zlim([0 10]);
caxis([0 10]);