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

%% Plant using Analytical Equations
% Stewart platform definition
k = 1e6; % Actuator stiffness [N/m]
c = 1e1; % Actuator damping [N/(m/s)]

stewart = initializeSimplifiedNanoHexapod(...
    'Mpm', 1e-3, ...
    'actuator_type', '1dof', ...
    'actuator_k', k, ...
    'actuator_kp', 0, ...
    'actuator_c', c ...
);

% Payload: Cylinder
h = 300e-3; % Height of the cylinder [m]
r = 110e-3; % Radius of the cylinder [m]
m = 10; % Mass of the payload [kg]
initializeSample('type', 'cylindrical', 'm', m, 'H', h, 'R', r);

% Mass Matrix
M = zeros(6,6);
M(1,1) = m;
M(2,2) = m;
M(3,3) = m;
M(4,4) = 1/12*m*(3*r^2 + h^2);
M(5,5) = 1/12*m*(3*r^2 + h^2);
M(6,6) = 1/2*m*r^2;

% Stiffness and Damping matrices
K = k*eye(6);
C = c*eye(6);

% Compute plant in the frame of the struts
G_analytical = inv(ss(inv(stewart.geometry.J')*M*inv(stewart.geometry.J)*s^2 + C*s + K));

% Compare with Simscape model
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', [], 'dL');  io_i = io_i + 1; % Encoders [m]

G_simscape = linearize(mdl, io);
G_simscape.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G_simscape.OutputName = {'dL1', 'dL2', 'dL3', 'dL4', 'dL5', 'dL6'};

%% Comparison of the analytical transfer functions and the multi-body model
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:6
plot(freqs, abs(squeeze(freqresp(G_simscape(i,1), freqs, 'Hz'))), 'color', [colors(i,:), 0.5], 'linewidth', 2.5, ...
     'DisplayName', sprintf('$l_%i/f_1$ - Multi-Body', i))
end
for i = 1:6
plot(freqs, abs(squeeze(freqresp(G_analytical(i,1), freqs, 'Hz'))), '--', 'color', [colors(i,:)], ...
     'DisplayName', sprintf('$l_%i/f_1$ - Analytical', i))
end
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', 6, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

ax2 = nexttile;
hold on;
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_simscape(i,1), freqs, 'Hz'))), 'color', [colors(i,:),0.5], 'linewidth', 2.5);
end
for i = 1:6
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_analytical(i,1), 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)]);

%% Multi-Body model of the Nano-Hexapod
% Initialize 1DoF
initializeSimplifiedNanoHexapod('flex_type_F', '2dof', 'flex_type_M', '3dof', 'actuator_type', '1dof');
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', [], 'dL');  io_i = io_i + 1; % Encoders [m]
io(io_i) = linio([mdl, '/plant'],  2, 'openoutput', [], 'fn');  io_i = io_i + 1; % Force Sensors [N]

% With no payload
G = linearize(mdl, io);
G.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G.OutputName = {'dL1', 'dL2', 'dL3', 'dL4', 'dL5', 'dL6', ...
                'fn1', 'fn2', 'fn3', 'fn4', 'fn5', 'fn6'};

%% Multi-Body model of the Nano-Hexapod without parallel stiffness
% Initialize 1DoF
initializeSimplifiedNanoHexapod('flex_type_F', '2dof', 'flex_type_M', '3dof', 'actuator_type', '1dof', 'actuator_kp', 0);

% With no payload
G_no_kp = linearize(mdl, io);
G_no_kp.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G_no_kp.OutputName = {'dL1', 'dL2', 'dL3', 'dL4', 'dL5', 'dL6', ...
                      'fn1', 'fn2', 'fn3', 'fn4', 'fn5', 'fn6'};

%% Transfer function from actuator force inputs to displacement of each strut
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(i,j), freqs, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
             'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G(1,1), freqs, 'Hz'))), 'color', colors(1,:), ...
     'DisplayName', '$l_i/f_i$')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G(i,i), freqs, 'Hz'))), 'color', colors(1,:), ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(G(1,2), freqs, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
     'DisplayName', '$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(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)]);

%% Transfer function from actuator force inputs to force sensor in each strut
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(6+i,j), freqs, 'Hz'))), 'color', [0, 0, 0, 0.2], ...
             'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G(7,1), freqs, 'Hz'))), 'color', colors(1,:), ...
     'DisplayName', '$f_{ni}/f_i$')
plot(freqs, abs(squeeze(freqresp(G_no_kp(7,1), freqs, 'Hz'))), 'color', colors(2,:), ...
     'DisplayName', '$f_{ni}/f_i$ (no $k_p$)')
for i = 2:6
    plot(freqs, abs(squeeze(freqresp(G(6+i,i), freqs, 'Hz'))), 'color', colors(1,:), ...
         'HandleVisibility', 'off');
    plot(freqs, abs(squeeze(freqresp(G_no_kp(6+i,i), freqs, 'Hz'))), 'color', colors(2,:), ...
         'HandleVisibility', 'off');
end
plot(freqs, abs(squeeze(freqresp(G(7,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', '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(6+i,i), freqs, 'Hz'))), 'color', colors(1,:));
    plot(freqs, 180/pi*angle(squeeze(freqresp(G_no_kp(6+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)]);