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

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

% Simulink Model name
mdl = 'nano_hexapod_model';

%% Colors for the figures
colors = colororder;

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

%% Identify plant from actuator forces to external metrology
stewart = initializeSimplifiedNanoHexapod();
initializeSample('type', 'cylindrical', 'm', 10, 'H', 300e-3);
initializeLoggingConfiguration('log', 'none');
initializeController('type', 'open-loop');

% 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, '/plant'],      1, 'openoutput');  io_i = io_i + 1; % External Metrology [m, rad]

% With no payload
G = linearize(mdl, io);
G.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G.OutputName = {'Dx', 'Dy', 'Dz', 'Rx', 'Ry', 'Rz'};

%% Plant in the Cartesian Frame
G_cart = G*inv(stewart.geometry.J');
G_cart.InputName  = {'Fx', 'Fy', 'Fz', 'Mx', 'My', 'Mz'};

%% Plant in the frame of the struts
G_struts = stewart.geometry.J*G;
G_struts.OutputName  = {'D1', 'D2', 'D3', 'D4', 'D5', 'D6'};

%% Bode plot of the plant projected in the frame of the struts
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_struts(i,j), freqs, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
             'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G_struts(1,1), freqs, 'Hz'))), 'color', colors(1,:), ...
     'DisplayName', '$-\epsilon_{\mathcal{L}i}/f_i$')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G_struts(i,i), freqs, 'Hz'))), 'color', colors(1,:), ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(G_struts(1,2), freqs, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
     'DisplayName', '$-\epsilon_{\mathcal{L}i}/f_j$')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-9, 1e-4]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

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

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

%% Bode plot of the plant projected in the Cartesian frame
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_cart(i,j), freqs, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
             'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G_cart(1,1), freqs, 'Hz'))), 'color', colors(1,:), ...
    'DisplayName', '$\epsilon_{D_x}/\mathcal{F}_x$ [m/N]')
plot(freqs, abs(squeeze(freqresp(G_cart(2,2), freqs, 'Hz'))), 'color', colors(2,:), ...
    'DisplayName', '$\epsilon_{D_y}/\mathcal{F}_y$ [m/N]')
plot(freqs, abs(squeeze(freqresp(G_cart(3,3), freqs, 'Hz'))), 'color', colors(3,:), ...
    'DisplayName', '$\epsilon_{D_z}/\mathcal{F}_z$ [m/N]')
plot(freqs, abs(squeeze(freqresp(G_cart(4,4), freqs, 'Hz'))), 'color', colors(4,:), ...
    'DisplayName', '$\epsilon_{R_x}/\mathcal{M}_x$ [rad/Nm]')
plot(freqs, abs(squeeze(freqresp(G_cart(5,5), freqs, 'Hz'))), 'color', colors(5,:), ...
    'DisplayName', '$\epsilon_{R_y}/\mathcal{M}_y$ [rad/Nm]')
plot(freqs, abs(squeeze(freqresp(G_cart(6,6), freqs, 'Hz'))), 'color', colors(6,:), ...
    'DisplayName', '$\epsilon_{R_z}/\mathcal{M}_z$ [rad/Nm]')
plot(freqs, abs(squeeze(freqresp(G_cart(1,5), freqs, 'Hz'))), 'color', [0, 0, 0, 0.5], ...
     'DisplayName', 'Coupling')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude'); set(gca, 'XTickLabel',[]);
ylim([1e-9, 4e-3]);
leg = legend('location', 'southwest', 'FontSize', 7, 'NumColumns', 3);
leg.ItemTokenSize(1) = 15;

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

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

%% Identify the IFF Plant
stewart = initializeSimplifiedNanoHexapod('actuator_kp', 0); % Ignoring parallel stiffness for now
initializeSample('type', 'cylindrical', 'm', 10, 'H', 300e-3);
initializeLoggingConfiguration('log', 'none');
initializeController('type', 'open-loop');

% 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, '/plant'],      2, 'openoutput', [], 'fn');  io_i = io_i + 1; % Force Sensors [N]

% With no payload
G_iff = linearize(mdl, io);
G_iff.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G_iff.OutputName = {'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};

%% IFF Controller Design
Kiff = -500/s * ... % 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'};

%% Root Locus plot of the Decentralized IFF Control
gains = logspace(-2, 1, 200);

figure;
tiledlayout(1, 1, 'TileSpacing', 'compact', 'Padding', 'None');
nexttile();
hold on;
plot(real(pole(G_iff)),  imag(pole(G_iff)),  'x', 'color', colors(1,:), ...
    'DisplayName', '$g = 0$');
plot(real(tzero(G_iff)), imag(tzero(G_iff)), 'o', 'color', colors(1,:), ...
    'HandleVisibility', 'off');

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

% Optimal gain
clpoles = pole(feedback(G_iff, Kiff, +1));
plot(real(clpoles), imag(clpoles), 'kx', ...
    'DisplayName', '$g_{opt}$');
hold off;
axis equal;
xlim([-600, 50]); ylim([-50, 600]);
xticks([-600:100:0]);
yticks([0:100:600]);
set(gca, 'XTickLabel',[]); set(gca, 'YTickLabel',[]);
xlabel('Real part'); ylabel('Imaginary part');

%% 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(-G_iff(1,1)*Kiff(1,1), freqs, 'Hz'))));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Loop Gain'); set(gca, 'XTickLabel',[]);
ylim([1e-2, 1e2]);
% leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
% leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
plot(freqs, 180/pi*angle(squeeze(freqresp(-G_iff(1,1)*Kiff(1,1), freqs, 'Hz'))));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
hold off;
yticks(-360:90:360);
ylim([-180, 180])

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

%% Identify the IFF Plant
initializeController('type', 'iff');

% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/Controller'], 1, 'input');       io_i = io_i + 1; % Actuator Inputs [N]
io(io_i) = linio([mdl, '/plant'],      1, 'openoutput');  io_i = io_i + 1; % External Metrology [m,rad]

% With no payload
G_hac = linearize(mdl, io);
G_hac.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G_hac.OutputName = {'Dx', 'Dy', 'Dz', 'Rx', 'Ry', 'Rz'};

%% Plant in the frame of the struts
G_hac_struts = stewart.geometry.J*G_hac;
G_hac_struts.OutputName  = {'D1', 'D2', 'D3', 'D4', 'D5', 'D6'};

%% Bode plot of the plant projected in the frame of the struts
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_struts(i,j), freqs, 'Hz'))), 'color', [0,0,0,0.1], ...
             'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G_struts(1,1), freqs, 'Hz'))), 'color', colors(1,:), ...
     'DisplayName', '$-\epsilon_{\mathcal{L}i}/f_i$')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G_struts(i,i), freqs, 'Hz'))), 'color', colors(1,:), ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(G_struts(1,2), freqs, 'Hz'))), 'color', [0,0,0,0.1], ...
     'DisplayName', '$-\epsilon_{\mathcal{L}i}/f_j$')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-9, 1e-4]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

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

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

%% Bode plot of the plant projected in the frame of the struts
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_hac_struts(i,j), freqs, 'Hz'))), 'color', [0,0,0,0.1], ...
             'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G_struts(1,1), freqs, 'Hz'))), 'color', [colors(1,:), 0.2], ...
     'DisplayName', '$-\epsilon_{\mathcal{L}i}/f_i$')
plot(freqs, abs(squeeze(freqresp(G_hac_struts(1,1), freqs, 'Hz'))), 'color', colors(2,:), ...
     'DisplayName', '$-\epsilon_{\mathcal{L}i}/f_i^\prime$')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G_struts(i,i), freqs, 'Hz'))), 'color', [colors(1,:), 0.2], ...
         'HandleVisibility', 'off');
    plot(freqs, abs(squeeze(freqresp(G_hac_struts(i,i), freqs, 'Hz'))), 'color', colors(2,:), ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(G_hac_struts(1,2), freqs, 'Hz'))), 'color', [0,0,0,0.1], ...
     'DisplayName', '$-\epsilon_{\mathcal{L}i}/f_j^\prime$')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-9, 1e-4]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_struts(i,i), freqs, 'Hz'))), 'color', [colors(1,:), 0.2]);
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_hac_struts(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([-180, 180]);
yticks([-180, -90, 0, 90, 180]);

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

%% High Authority Controller - Mid Stiffness Nano-Hexapod
% Wanted crossover
wc = 2*pi*20; % [rad/s]

% Integrator
H_int = wc/s;

% Lead to increase phase margin
a  = 2;  % Amount of phase lead / width of the phase lead / high frequency gain
H_lead = 1/sqrt(a)*(1 + s/(wc/sqrt(a)))/(1 + s/(wc*sqrt(a)));

% Low Pass filter to increase robustness
H_lpf = 1/(1 + s/2/pi/200);

% Gain to have unitary crossover at 5Hz
H_gain = 1./abs(evalfr(G_hac_struts(1, 1), 1j*wc));

% Decentralized HAC
Khac = H_gain  * ... % Gain
       H_int   * ... % Integrator
       H_lpf   * ... % Low Pass filter
       eye(6);       % 6x6 Diagonal

%% Plot of the eigenvalues of L in the complex plane
Ldet = zeros(6, length(freqs));
Lmimo = squeeze(freqresp(G_hac_struts*Khac, freqs, 'Hz'));
for i_f = 2:length(freqs)
    Ldet(:, i_f) = eig(squeeze(Lmimo(:,:,i_f)));
end
mod_margin = min(min(abs(Ldet + ones(size(Ldet)))));

figure;
hold on;
for i = 1:6
plot(real(squeeze(Ldet(i,:))), imag(squeeze(Ldet(i,:))), ...
        '.', 'color', colors(1, :), ...
        'HandleVisibility', 'off');
plot(real(squeeze(Ldet(i,:))), -imag(squeeze(Ldet(i,:))), ...
        '.', 'color', colors(1, :), ...
        'HandleVisibility', 'off');
end
plot(-1, 0, 'kx', 'HandleVisibility', 'off');
patch(-1 + mod_margin*cos([0:0.1:2*pi+0.1]), mod_margin*sin([0:0.1:2*pi+0.1]), colors(5,:), 'linestyle', '--', 'EdgeColor','black', 'FaceAlpha', 0.5, 'HandleVisibility', 'off');
text(-1,0.1, 'Robustness', 'FontSize', 8, 'horizontalalignment', 'center')
hold off;
set(gca, 'XScale', 'lin'); set(gca, 'YScale', 'lin');
xlabel('Real Part'); ylabel('Imaginary Part');
axis square
xlim([-1.8, 0.2]); ylim([-1, 1]);

%% Loop gain for the Decentralized HAC_IFF
i_fb = find(abs(squeeze(freqresp(G_hac_struts(1,1)*Khac(1,1), freqs, 'Hz')))<1, 1);

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

ax1 = nexttile([2,1]);
hold on;
patch([freqs(1:i_fb), freqs(i_fb), freqs(1)], [abs(squeeze(freqresp(G_hac_struts(1,1)*Khac(1,1), [freqs(1:i_fb)], 'Hz'))); 1; 1], colors(5,:), 'EdgeColor','none', 'FaceAlpha', 0.5)
plot(freqs, abs(squeeze(freqresp(G_hac_struts(1,1)*Khac(1,1), freqs, 'Hz'))));
text(1.2,2.5,{'Disturbance', 'rejection'}, 'FontSize', 8)
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Loop Gain'); set(gca, 'XTickLabel',[]);
ylim([1e-2, 1e2]);

ax2 = nexttile;
hold on;
plot(freqs, 180/pi*angle(squeeze(freqresp(G_hac_struts(1,1)*Khac(1,1), freqs, 'Hz'))));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
hold off;
yticks(-360:90:360);
ylim([-180, 180])

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