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

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

%% Path for functions, data and scripts
addpath('./src/'); % Path for scripts
addpath('./mat/'); % Path for data
addpath('./STEPS/'); % Path for Simscape Model
addpath('./subsystems/'); % Path for Subsystems Simulink files

%% Linearization options
opts = linearizeOptions;
opts.SampleTime = 0;

%% Open Simscape Model
mdl = 'detail_fem_nass'; % Name of the Simulink File
open(mdl); % Open Simscape Model

%% Colors for the figures
colors = colororder;
freqs = logspace(1,3,500); % Frequency vector [Hz]

%% Identify the dynamics for several considered bending stiffnesses
% Let's initialize all the stages with default parameters.
initializeGround('type', 'rigid');
initializeGranite('type', 'rigid');
initializeTy('type', 'rigid');
initializeRy('type', 'rigid');
initializeRz('type', 'rigid');
initializeMicroHexapod('type', 'rigid');
initializeSample('m', 50);

initializeSimscapeConfiguration();
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
io(io_i) = linio([mdl, '/Tracking Error'], 1, 'openoutput', [], 'EdL');  io_i = io_i + 1; % Errors in the frame of the struts
io(io_i) = linio([mdl, '/NASS'],       3, 'openoutput', [], 'fn');  io_i = io_i + 1; % Force Sensors

% Effect of bending stiffness
Kf = [0, 50, 100, 500]; % [Nm/rad]
G_Kf = {zeros(length(Kf), 1)};

for i = 1:length(Kf)
    % Limited joint axial compliance
    initializeSimplifiedNanoHexapod('actuator_type', '1dof', ...
                          'flex_type_F', '2dof', ...
                          'flex_type_M', '3dof', ...
                          'actuator_k', 1e6, ...
                          'actuator_c', 1e1, ...
                          'actuator_kp', 0, ...
                          'actuator_cp', 0, ...
                          'Fsm', 56e-3, ... % APA300ML weight 112g
                          'Msm', 56e-3, ...
                          'Kf_F', Kf(i), ...
                          'Kf_M', Kf(i));

    G_Kf(i) = {linearize(mdl, io)};
    G_Kf{i}.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
    G_Kf{i}.OutputName = {'l1', 'l2', 'l3', 'l4', 'l5', 'l6', 'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};
end

freqs = logspace(0, 3, 1000);

%% Effect of the flexible joint bending stiffness on the HAC-plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:length(Kf)
    for j = 1:5
        for k = j+1:6
            plot(freqs, abs(squeeze(freqresp(G_Kf{i}("l"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(i,:), 0.1], ...
                 'HandleVisibility', 'off');
        end
    end
    plot(freqs, abs(squeeze(freqresp(G_Kf{i}("l1","f1"), freqs, 'Hz'))), 'color', colors(i,:), 'DisplayName', sprintf('$k_f = %.0f $ [Nm/rad]', Kf(i)));
    for j = 2:6
        plot(freqs, abs(squeeze(freqresp(G_Kf{i}("l"+j,"f"+j), freqs, 'Hz'))), 'color', colors(i,:), 'HandleVisibility', 'off');
    end
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
ylim([1e-10, 1e-4]);

ax2 = nexttile;
hold on;
for i = 1:length(Kf)
    plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_Kf{i}(1, 1), freqs, 'Hz')))));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-200, 20]);
yticks([-360:45:360]);

linkaxes([ax1,ax2],'x');

%% Effect of the flexible joint bending stiffness on the IFF plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:length(Kf)
    for j = 1:5
        for k = j+1:6
            plot(freqs, abs(squeeze(freqresp(G_Kf{i}("fm"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(i,:), 0.1], ...
                 'HandleVisibility', 'off');
        end
    end
    plot(freqs, abs(squeeze(freqresp(G_Kf{i}("fm1","f1"), freqs, 'Hz'))), 'color', colors(i,:), 'DisplayName', sprintf('$k_f = %.0f $ [Nm/rad]', Kf(i)));
    for j = 2:6
        plot(freqs, abs(squeeze(freqresp(G_Kf{i}("fm"+j,"f"+j), freqs, 'Hz'))), 'color', colors(i,:), 'HandleVisibility', 'off');
    end
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [N/N]'); set(gca, 'XTickLabel',[]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
ylim([1e-4, 1e2]);

ax2 = nexttile();
hold on;
for i = 1:length(Kf)
    plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_Kf{i}("fm1", "f1"), freqs, 'Hz')))));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([-360:45:360]);

linkaxes([ax1,ax2],'x');

%% Decentalized IFF
Kiff = -200 * ...              % Gain
       1/s * ... % LPF: provides integral action
       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 for decentralized IFF - 1dof actuator - Effect of joint bending stiffness
gains = logspace(-1, 2, 400);

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

for i = 1:length(Kf)
    plot(real(pole(G_Kf{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))),  imag(pole(G_Kf{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))),  'x', 'color', colors(i,:), ...
        'DisplayName', sprintf('$k_f = %.0f$ Nm/rad', Kf(i)));
    plot(real(tzero(G_Kf{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))), imag(tzero(G_Kf{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))), 'o', 'color', colors(i,:), ...
        'HandleVisibility', 'off');

    for g = gains
        clpoles = pole(feedback(G_Kf{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}), g*Kiff, +1));
        plot(real(clpoles), imag(clpoles), '.', 'color', colors(i,:), ...
            'HandleVisibility', 'off');
    end

end

xline(0, 'HandleVisibility', 'off'); yline(0, 'HandleVisibility', 'off');
hold off;
axis equal;
xlim(1.1*[-900, 100]); ylim(1.1*[-100, 900]);
xticks(1.1*[-900:100:0]);
yticks(1.1*[0:100:900]);
set(gca, 'XTickLabel',[]); set(gca, 'YTickLabel',[]);
xlabel('Real part'); ylabel('Imaginary part');
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Identify the dynamics for several considered bending stiffnesses - APA300ML
G_Kf_apa300ml = {zeros(length(Kf), 1)};

for i = 1:length(Kf)
    % Limited joint axial compliance
    initializeSimplifiedNanoHexapod('actuator_type', 'apa300ml', ...
                          'flex_type_F', '2dof', ...
                          'flex_type_M', '3dof', ...
                          'Fsm', 56e-3, ... % APA300ML weight 112g
                          'Msm', 56e-3, ...
                          'Kf_F', Kf(i), ...
                          'Kf_M', Kf(i));

    G_Kf_apa300ml(i) = {linearize(mdl, io)};
    G_Kf_apa300ml{i}.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
    G_Kf_apa300ml{i}.OutputName = {'l1', 'l2', 'l3', 'l4', 'l5', 'l6', 'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};
end

Kiff = -1000 * ...              % Gain
       1/(s) * ... % LPF: provides integral action
       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 for decentralized IFF - APA300ML actuator - Effect of joint bending stiffness
gains = logspace(-1, 2, 300);

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

for i = 1:length(Kf)
    plot(real(pole(G_Kf_apa300ml{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))),  imag(pole(G_Kf_apa300ml{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))),  'x', 'color', colors(i,:), ...
        'DisplayName', sprintf('$k_f = %.0f$ [Nm/rad]', Kf(i)));
    plot(real(tzero(G_Kf_apa300ml{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))), imag(tzero(G_Kf_apa300ml{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))), 'o', 'color', colors(i,:), ...
        'HandleVisibility', 'off');

    for g = gains
        clpoles = pole(feedback(G_Kf_apa300ml{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}), g*Kiff, +1));
        plot(real(clpoles), imag(clpoles), '.', 'color', colors(i,:), ...
            'HandleVisibility', 'off');
    end

end

xline(0, 'HandleVisibility', 'off'); yline(0, 'HandleVisibility', 'off');
hold off;
axis equal;
xlim(1.4*[-900, 100]); ylim(1.4*[-100, 900]);
xticks(1.4*[-900:100:0]);
yticks(1.4*[0:100:900]);
set(gca, 'XTickLabel',[]); set(gca, 'YTickLabel',[]);
xlabel('Real part'); ylabel('Imaginary part');
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Identify the dynamics for several considered axial stiffnesses
% Let's initialize all the stages with default parameters.
initializeGround('type', 'rigid');
initializeGranite('type', 'rigid');
initializeTy('type', 'rigid');
initializeRy('type', 'rigid');
initializeRz('type', 'rigid');
initializeMicroHexapod('type', 'rigid');
initializeSample('m', 50);

initializeSimscapeConfiguration();
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
io(io_i) = linio([mdl, '/Tracking Error'], 1, 'openoutput', [], 'EdL');  io_i = io_i + 1; % Errors in the frame of the struts
io(io_i) = linio([mdl, '/NASS'],       3, 'openoutput', [], 'fn');  io_i = io_i + 1; % Force Sensors

% Effect of bending stiffness
Ka = 1e6*[1000, 100, 10, 1]; % [Nm/rad]
G_Ka = {zeros(length(Ka), 1)};

for i = 1:length(Ka)
    % Limited joint axial compliance
    initializeSimplifiedNanoHexapod('actuator_type', '1dof', ...
                          'flex_type_F', '2dof_axial', ...
                          'flex_type_M', '4dof', ...
                          'actuator_k', 1e6, ...
                          'actuator_c', 1e1, ...
                          'actuator_kp', 0, ...
                          'actuator_cp', 0, ...
                          'Fsm', 56e-3, ... % APA300ML weight 112g
                          'Msm', 56e-3, ...
                          'Ca_F', 1, ...
                          'Ca_M', 1, ...
                          'Ka_F', Ka(i), ...
                          'Ka_M', Ka(i));

    G_Ka(i) = {linearize(mdl, io)};
    G_Ka{i}.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
    G_Ka{i}.OutputName = {'l1', 'l2', 'l3', 'l4', 'l5', 'l6', 'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};
end

freqs = logspace(1, 4, 1000);

%% Effect of the flexible joint axial stiffness on the HAC-plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:length(Ka)
    for j = 1:5
        for k = j+1:6
            plot(freqs, abs(squeeze(freqresp(G_Ka{i}("l"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(i,:), 0.1], ...
                 'HandleVisibility', 'off');
        end
    end
end
for i = 1:length(Ka)
    plot(freqs, abs(squeeze(freqresp(G_Ka{i}("l1","f1"), freqs, 'Hz'))), 'color', colors(i,:), 'DisplayName', sprintf('$k_a = %.0f$ [N/$\\mu$m]', 1e-6*Ka(i)));
    % for j = 2:6
    %     plot(freqs, abs(squeeze(freqresp(G_Ka{i}("l"+j,"f"+j), freqs, 'Hz'))), 'color', colors(i,:), 'HandleVisibility', 'off');
    % end
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
ylim([1e-10, 1e-4]);

ax2 = nexttile;
hold on;
for i = 1:length(Ka)
    plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_Ka{i}(1, 1), freqs, 'Hz')))));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-200, 20]);
yticks([-360:45:360]);

linkaxes([ax1,ax2],'x');

%% Effect of the flexible joint axial stiffness on the IFF plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:length(Ka)
    for j = 1:5
        for k = j+1:6
            plot(freqs, abs(squeeze(freqresp(G_Ka{i}("fm"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(i,:), 0.1], ...
                 'HandleVisibility', 'off');
        end
    end
end
for i = 1:length(Ka)
    plot(freqs, abs(squeeze(freqresp(G_Ka{i}("fm1","f1"), freqs, 'Hz'))), 'color', colors(i,:), 'DisplayName', sprintf('$k_a = %.0f$ [N/$\\mu$m]', 1e-6*Ka(i)));
    % for j = 2:6
    %     plot(freqs, abs(squeeze(freqresp(G_Ka{i}("fm"+j,"f"+j), freqs, 'Hz'))), 'color', colors(i,:), 'HandleVisibility', 'off');
    % end
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [N/N]'); set(gca, 'XTickLabel',[]);
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
ylim([1e-4, 1e2]);

ax2 = nexttile();
hold on;
for i = 1:length(Ka)
    plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_Ka{i}("fm1", "f1"), freqs, 'Hz')))));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([-360:45:360]);

linkaxes([ax1,ax2],'x');

%% Decentalized IFF
Kiff = -200 * ...              % Gain
       1/s * ... % LPF: provides integral action
       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 for decentralized IFF - 1dof actuator - Effect of joint bending stiffness
gains = logspace(-1, 2, 400);

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

for i = 1:length(Ka)
    plot(real(pole(G_Ka{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))),  imag(pole(G_Ka{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))),  'x', 'color', colors(i,:), ...
        'DisplayName', sprintf('$k_a = %.0f$ N/$\\mu$m', 1e-6*Ka(i)));
    plot(real(tzero(G_Ka{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))), imag(tzero(G_Ka{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}))), 'o', 'color', colors(i,:), ...
        'HandleVisibility', 'off');

    for g = gains
        clpoles = pole(feedback(G_Ka{i}({"fm1", "fm2", "fm3", "fm4", "fm5", "fm6"}, {"f1", "f2", "f3", "f4", "f5", "f6"}), g*Kiff, +1));
        plot(real(clpoles), imag(clpoles), '.', 'color', colors(i,:), ...
            'HandleVisibility', 'off');
    end

end

xline(0, 'HandleVisibility', 'off'); yline(0, 'HandleVisibility', 'off');
hold off;
axis equal;
xlim(1.1*[-900, 100]); ylim(1.1*[-100, 900]);
xticks(1.1*[-900:100:0]);
yticks(1.1*[0:100:900]);
set(gca, 'XTickLabel',[]); set(gca, 'YTickLabel',[]);
xlabel('Real part'); ylabel('Imaginary part');
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Compute the damped plants
Kiff = -500 * ...              % Gain
       1/(s + 2*pi*0.1) * ... % LPF: provides integral action
       eye(6);                 % Diagonal 6x6 controller (i.e. decentralized)

Kiff.InputName = {'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};
Kiff.OutputName = {'u1iff', 'u2iff', 'u3iff', 'u4iff', 'u5iff', 'u6iff'};

% New damped plant input
S1 = sumblk("f1 = u1iff + u1");
S2 = sumblk("f2 = u2iff + u2");
S3 = sumblk("f3 = u3iff + u3");
S4 = sumblk("f4 = u4iff + u4");
S5 = sumblk("f5 = u5iff + u5");
S6 = sumblk("f6 = u6iff + u6");

G_Ka_iff = {zeros(1,length(Ka))};
for i=1:length(Ka)
    G_Ka_iff(i) = {connect(G_Ka{i}, Kiff, S1, S2, S3, S4, S5, S6, {'u1', 'u2', 'u3', 'u4', 'u5', 'u6'}, {'l1', 'l2', 'l3', 'l4', 'l5', 'l6'})};
end

%% Interaction Analysis - RGA Number
rga = zeros(length(Ka), length(freqs));
for i = 1:length(Ka)
    for j = 1:length(freqs)
        rga(i,j) = sum(sum(abs(inv(evalfr(G_Ka_iff{i}({"l1", "l2", "l3", "l4", "l5", "l6"}, {"u1", "u2", "u3", "u4", "u5", "u6"}), 1j*2*pi*freqs(j)).').*evalfr(G_Ka_iff{i}({"l1", "l2", "l3", "l4", "l5", "l6"}, {"u1", "u2", "u3", "u4", "u5", "u6"}), 1j*2*pi*freqs(j)) - eye(6))));
    end
end

%% RGA number for the damped plants - Effect of the flexible joint axial stiffness
figure;
hold on;
for i = 1:length(Ka)
    plot(freqs, rga(i,:), 'DisplayName', sprintf('$k_a = %.0f$ N/$\\mu$m', 1e-6*Ka(i)))
end
hold off;
xlabel('Frequency [Hz]'); ylabel('RGA number');
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylim([0, 10]); xlim([10, 5e3]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Extract stiffness of the joint from the reduced order model
% We first extract the stiffness and mass matrices.
K = readmatrix('flex025_mat_K.CSV');
M = readmatrix('flex025_mat_M.CSV');
% Then, we extract the coordinates of the interface nodes.
[int_xyz, int_i, n_xyz, n_i, nodes] = extractNodes('flex025_out_nodes_3D.txt');

m = 1;

%% Name of the Simulink File
mdl = 'detail_fem_joint';

%% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/T'], 1, 'openinput');  io_i = io_i + 1; % Forces and Torques
io(io_i) = linio([mdl, '/D'], 1, 'openoutput'); io_i = io_i + 1; % Translations and Rotations

G = linearize(mdl, io);

% Stiffness extracted from the Simscape model
k_a = 1/dcgain(G(3,3)); % Axial stiffness [N/m]
k_f = 1/dcgain(G(4,4)); % Bending stiffness [N/m]
k_t = 1/dcgain(G(6,6)); % Torsion stiffness [N/m]


% Stiffness extracted from the Stiffness matrix
k_s = K(1,1); % shear [N/m]
% k_s = K(2,2); % shear [N/m]
k_a = K(3,3); % axial [N/m]
k_f = K(4,4); % bending [Nm/rad]
% k_f = K(5,5); % bending [Nm/rad]
k_t =  K(6,6); % torsion [Nm/rad]

%% Compare Dynamics between "Reduced Order" flexible joints and "2-dof and 3-dof" joints
% Let's initialize all the stages with default parameters.
initializeGround('type', 'rigid');
initializeGranite('type', 'rigid');
initializeTy('type', 'rigid');
initializeRy('type', 'rigid');
initializeRz('type', 'rigid');
initializeMicroHexapod('type', 'rigid');
initializeSample('m', 50);

initializeSimscapeConfiguration();
initializeDisturbances('enable', false);
initializeLoggingConfiguration('log', 'none');
initializeController('type', 'open-loop');
initializeReferences();

mdl = 'detail_fem_nass';

% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/Controller'],     1, 'openinput');              io_i = io_i + 1; % Actuator Inputs
io(io_i) = linio([mdl, '/Tracking Error'], 1, 'openoutput', [], 'EdL');  io_i = io_i + 1; % Errors in the frame of the struts
io(io_i) = linio([mdl, '/NASS'],       3, 'openoutput', [], 'fn');  io_i = io_i + 1; % Force Sensors

% Fully flexible joints
initializeSimplifiedNanoHexapod('actuator_type', 'apa300ml', ...
            'flex_type_F', 'flexible', ...
            'flex_type_M', 'flexible', ...
            'Fsm', 56e-3, ... % APA300ML weight 112g
            'Msm', 56e-3);

G_flex = linearize(mdl, io);
G_flex.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G_flex.OutputName = {'l1', 'l2', 'l3', 'l4', 'l5', 'l6', 'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};

% Flexible joints modelled by 2DoF and 3DoF joints
initializeSimplifiedNanoHexapod('actuator_type', 'apa300ml', ...
            'flex_type_F', '2dof_axial', ...
            'flex_type_M', '4dof', ...
            'Kf_F', k_f, ...
            'Kt_F', k_t, ...
            'Ka_F', k_a, ...
            'Kf_M', k_f, ...
            'Kt_M', k_t, ...
            'Ka_M', k_a, ...
            'Cf_F', 1e-2, ...
            'Ct_F', 1e-2, ...
            'Ca_F', 1e-2, ...
            'Cf_M', 1e-2, ...
            'Ct_M', 1e-2, ...
            'Ca_M', 1e-2, ...
            'Fsm', 56e-3, ... % APA300ML weight 112g
            'Msm', 56e-3);

G_ideal = linearize(mdl, io);
G_ideal.InputName  = {'f1', 'f2', 'f3', 'f4', 'f5', 'f6'};
G_ideal.OutputName = {'l1', 'l2', 'l3', 'l4', 'l5', 'l6', 'fm1', 'fm2', 'fm3', 'fm4', 'fm5', 'fm6'};

%% Comparison of the dynamics with joints modelled with FEM and modelled with "ideal joints" - HAC plant
freqs = logspace(1, 4, 1000);

%% Effect of the flexible joint axial stiffness on the HAC-plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for j = 1:5
    for k = j+1:6
        plot(freqs, abs(squeeze(freqresp(G_flex("l"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(1,:), 0.1], ...
            'HandleVisibility', 'off');
        plot(freqs, abs(squeeze(freqresp(G_ideal("l"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(2,:), 0.1], ...
            'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G_flex("l1","f1"), freqs, 'Hz'))), 'color', colors(1,:), 'DisplayName', 'Reduced Order Flexible Joints');
plot(freqs, abs(squeeze(freqresp(G_ideal("l1","f1"), freqs, 'Hz'))), 'color', colors(2,:), 'DisplayName', 'Bot: $k_f$, $k_a$, Top: $k_f$, $k_t$, $k_a$');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
ylim([1e-10, 1e-4]);

ax2 = nexttile;
hold on;
plot(freqs, 180/pi*angle(squeeze(freqresp(G_flex("l1","f1"), freqs, 'Hz'))));
plot(freqs, 180/pi*angle(squeeze(freqresp(G_ideal("l1","f1"), freqs, 'Hz'))));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([-360:45:360]);

linkaxes([ax1,ax2],'x');

freqs = logspace(0, 3, 1000);

%% Effect of the flexible joint axial stiffness on the HAC-plant
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for j = 1:5
    for k = j+1:6
        plot(freqs, abs(squeeze(freqresp(G_flex("fm"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(1,:), 0.1], ...
            'HandleVisibility', 'off');
        plot(freqs, abs(squeeze(freqresp(G_ideal("fm"+k,"f"+j), freqs, 'Hz'))), 'color', [colors(2,:), 0.1], ...
            'HandleVisibility', 'off');
    end
end
plot(freqs, abs(squeeze(freqresp(G_flex("fm1","f1"), freqs, 'Hz'))), 'color', colors(1,:), 'DisplayName', 'Reduced Order Flexible Joints');
plot(freqs, abs(squeeze(freqresp(G_ideal("fm1","f1"), freqs, 'Hz'))), 'color', colors(2,:), 'DisplayName', 'Bot: $k_f$, $k_a$, Top: $k_f$, $k_t$, $k_a$');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [N/N]'); set(gca, 'XTickLabel',[]);
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
ylim([1e-5, 1e1]);

ax2 = nexttile;
hold on;
plot(freqs, 180/pi*angle(squeeze(freqresp(G_flex("fm1","f1"), freqs, 'Hz'))));
plot(freqs, 180/pi*angle(squeeze(freqresp(G_ideal("fm1","f1"), freqs, 'Hz'))));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-20, 200]);
yticks([-360:45:360]);

linkaxes([ax1,ax2],'x');