%% 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
addpath('./STEPS/'); % Path for STEPS
addpath('./subsystems/'); % Path for Subsystems Simulink files

%% Data directory
data_dir = './mat/';

% Simulink Model name
mdl = 'nass_model';

%% Colors for the figures
colors = colororder;

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

%% Identify the IFF plant dynamics using the Simscape model

% Initialize each Simscape model elements
initializeGround();
initializeGranite();
initializeTy();
initializeRy();
initializeRz();
initializeMicroHexapod();
initializeSimplifiedNanoHexapod();

% Initial Simscape Configuration
initializeSimscapeConfiguration('gravity', false);
initializeDisturbances('enable', false);
initializeLoggingConfiguration('log', 'none');
initializeController('type', 'open-loop');
initializeReferences();

% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/Controller'], 1, 'openinput');             io_i = io_i + 1; % Actuator Inputs [N]
io(io_i) = linio([mdl, '/NASS'],       3, 'openoutput', [], 'fn');  io_i = io_i + 1; % Force Sensors [N]

% Identify for multi payload masses (no rotation)
initializeReferences(); % No Spindle Rotation
% 1kg Sample
initializeSample('type', 'cylindrical', 'm', 1);
G_iff_m1 = linearize(mdl, io);
G_iff_m1.InputName  = {'f1',  'f2',  'f3',  'f4',  'f5',  'f6'};
G_iff_m1.OutputName = {'fn1', 'fn2', 'fn3', 'fn4', 'fn5', 'fn6'};

% 25kg Sample
initializeSample('type', 'cylindrical', 'm', 25);
G_iff_m25 = linearize(mdl, io);
G_iff_m25.InputName  = {'f1',  'f2',  'f3',  'f4',  'f5',  'f6'};
G_iff_m25.OutputName = {'fn1', 'fn2', 'fn3', 'fn4', 'fn5', 'fn6'};

% 50kg Sample
initializeSample('type', 'cylindrical', 'm', 50);
G_iff_m50 = linearize(mdl, io);
G_iff_m50.InputName  = {'f1',  'f2',  'f3',  'f4',  'f5',  'f6'};
G_iff_m50.OutputName = {'fn1', 'fn2', 'fn3', 'fn4', 'fn5', 'fn6'};

% Effect of Rotation
initializeReferences(...
    'Rz_type', 'rotating', ...
    'Rz_period', 1); % 360 deg/s
initializeSample('type', 'cylindrical', 'm', 25);
G_iff_m25_Rz = linearize(mdl, io, 0.1);
G_iff_m25_Rz.InputName  = {'f1',  'f2',  'f3',  'f4',  'f5',  'f6'};
G_iff_m25_Rz.OutputName = {'fn1', 'fn2', 'fn3', 'fn4', 'fn5', 'fn6'};

% Effect of Rotation - No added parallel stiffness
initializeSimplifiedNanoHexapod('actuator_kp', 0);
initializeReferences(...
    'Rz_type', 'rotating', ...
    'Rz_period', 1); % 360 deg/s
initializeSample('type', 'cylindrical', 'm', 25);
G_iff_m25_Rz_no_kp = linearize(mdl, io, 0.1);
G_iff_m25_Rz_no_kp.InputName  = {'f1',  'f2',  'f3',  'f4',  'f5',  'f6'};
G_iff_m25_Rz_no_kp.OutputName = {'fn1', 'fn2', 'fn3', 'fn4', 'fn5', 'fn6'};

%% IFF Plant - Without parallel stiffness
f = logspace(-1,3,1000);
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:5
    for j = i+1:6
        plot(f, abs(squeeze(freqresp(G_iff_m25_Rz_no_kp(i,j), f, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
             'HandleVisibility', 'off');
    end
end
plot(f, abs(squeeze(freqresp(G_iff_m25_Rz_no_kp(1,1), f, 'Hz'))), 'color', colors(1,:), ...
     'DisplayName', '$f_{ni}/f_i$ - $k_p = 0$')
for i = 2:6
    plot(f, abs(squeeze(freqresp(G_iff_m25_Rz_no_kp(i,i), f, 'Hz'))), 'color', colors(1,:), ...
         'HandleVisibility', 'off');
end
plot(f, abs(squeeze(freqresp(G_iff_m25_Rz_no_kp(1,2), f, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
     'DisplayName', '$f_{ni}/f_j$ - $k_p = 0$')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [N/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-4, 1e1]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
for i = 1:6
    plot(f, 180/pi*unwrap(angle(squeeze(freqresp(G_iff_m25_Rz_no_kp(i,i), f, 'Hz')))), 'color', colors(1,:));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([0:45:180]);

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

%% IFF Plant - With added parallel stiffness
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:5
    for j = i+1:6
        plot(f, abs(squeeze(freqresp(G_iff_m25_Rz(i,j), f, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
             'HandleVisibility', 'off');
    end
end
plot(f, abs(squeeze(freqresp(G_iff_m25_Rz(1,1), f, 'Hz'))), 'color', colors(1,:), ...
     'DisplayName', '$f_{ni}/f_i$ - $k_p = 50N/mm$')
for i = 2:6
    plot(f, abs(squeeze(freqresp(G_iff_m25_Rz(i,i), f, 'Hz'))), 'color', colors(1,:), ...
         'HandleVisibility', 'off');
end
plot(f, abs(squeeze(freqresp(G_iff_m25_Rz(1,2), f, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
     'DisplayName', '$f_{ni}/f_j$ - $k_p = 50N/mm$')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [N/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-4, 1e1]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
for i = 1:6
    plot(f, 180/pi*angle(squeeze(freqresp(G_iff_m25_Rz(i,i), f, 'Hz'))), 'color', colors(1,:));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([0:45:180]);

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

%% Effect of spindle's rotation on the IFF Plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:5
    for j = i+1:6
        plot(freqs, abs(squeeze(freqresp(G_iff_m25(i,j), freqs, 'Hz'))), 'color', [colors(1,:), 0.1], ...
             'HandleVisibility', 'off');
        plot(freqs, abs(squeeze(freqresp(G_iff_m25_Rz(i,j), freqs, 'Hz'))), 'color', [colors(2,:), 0.1], ...
             'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G_iff_m25(1,1), freqs, 'Hz'))), 'color', colors(1,:), ...
     'DisplayName', '$f_{ni}/f_i$ - $\Omega_z = 0$ deg/s')
plot(freqs, abs(squeeze(freqresp(G_iff_m25_Rz(1,1), freqs, 'Hz'))), 'color', colors(2,:), ...
     'DisplayName', '$f_{ni}/f_i$ - $\Omega_z = 360$ deg/s')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G_iff_m25(i,i), freqs, 'Hz'))), 'color', colors(1,:), ...
         'HandleVisibility', 'off');
    plot(freqs, abs(squeeze(freqresp(G_iff_m25_Rz(i,i), freqs, 'Hz'))), 'color', colors(2,:), ...
         'HandleVisibility', 'off');
end
% plot(freqs, abs(squeeze(freqresp(G_iff_m25_Rz(1,2), freqs, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
%      'DisplayName', '$f_{ni}/f_j$')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [N/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-4, 1e2]);
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_iff_m25(i,i), freqs, 'Hz'))), 'color', colors(1,:));
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_iff_m25_Rz(i,i), freqs, 'Hz'))), 'color', colors(2,:));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([0:45:180]);

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

%% Effect of the payload's mass on the IFF Plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G_iff_m1(1,1), freqs, 'Hz'))), 'color', [colors(1,:), 0.5], ...
     'DisplayName', '$f_{ni}/f_i$ - 1kg')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G_iff_m1(i,i), freqs, 'Hz'))), 'color', [colors(1,:), 0.5], ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(G_iff_m25(1,1), freqs, 'Hz'))), 'color', [colors(2,:), 0.5], ...
     'DisplayName', '$f_{ni}/f_i$ - 25kg')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G_iff_m25(i,i), freqs, 'Hz'))), 'color', [colors(2,:), 0.5], ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(G_iff_m50(1,1), freqs, 'Hz'))), 'color', [colors(3,:), 0.5], ...
     'DisplayName', '$f_{ni}/f_i$ - 50kg')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G_iff_m50(i,i), freqs, 'Hz'))), 'color', [colors(3,:), 0.5], ...
         'HandleVisibility', 'off');
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [N/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-4, 1e2]);
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_iff_m1(i,i), freqs, 'Hz'))), 'color', [colors(1,:), 0.5]);
end
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_iff_m25(i,i), freqs, 'Hz'))), 'color', [colors(2,:), 0.5]);
end
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_iff_m50(i,i), freqs, 'Hz'))), 'color', [colors(3,:), 0.5]);
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([0:45:180]);

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

%% Verify that parallel stiffness permits to have a stable plant
Kiff_pure_int = -200/s*eye(6);

if not(isstable(feedback(G_iff_m25_Rz, Kiff_pure_int, 1)))
    disp("Decentralized IFF is not stable with rotation")
end

if isstable(feedback(G_iff_m25_Rz_no_kp, Kiff_pure_int, 1))
    disp("Parallel stiffness makes the decentralized IFF stable")
else
    warning("Decentralized IFF is not stable even with the parallel stiffness")
end

%% IFF Controller Design
% Second order high pass filter
wz = 2*pi*2;
xiz = 0.7;
Ghpf = (s^2/wz^2)/(s^2/wz^2 + 2*xiz*s/wz + 1);

Kiff = -200 * ...              % Gain
       1/(0.01*2*pi + s) * ... % LPF: provides integral action
       Ghpf * ...              % 2nd order HPF (limit low frequency gain)
       eye(6);                 % Diagonal 6x6 controller (i.e. decentralized)

Kiff.InputName = {'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};
Kiff.OutputName = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};

% The designed IFF controller is saved
save('./mat/nass_K_iff.mat', 'Kiff');

%% Loop gain for the decentralized IFF
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
plot(freqs, abs(squeeze(freqresp(Kiff(1,1)*G_iff_m1(1,1), freqs, 'Hz'))), 'color', [colors(1,:), 0.5], ...
     'DisplayName', '1kg')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(Kiff(1,1)*G_iff_m1(i,i), freqs, 'Hz'))), 'color', [colors(1,:), 0.5], ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(Kiff(1,1)*G_iff_m25(1,1), freqs, 'Hz'))), 'color', [colors(2,:), 0.5], ...
     'DisplayName', '25kg')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(Kiff(1,1)*G_iff_m25(i,i), freqs, 'Hz'))), 'color', [colors(2,:), 0.5], ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(Kiff(1,1)*G_iff_m50(1,1), freqs, 'Hz'))), 'color', [colors(3,:), 0.5], ...
     'DisplayName', '50kg')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(Kiff(1,1)*G_iff_m50(i,i), freqs, 'Hz'))), 'color', [colors(3,:), 0.5], ...
         'HandleVisibility', 'off');
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Loop Gain'); set(gca, 'XTickLabel',[]);
ylim([1e-4, 1e2]);
leg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(-Kiff(1,1)*G_iff_m1(i,i), freqs, 'Hz'))), 'color', [colors(1,:), 0.5]);
end
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(-Kiff(1,1)*G_iff_m25(i,i), freqs, 'Hz'))), 'color', [colors(2,:), 0.5]);
end
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(-Kiff(1,1)*G_iff_m50(i,i), freqs, 'Hz'))), 'color', [colors(3,:), 0.5]);
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-110, 200]);
yticks([-180, -90, 0, 90, 180]);

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

%% Root Locus for the Decentralized IFF controller - 1kg Payload
figure;
gains = logspace(-2, 1, 200);

figure;
tiledlayout(1, 1, 'TileSpacing', 'compact', 'Padding', 'None');
nexttile();
hold on;

plot(real(pole(G_iff_m1)),  imag(pole(G_iff_m1)),  'x', 'color', colors(1,:), ...
    'DisplayName', '$g = 0$');
plot(real(tzero(G_iff_m1)), imag(tzero(G_iff_m1)), 'o', 'color', colors(1,:), ...
    'HandleVisibility', 'off');

for g = gains
    clpoles = pole(feedback(G_iff_m1, g*Kiff, +1));
    plot(real(clpoles), imag(clpoles), '.', 'color', colors(1,:), ...
        'HandleVisibility', 'off');
end

% Optimal gain
clpoles = pole(feedback(G_iff_m1, Kiff, +1));
plot(real(clpoles), imag(clpoles), 'kx', ...
    'DisplayName', '$g_{opt}$');

xline(0);
yline(0);
hold off;
axis equal;
xlim([-900, 100]); ylim([-100, 900]);
xticks([-900:100:0]);
yticks([0:100:900]);
set(gca, 'XTickLabel',[]); set(gca, 'YTickLabel',[]);
xlabel('Real part'); ylabel('Imaginary part');

%% Root Locus for the Decentralized IFF controller - 25kg Payload
gains = logspace(-2, 1, 200);

figure;
tiledlayout(1, 1, 'TileSpacing', 'compact', 'Padding', 'None');
nexttile();
hold on;

plot(real(pole(G_iff_m25)),  imag(pole(G_iff_m25)),  'x', 'color', colors(2,:), ...
    'DisplayName', '$g = 0$');
plot(real(tzero(G_iff_m25)), imag(tzero(G_iff_m25)), 'o', 'color', colors(2,:), ...
    'HandleVisibility', 'off');

for g = gains
    clpoles = pole(feedback(G_iff_m25, g*Kiff, +1));
    plot(real(clpoles), imag(clpoles), '.', 'color', colors(2,:), ...
        'HandleVisibility', 'off');
end

% Optimal gain
clpoles = pole(feedback(G_iff_m25, Kiff, +1));
plot(real(clpoles), imag(clpoles), 'kx', ...
    'DisplayName', '$g_{opt}$');

xline(0);
yline(0);
hold off;
axis equal;
xlim([-900, 100]); ylim([-100, 900]);
xticks([-900:100:0]);
yticks([0:100:900]);
set(gca, 'XTickLabel',[]); set(gca, 'YTickLabel',[]);
xlabel('Real part'); ylabel('Imaginary part');

%% Root Locus for the Decentralized IFF controller - 50kg Payload
gains = logspace(-2, 1, 200);

figure;
tiledlayout(1, 1, 'TileSpacing', 'compact', 'Padding', 'None');
nexttile();
hold on;

plot(real(pole(G_iff_m50)),  imag(pole(G_iff_m50)),  'x', 'color', colors(3,:), ...
    'DisplayName', '$g = 0$');
plot(real(tzero(G_iff_m50)), imag(tzero(G_iff_m50)), 'o', 'color', colors(3,:), ...
    'HandleVisibility', 'off');

for g = gains
    clpoles = pole(feedback(G_iff_m50, g*Kiff, +1));
    plot(real(clpoles), imag(clpoles), '.', 'color', colors(3,:), ...
        'HandleVisibility', 'off');
end

% Optimal gain
clpoles = pole(feedback(G_iff_m50, Kiff, +1));
plot(real(clpoles), imag(clpoles), 'kx', ...
    'DisplayName', '$g_{opt}$');

xline(0);
yline(0);
hold off;
axis equal;
xlim([-900, 100]); ylim([-100, 900]);
xticks([-900:100:0]);
yticks([0:100:900]);
set(gca, 'XTickLabel',[]); set(gca, 'YTickLabel',[]);
xlabel('Real part'); ylabel('Imaginary part');