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

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

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

%% Colors for the figures
colors = colororder;

%% Simscape model name
mdl = 'rotating_model';

%% Load "Generic" system dynamics
load('rotating_generic_plants.mat', 'Gs', 'Wzs');

%% Modified IFF - parameters
g = 2; % Controller gain
wi = 0.1; % HPF Cut-Off frequency [rad/s]

Kiff     = (g/s)*eye(2); % Pure IFF
Kiff_hpf = (g/(wi+s))*eye(2); % IFF with added HPF

%% Loop gain for the IFF with pure integrator and modified IFF with added  high-pass filter
freqs = logspace(-2, 1, 1000);
Wz_i = 2;

figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

% Magnitude
ax1 = nexttile([2, 1]);
hold on;
plot(freqs, abs(squeeze(freqresp(Gs{Wz_i}('fu', 'Fu')*Kiff(1,1), freqs, 'rad/s'))))
plot(freqs, abs(squeeze(freqresp(Gs{Wz_i}('fu', 'Fu')*Kiff_hpf(1,1), freqs, 'rad/s'))))
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
set(gca, 'XTickLabel',[]); ylabel('Loop Gain');

% Phase
ax2 = nexttile;
hold on;
plot(freqs, 180/pi*angle(squeeze(freqresp(Gs{Wz_i}('fu', 'Fu')*Kiff(1,1), freqs, 'rad/s'))), 'DisplayName', 'IFF')
plot(freqs, 180/pi*angle(squeeze(freqresp(Gs{Wz_i}('fu', 'Fu')*Kiff_hpf(1,1), freqs, 'rad/s'))), 'DisplayName', 'IFF,HPF')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [rad/s]'); ylabel('Phase [deg]');
yticks(-180:90:180);
ylim([-90 180]);
xticks([1e-2,1e-1,1,1e1])
xticklabels({'', '$0.1 \omega_0$', '$\omega_0$', '$10 \omega_0$'})
leg = legend('location', 'southwest', 'FontSize', 8);
leg.ItemTokenSize(1) = 15;

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

%%  High-Pass Filter Cut-Off Frequency
wis = [0.01, 0.05, 0.1]; % [rad/s]

%% Root Locus for the initial IFF and the modified IFF
gains = logspace(-2, 4, 200);

figure;
hold on;
for wi_i = 1:length(wis)
    wi = wis(wi_i);
    Kiff = (1/(wi+s))*eye(2);
    L(wi_i) = plot(nan, nan, '.', 'color', colors(wi_i,:), 'DisplayName', sprintf('$\\omega_i = %.2f \\omega_0$', wi));
    for g = gains
        clpoles = pole(feedback(Gs{2}({'fu', 'fv'}, {'Fu', 'Fv'}), g*Kiff));
        plot(real(clpoles), imag(clpoles), '.', 'color', colors(wi_i,:),'MarkerSize',4);
    end
    plot(real(pole(Gs{2}({'fu', 'fv'}, {'Fu', 'Fv'})*Kiff)), ...
         imag(pole(Gs{2}({'fu', 'fv'}, {'Fu', 'Fv'})*Kiff)), ...
         'x', 'color', colors(wi_i,:),'MarkerSize',8);
    plot(real(tzero(Gs{2}({'fu', 'fv'}, {'Fu', 'Fv'})*Kiff)), ...
         imag(tzero(Gs{2}({'fu', 'fv'}, {'Fu', 'Fv'})*Kiff)), ...
         'o', 'color', colors(wi_i,:),'MarkerSize',8);
end
hold off;
axis equal;
xlim([-2.3, 0.1]); ylim([-1.2, 1.2]);
xticks([-2:1:2]); yticks([-2:1:2]);
leg = legend(L, 'location', 'southwest', 'FontSize', 8);
leg.ItemTokenSize(1) = 8;
xlabel('Real Part'); ylabel('Imaginary Part');
clear L

xlim([-1.25, 0.25]); ylim([-1.25, 1.25]);
xticks([-1, 0])
xticklabels({'$-\omega_0$', '$0$'})
yticks([-1, 0, 1])
yticklabels({'$-\omega_0$', '$0$', '$\omega_0$'})
ytickangle(90)

%% Compute the optimal control gain
wis = logspace(-2, 1, 100); % [rad/s]

opt_xi = zeros(1, length(wis)); % Optimal simultaneous damping
opt_gain = zeros(1, length(wis)); % Corresponding optimal gain

for wi_i = 1:length(wis)
    wi = wis(wi_i);
    Kiff = 1/(s + wi)*eye(2);

    fun = @(g)computeSimultaneousDamping(g, Gs{2}({'fu', 'fv'}, {'Fu', 'Fv'}), Kiff);

    [g_opt, xi_opt] = fminsearch(fun, 0.5*wi*((1/Wzs(2))^2 - 1));
    opt_xi(wi_i) = 1/xi_opt;
    opt_gain(wi_i) = g_opt;
end

%% Attainable damping ratio as a function of wi/w0. Corresponding control gain g_opt and g_max are also shown
figure;
yyaxis left
plot(wis, opt_xi, '-', 'DisplayName', '$\xi_{cl}$');
set(gca, 'YScale', 'lin');
ylim([0,1]);
ylabel('Damping Ratio $\xi$');

yyaxis right
hold on;
plot(wis, opt_gain, '-', 'DisplayName', '$g_{opt}$');
plot(wis, wis*((1/Wzs(2))^2 - 1), '--', 'DisplayName', '$g_{max}$');
hold off;
set(gca, 'YScale', 'lin');
ylim([0,10]);
ylabel('Controller gain $g$');

xlabel('$\omega_i/\omega_0$');
set(gca, 'XScale', 'log');
legend('location', 'northeast', 'FontSize', 8);

%% Compute damped plant
wis = [0.03, 0.1, 0.5]; % [rad/s]
g = 2; % Gain

Gs_iff_hpf  = {};

for i = 1:length(wis)
    Kiff_hpf = (g/(wis(i)+s))*eye(2); % IFF with added HPF
    Kiff_hpf.InputName = {'fu', 'fv'};
    Kiff_hpf.OutputName = {'Fu', 'Fv'};

    Gs_iff_hpf(i) = {feedback(Gs{2}, Kiff_hpf, 'name')};
end

%% Effect of $\omega_i$ on the damped plant coupling
freqs = logspace(-2, 1, 1000);

figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

% Magnitude
ax1 = nexttile([2, 1]);
hold on;
for i = 1:length(wis)
    plot(freqs, abs(squeeze(freqresp(Gs_iff_hpf{i}('Du', 'Fu'), freqs, 'rad/s'))), '-', 'color', [colors(i,:)], ...
         'DisplayName', sprintf('$d_u/F_u$, $\\omega_i = %.2f \\omega_0$', wis(i)))
    plot(freqs, abs(squeeze(freqresp(Gs_iff_hpf{i}('Dv', 'Fu'), freqs, 'rad/s'))), '-', 'color', [colors(i,:), 0.5], ...
         'DisplayName', sprintf('$d_v/F_u$, $\\omega_i = %.2f \\omega_0$', wis(i)))
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
set(gca, 'XTickLabel',[]); ylabel('Magnitude [m/N]');
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 20;

ax2 = nexttile;
hold on;
for i = 1:length(wis)
    plot(freqs, 180/pi*angle(squeeze(freqresp(Gs_iff_hpf{i}('Du', 'Fu'), freqs, 'rad/s'))), '-')
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [rad/s]'); ylabel('Phase [deg]');
yticks(-180:90:180);
ylim([-180 0]);
xticks([1e-2,1e-1,1,1e1])
xticklabels({'$0.01 \omega_0$', '$0.1 \omega_0$', '$\omega_0$', '$10 \omega_0$'})

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

%% Effect of $\omega_i$ on the obtained compliance
freqs = logspace(-2, 1, 1000);

figure;
tiledlayout(1, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

% Magnitude
ax1 = nexttile();
hold on;
for i = 1:length(wis)
    plot(freqs, abs(squeeze(freqresp(Gs_iff_hpf{i}('Du', 'Fdx'), freqs, 'rad/s'))), '-', 'color', [colors(i,:)], ...
         'DisplayName', sprintf('$d_{x}/F_{dx}$, $\\omega_i = %.2f \\omega_0$', wis(i)))
end
plot(freqs, abs(squeeze(freqresp(Gs{2}('Du', 'Fdx'), freqs, 'rad/s'))), 'k--', ...
     'DisplayName', '$d_{x}/F_{dx}$, OL')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [rad/s]'); ylabel('Compliance [m/N]');
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 20;
xticks([1e-2,1e-1,1,1e1])
xticklabels({'$0.01 \omega_0$', '$0.1 \omega_0$', '$\omega_0$', '$10 \omega_0$'})