%% 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');

%% Tuv Stage
mn = 0.5; % Tuv mass [kg]

%% Sample
ms = 0.5; % Sample mass [kg]

%% General Configuration
model_config = struct();
model_config.controller = "open_loop"; % Default: Open-Loop
model_config.Tuv_type   = "parallel_k";    % Default: 2DoF stage

%% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/controller'],        1, 'openinput');  io_i = io_i + 1; % [Fu, Fv]
io(io_i) = linio([mdl, '/fd'],                1, 'openinput');  io_i = io_i + 1; % [Fdu, Fdv]
io(io_i) = linio([mdl, '/translation_stage'], 1, 'openoutput'); io_i = io_i + 1; % [Fmu, Fmv]
io(io_i) = linio([mdl, '/translation_stage'], 2, 'openoutput'); io_i = io_i + 1; % [Du, Dv]
io(io_i) = linio([mdl, '/ext_metrology'],     1, 'openoutput'); io_i = io_i + 1; % [Dx, Dy]

Wz = 0.1; % The rotation speed [rad/s]

%% No parallel Stiffness
kp = 0; % Parallel Stiffness [N/m]
cp = 0.001*2*sqrt(kp*(mn+ms)); % Small parallel damping [N/(m/s)]
kn = 1 - kp; % Stiffness [N/m]
cn = 0.01*2*sqrt(kn*(mn+ms)); % Damping [N/(m/s)]

G_no_kp = linearize(mdl, io, 0);
G_no_kp.InputName = {'Fu', 'Fv', 'Fdx', 'Fdy'};
G_no_kp.OutputName = {'fu', 'fv', 'Du', 'Dv', 'Dx', 'Dy'};

%% Small parallel Stiffness
kp = 0.5*(mn+ms)*Wz^2; % Parallel Stiffness [N/m]
cp = 0.001*2*sqrt(kp*(mn+ms)); % Small parallel damping [N/(m/s)]
kn = 1 - kp; % Stiffness [N/m]
cn = 0.01*2*sqrt(kn*(mn+ms)); % Damping [N/(m/s)]

G_low_kp = linearize(mdl, io, 0);
G_low_kp.InputName = {'Fu', 'Fv', 'Fdx', 'Fdy'};
G_low_kp.OutputName = {'fu', 'fv', 'Du', 'Dv', 'Dx', 'Dy'};

%% Large parallel Stiffness
kp = 1.5*(mn+ms)*Wz^2; % Parallel Stiffness [N/m]
cp = 0.001*2*sqrt(kp*(mn+ms)); % Small parallel damping [N/(m/s)]
kn = 1 - kp; % Stiffness [N/m]
cn = 0.01*2*sqrt(kn*(mn+ms)); % Damping [N/(m/s)]

G_high_kp = linearize(mdl, io, 0);
G_high_kp.InputName = {'Fu', 'Fv', 'Fdx', 'Fdy'};
G_high_kp.OutputName = {'fu', 'fv', 'Du', 'Dv', 'Dx', 'Dy'};

%% Effect of the parallel stiffness on the IFF plant
freqs = logspace(-2, 1, 1000);

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

% Magnitude
ax1 = nexttile([2, 1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G_no_kp(  'fu', 'Fu'), freqs, 'rad/s'))), '-', ...
     'DisplayName', '$k_p = 0$')
plot(freqs, abs(squeeze(freqresp(G_low_kp( 'fu', 'Fu'), freqs, 'rad/s'))), '-', ...
     'DisplayName', '$k_p < m\Omega^2$')
plot(freqs, abs(squeeze(freqresp(G_high_kp('fu', 'Fu'), freqs, 'rad/s'))), '-', ...
     'DisplayName', '$k_p > m\Omega^2$')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
set(gca, 'XTickLabel',[]); ylabel('Magnitude [N/N]');
ylim([1e-4, 5e1]);
legend('location', 'southeast', 'FontSize', 8);

% Phase
ax2 = nexttile;
hold on;
plot(freqs, 180/pi*angle(squeeze(freqresp(G_no_kp(  'fu', 'Fu'), freqs, 'rad/s'))), '-')
plot(freqs, 180/pi*angle(squeeze(freqresp(G_low_kp( 'fu', 'Fu'), freqs, 'rad/s'))), '-')
plot(freqs, 180/pi*angle(squeeze(freqresp(G_high_kp('fu', 'Fu'), freqs, 'rad/s'))), '-')
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [rad/s]'); ylabel('Phase [deg]');
yticks(-180:90:180);
ylim([0 180]);
hold off;
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)]);

%% Root Locus for IFF without parallel spring, with small parallel spring and with large parallel spring
gains = logspace(-2, 2, 200);

figure;
hold on;
plot(real(pole(G_no_kp({'fu','fv'},{'Fu','Fv'})*(1/s))),  imag(pole(G_no_kp({'fu','fv'},{'Fu','Fv'})*(1/s))), 'x', 'color', colors(1,:), ...
     'DisplayName', '$k_p = 0$','MarkerSize',8);
plot(real(tzero(G_no_kp({'fu','fv'},{'Fu','Fv'})*(1/s))),  imag(tzero(G_no_kp({'fu','fv'},{'Fu','Fv'})*(1/s))), 'o', 'color', colors(1,:), ...
     'HandleVisibility', 'off','MarkerSize',8);
for g = gains
    cl_poles = pole(feedback(G_no_kp({'fu','fv'},{'Fu','Fv'}), (g/s)*eye(2)));
    plot(real(cl_poles), imag(cl_poles), '.', 'color', colors(1,:), ...
         'HandleVisibility', 'off','MarkerSize',4);
end

plot(real(pole(G_low_kp({'fu','fv'},{'Fu','Fv'})*(1/s))),  imag(pole(G_low_kp({'fu','fv'},{'Fu','Fv'})*(1/s))), 'x', 'color', colors(2,:), ...
     'DisplayName', '$k_p < m\Omega^2$','MarkerSize',8);
plot(real(tzero(G_low_kp({'fu','fv'},{'Fu','Fv'})*(1/s))),  imag(tzero(G_low_kp({'fu','fv'},{'Fu','Fv'})*(1/s))), 'o', 'color', colors(2,:), ...
     'HandleVisibility', 'off','MarkerSize',8);
for g = gains
    cl_poles = pole(feedback(G_low_kp({'fu','fv'},{'Fu','Fv'}), (g/s)*eye(2)));
    plot(real(cl_poles), imag(cl_poles), '.', 'color', colors(2,:), ...
         'HandleVisibility', 'off','MarkerSize',4);
end

plot(real(pole(G_high_kp({'fu','fv'},{'Fu','Fv'})*(1/s))),  imag(pole(G_high_kp({'fu','fv'},{'Fu','Fv'})*(1/s))), 'x', 'color', colors(3,:), ...
     'DisplayName', '$k_p > m\Omega^2$','MarkerSize',8);
plot(real(tzero(G_high_kp({'fu','fv'},{'Fu','Fv'})*(1/s))),  imag(tzero(G_high_kp({'fu','fv'},{'Fu','Fv'})*(1/s))), 'o', 'color', colors(3,:), ...
     'HandleVisibility', 'off','MarkerSize',8);
for g = gains
    cl_poles = pole(feedback(G_high_kp({'fu','fv'},{'Fu','Fv'}), (g/s)*eye(2)));
    plot(real(cl_poles), imag(cl_poles), '.', 'color', colors(3,:), ...
         'HandleVisibility', 'off','MarkerSize',4);
end
hold off;
axis square;
xlim([-2.25, 0.25]); ylim([-1.25, 1.25]);
xticks([-2, -1, 0])
xticklabels({'$-2\omega_0$', '$-\omega_0$', '$0$'})
yticks([-1, 0, 1])
yticklabels({'$-\omega_0$', '$0$', '$\omega_0$'})

xlabel('Real Part'); ylabel('Imaginary Part');
leg = legend('location', 'northwest', 'FontSize', 8);
leg.ItemTokenSize(1) = 8;

%% Tested parallel stiffnesses
kps = [2, 20, 40]*(mn + ms)*Wz^2;

%% Root Locus: Effect of the parallel stiffness on the attainable damping
gains = logspace(-2, 4, 500);

figure;
hold on;
for kp_i = 1:length(kps)
    kp = kps(kp_i); % Parallel Stiffness [N/m]
    cp = 0.001*2*sqrt(kp*(mn+ms)); % Small parallel damping [N/(m/s)]
    kn = 1 - kp; % Stiffness [N/m]
    cn = 0.01*2*sqrt(kn*(mn+ms)); % Damping [N/(m/s)]

    % Identify dynamics
    G = linearize(mdl, io, 0);
    G.InputName = {'Fu', 'Fv', 'Fdx', 'Fdy'};
    G.OutputName = {'fu', 'fv', 'Du', 'Dv', 'Dx', 'Dy'};

    plot(real(pole(G({'fu', 'fv'}, {'Fu', 'Fv'})*(1/s*eye(2)))),  imag(pole(G({'fu', 'fv'}, {'Fu', 'Fv'})*(1/s*eye(2)))), 'x', 'color', colors(kp_i,:), ...
         'DisplayName', sprintf('$k_p = %.1f m \\Omega^2$', kp/((mn+ms)*Wz^2)),'MarkerSize',8);
    plot(real(tzero(G({'fu', 'fv'}, {'Fu', 'Fv'})*(1/s*eye(2)))),  imag(tzero(G({'fu', 'fv'}, {'Fu', 'Fv'})*(1/s*eye(2)))), 'o', 'color', colors(kp_i,:), ...
         'HandleVisibility', 'off','MarkerSize',8);
    for g = gains
        cl_poles = pole(feedback(G({'fu', 'fv'}, {'Fu', 'Fv'}), (g/s)*eye(2)));
        plot(real(cl_poles), imag(cl_poles), '.', 'color', colors(kp_i,:),'MarkerSize',4, ...
             'HandleVisibility', 'off');
    end
end
hold off;
axis square;
% xlim([-1.15, 0.05]); ylim([0, 1.2]);
xlim([-2.25, 0.25]); ylim([-1.25, 1.25]);
xticks([-2, -1, 0])
xticklabels({'$-2\omega_0$', '$-\omega_0$', '$0$'})
yticks([-1, 0, 1])
yticklabels({'$-\omega_0$', '$0$', '$\omega_0$'})

xlabel('Real Part'); ylabel('Imaginary Part');
leg = legend('location', 'northwest', 'FontSize', 8);
leg.ItemTokenSize(1) = 12;

%% Computes the optimal parameters and attainable simultaneous damping
alphas = logspace(-2, 0, 100);
alphas(end) = []; % Remove last point

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

Kiff = 1/s*eye(2);

for alpha_i = 1:length(alphas)
    kp = alphas(alpha_i);
    cp = 0.001*2*sqrt(kp*(mn+ms)); % Small parallel damping [N/(m/s)]
    kn = 1 - kp; % Stiffness [N/m]
    cn = 0.01*2*sqrt(kn*(mn+ms)); % Damping [N/(m/s)]

    % Identify dynamics
    G = linearize(mdl, io, 0);
    G.InputName = {'Fu', 'Fv', 'Fdx', 'Fdy'};
    G.OutputName = {'fu', 'fv', 'Du', 'Dv', 'Dx', 'Dy'};

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

    [g_opt, xi_opt] = fminsearch(fun, 2);
    opt_xi(alpha_i) = 1/xi_opt;
    opt_gain(alpha_i) = g_opt;
end

%% Attainable damping as a function of the stiffness ratio
figure;
yyaxis left
plot(alphas, opt_xi, '-', 'DisplayName', '$\xi_{cl}$');
set(gca, 'YScale', 'lin');
ylim([0,1]);
ylabel('Damping Ratio $\xi$');

yyaxis right
hold on;
plot(alphas, opt_gain, '-', 'DisplayName', '$g_{opt}$');
set(gca, 'YScale', 'lin');
ylim([0,2.5]);
ylabel('Controller gain $g$');

set(gca, 'XScale', 'log');
legend('location', 'northeast', 'FontSize', 8);

xlabel('$k_p$');
xlim([0.01, 1]);
xticks([0.01, 0.1, 1])
xticklabels({'$m\Omega^2$', '$10m\Omega^2$', '$100m\Omega^2$'})

%% Identify dynamics with parallel stiffness = 2mW^2
Wz = 0.1; % [rad/s]
kp = 2*(mn + ms)*Wz^2; % Parallel Stiffness [N/m]
cp = 0.001*2*sqrt(kp*(mn+ms)); % Small parallel damping [N/(m/s)]
kn = 1 - kp; % Stiffness [N/m]
cn = 0.01*2*sqrt(kn*(mn+ms)); % Damping [N/(m/s)]

% Identify dynamics
G = linearize(mdl, io, 0);
G.InputName = {'Fu', 'Fv', 'Fdx', 'Fdy'};
G.OutputName = {'fu', 'fv', 'Du', 'Dv', 'Dx', 'Dy'};

%% IFF controller with pure integrator
Kiff_kp = (2.2/s)*eye(2);
Kiff_kp.InputName = {'fu', 'fv'};
Kiff_kp.OutputName = {'Fu', 'Fv'};

%% Compute the damped plant
G_cl_iff_kp = feedback(G, Kiff_kp, 'name');

w0 = sqrt((kn+kp)/(mn+ms)); % Resonance frequency [rad/s]
wis = w0*logspace(-2, 0, 100); % LPF cut-off [rad/s]

%% Computes the obtained damping as a function of the HPF cut-off frequency
opt_xi = zeros(1, length(wis)); % Optimal simultaneous damping

for wi_i = 1:length(wis)
    Kiff_kp_hpf = (2.2/(s + wis(wi_i)))*eye(2);
    Kiff_kp_hpf.InputName = {'fu', 'fv'};
    Kiff_kp_hpf.OutputName = {'Fu', 'Fv'};

    [~, xi] = damp(feedback(G, Kiff_kp_hpf, 'name'));
    opt_xi(wi_i) = min(xi);
end

%% Effect of the  high-pass filter cut-off frequency on the obtained damping
figure;
plot(wis, opt_xi, '-');
set(gca, 'XScale', 'log');
set(gca, 'YScale', 'lin');
ylim([0,1]);
ylabel('Damping Ratio $\xi$');
xlabel('$\omega_i/\omega_0$');

%% Compute the damped plant with added  High-Pass Filter
Kiff_kp_hpf = (2.2/(s + 0.1*w0))*eye(2);
Kiff_kp_hpf.InputName = {'fu', 'fv'};
Kiff_kp_hpf.OutputName = {'Fu', 'Fv'};

G_cl_iff_hpf_kp = feedback(G, Kiff_kp_hpf, 'name');

%% Bode plot of the direct and coupling terms for several rotating velocities
freqs = logspace(-3, 1, 1000);

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

% Magnitude
ax1 = nexttile([2, 1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G(              'Du', 'Fu'), freqs, 'rad/s'))), '-', 'color', zeros( 1,3), ...
    'DisplayName', '$d_u/F_u$ - OL')
plot(freqs, abs(squeeze(freqresp(G_cl_iff_kp(    'Du', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(1,:), ...
    'DisplayName', '$d_u/F_u$ - IFF + $k_p$')
plot(freqs, abs(squeeze(freqresp(G_cl_iff_hpf_kp('Du', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(2,:), ...
    'DisplayName', '$d_u/F_u$ - IFF + $k_p$ + HPF')
plot(freqs, abs(squeeze(freqresp(G(              'Dv', 'Fu'), freqs, 'rad/s'))), '-', 'color', [zeros( 1,3), 0.5], ...
    'HandleVisibility', 'off')
plot(freqs, abs(squeeze(freqresp(G_cl_iff_kp(    'Dv', 'Fu'), freqs, 'rad/s'))), '-', 'color', [colors(1,:), 0.5], ...
    'HandleVisibility', 'off')
plot(freqs, abs(squeeze(freqresp(G_cl_iff_hpf_kp('Dv', 'Fu'), freqs, 'rad/s'))), '-', 'color', [colors(2,:), 0.5], ...
    'HandleVisibility', 'off')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
set(gca, 'XTickLabel',[]); ylabel('Magnitude [m/N]');
ldg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
ldg.ItemTokenSize(1) = 10;

ax2 = nexttile;
hold on;
plot(freqs, 180/pi*angle(squeeze(freqresp(G(              'Du', 'Fu'), freqs, 'rad/s'))), '-', 'color', zeros( 1,3))
plot(freqs, 180/pi*angle(squeeze(freqresp(G_cl_iff_kp(    'Du', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(1,:))
plot(freqs, 180/pi*angle(squeeze(freqresp(G_cl_iff_hpf_kp('Du', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(2,:))
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [rad/s]'); ylabel('Phase [deg]');
yticks(-180:90:180);
ylim([-180 90]);
xticks([1e-3,1e-2,1e-1,1,1e1])
xticklabels({'$0.001 \omega_0$', '$0.01 \omega_0$', '$0.1 \omega_0$', '$\omega_0$', '$10 \omega_0$'})

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