%% Clear Workspace and Close figures
clear; close all; clc;

%% Intialize Laplace variable
s = zpk('s');

%% Path for functions, data and scripts
addpath('./src/'); % Path for functions

%% Colors for the figures
colors = colororder;

%% Initialize Frequency Vector
freqs = logspace(-1, 3, 1000);

%% Analytical Complementary Filters - Effect of alpha
freqs_study = logspace(-2, 2, 1000);
alphas = [0.1, 1, 10];
w0 = 2*pi*1;
s = tf('s');

figure;
hold on;
for i = 1:length(alphas)
    alpha = alphas(i);
    Hh2 = (s/w0)^2*((s/w0)+1+alpha)/(((s/w0)+1)*((s/w0)^2 + alpha*(s/w0) + 1));
    Hl2 = ((1+alpha)*(s/w0)+1)/(((s/w0)+1)*((s/w0)^2 + alpha*(s/w0) + 1));
    plot(freqs_study, abs(squeeze(freqresp(Hh2, freqs_study, 'Hz'))), 'color', colors(i,:), 'DisplayName', sprintf('$\\alpha = %g$', alphas(i)));
    plot(freqs_study, abs(squeeze(freqresp(Hl2, freqs_study, 'Hz'))), 'color', colors(i,:), 'HandleVisibility', 'off');
end
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Relative Frequency $\frac{\omega}{\omega_0}$'); ylabel('Magnitude');
hold off;
ylim([1e-3, 20]);
leg = legend('location', 'northeast', 'FontSize', 8);
leg.ItemTokenSize(1) = 18;

%% Analytical Complementary Filters - Effect of w0
freqs_study = logspace(-1, 3, 1000);
alpha = [1];
w0s = [2*pi*1, 2*pi*10, 2*pi*100];
s = tf('s');

figure;
hold on;
for i = 1:length(w0s)
    w0 =w0s(i);
    Hh2 = (s/w0)^2*((s/w0)+1+alpha)/(((s/w0)+1)*((s/w0)^2 + alpha*(s/w0) + 1));
    Hl2 = ((1+alpha)*(s/w0)+1)/(((s/w0)+1)*((s/w0)^2 + alpha*(s/w0) + 1));
    plot(freqs_study, abs(squeeze(freqresp(Hh2, freqs_study, 'Hz'))), 'color', colors(i,:), 'DisplayName', sprintf('$\\omega_0 = %g$ Hz', w0/2/pi));
    plot(freqs_study, abs(squeeze(freqresp(Hl2, freqs_study, 'Hz'))), 'color', colors(i,:), 'HandleVisibility', 'off');
end
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [Hz]'); ylabel('Magnitude');
hold off;
xlim([freqs_study(1), freqs_study(end)]); ylim([1e-3, 20]);
leg = legend('location', 'southeast', 'FontSize', 8);
leg.ItemTokenSize(1) = 18;

%% Test model
freqs = logspace(0, 3, 1000); % Frequency Vector [Hz]

m = 20;  % mass [kg]
k = 1e6; % stiffness [N/m]
c = 1e2; % damping [N/(m/s)]

% Plant dynamics
G = 1/(m*s^2 + c*s + k);

% Uncertainty weight
wI = generateWF('n', 2, 'w0', 2*pi*50, 'G0', 0.1, 'Ginf', 10, 'Gc', 1);

%% Bode plot of the plant with dynamical uncertainty
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

% Magnitude
ax1 = nexttile([2,1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G, freqs, 'Hz'))), 'k-', 'DisplayName', 'G');
plotMagUncertainty(wI, freqs, 'G', G, 'DisplayName', '$\Pi_i$');
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Magnitude [m/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-8, 7e-5]);
hold off;
leg = legend('location', 'northeast', 'FontSize', 8);
leg.ItemTokenSize(1) = 18;

% Phase
ax2 = nexttile;
hold on;
plotPhaseUncertainty(wI, freqs, 'G', G);
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G, freqs, 'Hz')))), 'k-');
set(gca,'xscale','log');
yticks(-360:90:90);
ylim([-270 45]);
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
hold off;

linkaxes([ax1,ax2],'x');
xlim([freqs(1), freqs(end)]);

%% Analytical Complementary Filters
w0 = 2*pi*20;
alpha = 1;

Hh = (s/w0)^2*((s/w0)+1+alpha)/(((s/w0)+1)*((s/w0)^2 + alpha*(s/w0) + 1));
Hl = ((1+alpha)*(s/w0)+1)/(((s/w0)+1)*((s/w0)^2 + alpha*(s/w0) + 1));

%% Specifications
figure;
hold on;
plot([1, 100], [0.01, 100], ':', 'color', colors(2,:));
plot([300, 1000], [0.01, 0.01], ':', 'color', colors(1,:));
plot(freqs, 1./abs(squeeze(freqresp(wI, freqs, 'Hz'))), ':', 'color', colors(1,:));
plot(freqs, abs(squeeze(freqresp(Hl, freqs, 'Hz'))), '-', 'color', colors(1,:));
plot(freqs, abs(squeeze(freqresp(Hh, freqs, 'Hz'))), '-', 'color', colors(2,:));
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [Hz]'); ylabel('Magnitude');
hold off;
xlim([freqs(1), freqs(end)]);
ylim([1e-3, 10]);
xticks([0.1, 1, 10, 100, 1000]);

%% Obtained controller
omega = 2*pi*1000;

K = 1/(Hh*G) * 1/((1+s/omega+(s/omega)^2));
K = zpk(minreal(K));

%% Bode plot of the controller K
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

% Magnitude
ax1 = nexttile([2, 1]);
plot(freqs, abs(squeeze(freqresp(K*Hl, freqs, 'Hz'))), 'k-');
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Magnitude'); set(gca, 'XTickLabel',[]);
ylim([8e3, 1e8])

% Phase
ax2 = nexttile;
plot(freqs, 180/pi*angle(squeeze(freqresp(K*Hl, freqs, 'Hz'))), 'k-');
set(gca,'xscale','log');
yticks(-180:45:180);
ylim([-180 45]);
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');

linkaxes([ax1,ax2],'x');
xlim([freqs(1), freqs(end)]);

num_delta_points = 50;
theta = linspace(0, 2*pi, num_delta_points);
delta_points = exp(1j * theta);

% Get frequency responses for all components
G_resp  = squeeze(freqresp(G,  freqs, 'Hz'));
K_resp  = squeeze(freqresp(K,  freqs, 'Hz'));
Hl_resp = squeeze(freqresp(Hl, freqs, 'Hz'));
wI_resp = squeeze(freqresp(wI, freqs, 'Hz'));

% Calculate nominal responses
nom_L = G_resp .* K_resp .* Hl_resp;
nom_S = 1 ./ (1 + nom_L);
nom_T = nom_L ./ (1 + nom_L);

% Store all the points in the complex plane that L can take
loop_region_points = zeros(length(freqs), num_delta_points);

% Initialize arrays to store magnitude bounds
S_mag_min = ones(length(freqs), 1) * inf;
S_mag_max = zeros(length(freqs), 1);
T_mag_min = ones(length(freqs), 1) * inf;
T_mag_max = zeros(length(freqs), 1);

% Calculate magnitude bounds for all delta values
for i = 1:num_delta_points
    % Perturbed loop gain
    loop_perturbed = nom_L .* (1 + wI_resp .* delta_points(i));
    loop_region_points(:,i) = loop_perturbed;

    % Perturbed sensitivity function
    S_perturbed = 1 ./ (1 + loop_perturbed);
    S_mag = abs(S_perturbed);

    % Update S magnitude bounds
    S_mag_min = min(S_mag_min, S_mag);
    S_mag_max = max(S_mag_max, S_mag);

    % Perturbed complementary sensitivity function
    T_perturbed = loop_perturbed ./ (1 + loop_perturbed);
    T_mag = abs(T_perturbed);

    % Update T magnitude bounds
    T_mag_min = min(T_mag_min, T_mag);
    T_mag_max = max(T_mag_max, T_mag);
end

% At frequencies where |wI| > 1, T min is zero
T_mag_min(abs(wI_resp)>1) = 1e-10;

%% Nyquist plot to check Robust Stability
figure;
hold on;
plot(real(squeeze(freqresp(G*K*Hl, freqs, 'Hz'))), imag(squeeze(freqresp(G*K*Hl, freqs, 'Hz'))), 'k', 'DisplayName', '$L(j\omega)$ - Nominal');
plot(alphaShape(real(loop_region_points(:)), imag(loop_region_points(:)), 0.1), 'FaceColor', [0, 0, 0], 'EdgeColor', 'none', 'FaceAlpha', 0.3, 'DisplayName', '$L(j\omega)$ - $\forall G \in \Pi_i$');
plot(-1, 0, 'k+', 'MarkerSize', 5, 'HandleVisibility', 'off');
hold off;
grid on;
axis equal
xlim([-1.4, 0.2]); ylim([-1.2, 0.4]);
xticks(-1.4:0.2:0.2); yticks(-1.2:0.2:0.4);
xlabel('Real Part'); ylabel('Imaginary Part');
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 18;

%% Robust Performance
figure;
hold on;
plot(freqs, abs(nom_S), 'color', colors(2,:), 'DisplayName', '$|S|$ - Nom.');
plot(freqs, abs(nom_T), 'color', colors(1,:), 'DisplayName', '$|T|$ - Nom.');

patch([freqs, fliplr(freqs)], [S_mag_max', fliplr(S_mag_min')], colors(2,:), 'FaceAlpha', 0.2, 'EdgeColor', 'none', 'HandleVisibility', 'off');
patch([freqs, fliplr(freqs)], [T_mag_max', fliplr(T_mag_min')], colors(1,:), 'FaceAlpha', 0.2, 'EdgeColor', 'none', 'HandleVisibility', 'off');

plot([1, 100], [0.01, 100], ':', 'color', colors(2,:), 'DisplayName', 'Specs.');
plot([300, 1000], [0.01, 0.01], ':', 'color', colors(1,:), 'DisplayName', 'Specs.');
plot(freqs, 1./abs(squeeze(freqresp(wI, freqs, 'Hz'))), ':', 'color', colors(1,:), 'HandleVisibility', 'off');

set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
hold off;
xlabel('Frequency [Hz]'); ylabel('Magnitude');
xlim([freqs(1), freqs(end)]);
ylim([1e-4, 5]);
xticks([0.1, 1, 10, 100, 1000]);
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 3);
leg.ItemTokenSize(1) = 18;