%% ustation_2_modeling.m

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

% Simulink Model name
mdl = 'ustation_simscape';

load('nass_model_conf_simulink.mat');

%% Colors for the figures
colors = colororder;

%% Frequency Vector
freqs = logspace(log10(10), log10(2e3), 1000);

%% Indentify the model dynamics from the 3 hammer imapcts
% To the motion of each solid body at their CoM

% All stages are initialized
initializeGround(      'type', 'rigid');
initializeGranite(     'type', 'flexible');
initializeTy(          'type', 'flexible');
initializeRy(          'type', 'flexible');
initializeRz(          'type', 'flexible');
initializeMicroHexapod('type', 'flexible');

initializeLoggingConfiguration('log', 'none');

initializeReferences();
initializeDisturbances('enable', false);

% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/Micro-Station/Translation Stage/Flexible/F_hammer'],          1, 'openinput');  io_i = io_i + 1;
io(io_i) = linio([mdl, '/Micro-Station/Granite/Flexible/accelerometer'],               1, 'openoutput'); io_i = io_i + 1;
io(io_i) = linio([mdl, '/Micro-Station/Translation Stage/Flexible/accelerometer'],     1, 'openoutput'); io_i = io_i + 1;
io(io_i) = linio([mdl, '/Micro-Station/Tilt Stage/Flexible/accelerometer'],            1, 'openoutput'); io_i = io_i + 1;
io(io_i) = linio([mdl, '/Micro-Station/Spindle/Flexible/accelerometer'],               1, 'openoutput'); io_i = io_i + 1;
io(io_i) = linio([mdl, '/Micro-Station/Micro Hexapod/Flexible/accelerometer'],         1, 'openoutput'); io_i = io_i + 1;

% Run the linearization
G_ms = linearize(mdl, io, 0);
G_ms.InputName  = {'Fx', 'Fy', 'Fz'};
G_ms.OutputName = {'gtop_x', 'gtop_y', 'gtop_z', 'gtop_rx', 'gtop_ry', 'gtop_rz', ...
                   'ty_x', 'ty_y', 'ty_z', 'ty_rx', 'ty_ry', 'ty_rz', ...
                   'ry_x', 'ry_y', 'ry_z', 'ry_rx', 'ry_ry', 'ry_rz', ...
                   'rz_x', 'rz_y', 'rz_z', 'rz_rx', 'rz_ry', 'rz_rz', ...
                   'hexa_x', 'hexa_y', 'hexa_z', 'hexa_rx', 'hexa_ry', 'hexa_rz'};

% Load estimated FRF at CoM
load('ustation_frf_matrix.mat', 'freqs');
load('ustation_frf_com.mat', 'frfs_CoM');

% Initialization of some variables to the figures
dirs = {'x', 'y', 'z', 'rx', 'ry', 'rz'};
stages = {'gbot', 'gtop', 'ty', 'ry', 'rz', 'hexa'};
f = logspace(log10(10), log10(500), 1000);

%% Spindle - X response
n_stg = 5; % Measured Stage
n_dir = 1; % Measured DoF: x, y, z, Rx, Ry, Rz
n_exc = 1; % Hammer impact: x, y, z

figure;
hold on;
plot(freqs, abs(squeeze(frfs_CoM(6*(n_stg-1) + n_dir, n_exc, :))), 'DisplayName', 'Measurement');
plot(f, abs(squeeze(freqresp(G_ms([stages{n_stg}, '_', dirs{n_dir}], ['F', dirs{n_exc}]), f, 'Hz'))), 'DisplayName', 'Model');
text(50, 2e-2,{'$D_x/F_x$ - Spindle'},'VerticalAlignment','bottom','HorizontalAlignment','center')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [Hz]'); ylabel('Magnitude [$m/s^2/N$]');
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
xlim([10, 200]);
ylim([1e-6, 1e-1])

%% Hexapod - Y response
n_stg = 6; % Measured Stage
n_dir = 2; % Measured DoF: x, y, z, Rx, Ry, Rz
n_exc = 2; % Hammer impact: x, y, z

figure;
hold on;
plot(freqs, abs(squeeze(frfs_CoM(6*(n_stg-1) + n_dir, n_exc, :))), 'DisplayName', 'Measurement');
plot(f, abs(squeeze(freqresp(G_ms([stages{n_stg}, '_', dirs{n_dir}], ['F', dirs{n_exc}]), f, 'Hz'))), 'DisplayName', 'Model');
text(50, 2e-2,{'$D_y/F_y$ - Hexapod'},'VerticalAlignment','bottom','HorizontalAlignment','center')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [Hz]'); ylabel('Magnitude [$m/s^2/N$]');
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
xlim([10, 200]);
ylim([1e-6, 1e-1])

%% Tilt stage - Z response
n_stg = 4; % Measured Stage
n_dir = 3; % Measured DoF: x, y, z, Rx, Ry, Rz
n_exc = 3; % Hammer impact: x, y, z

figure;
hold on;
plot(freqs, abs(squeeze(frfs_CoM(6*(n_stg-1) + n_dir, n_exc, :))), 'DisplayName', 'Measurement');
plot(f, abs(squeeze(freqresp(G_ms([stages{n_stg}, '_', dirs{n_dir}], ['F', dirs{n_exc}]), f, 'Hz'))), 'DisplayName', 'Model');
text(50, 2e-2,{'$D_z/F_z$ - Tilt'},'VerticalAlignment','bottom','HorizontalAlignment','center')
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [Hz]'); ylabel('Magnitude [$m/s^2/N$]');
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;
xlim([10, 200]);
ylim([1e-6, 1e-1])

% Positions and orientation of accelerometers
% | *Num* | *Position* | *Orientation* | *Sensibility* | *Channels* |
% |-------+------------+---------------+---------------+------------|
% |     1 | [0, +A, 0] | [x, y, z]     | 1V/g          |        1-3 |
% |     2 | [-B, 0, 0] | [x, y, z]     | 1V/g          |        4-6 |
% |     3 | [0, -A, 0] | [x, y, z]     | 0.1V/g        |        7-9 |
% |     4 | [+B, 0, 0] | [x, y, z]     | 1V/g          |      10-12 |
%

% Measuerment channels
%
% | Acc Number | Dir | Channel Number |
% |------------+-----+----------------|
% |          1 | x   |              1 |
% |          1 | y   |              2 |
% |          1 | z   |              3 |
% |          2 | x   |              4 |
% |          2 | y   |              5 |
% |          2 | z   |              6 |
% |          3 | x   |              7 |
% |          3 | y   |              8 |
% |          3 | z   |              9 |
% |          4 | x   |             10 |
% |          4 | y   |             11 |
% |          4 | z   |             12 |
% |     Hammer |     |             13 |

% 10 measurements corresponding to 10 different Hammer impact locations
% | *Num* | *Direction* | *Position* | Accelerometer position | Jacobian number |
% |-------+-------------+------------+------------------------+-----------------|
% |     1 | -Y          | [0, +A, 0] |                      1 |              -2 |
% |     2 | -Z          | [0, +A, 0] |                      1 |              -3 |
% |     3 | X           | [-B, 0, 0] |                      2 |               4 |
% |     4 | -Z          | [-B, 0, 0] |                      2 |              -6 |
% |     5 | Y           | [0, -A, 0] |                      3 |               8 |
% |     6 | -Z          | [0, -A, 0] |                      3 |              -9 |
% |     7 | -X          | [+B, 0, 0] |                      4 |             -10 |
% |     8 | -Z          | [+B, 0, 0] |                      4 |             -12 |
% |     9 | -X          | [0, -A, 0] |                      3 |              -7 |
% |    10 | -X          | [0, +A, 0] |                      1 |              -1 |


%% Jacobian for accelerometers
% aX = Ja . aL
d = 0.14; % [m]
Ja = [1  0  0  0  0 -d;
      0  1  0  0  0  0;
      0  0  1  d  0  0;
      1  0  0  0  0  0;
      0  1  0  0  0 -d;
      0  0  1  0  d  0;
      1  0  0  0  0  d;
      0  1  0  0  0  0;
      0  0  1 -d  0  0;
      1  0  0  0  0  0;
      0  1  0  0  0  d;
      0  0  1  0 -d  0];

Ja_inv = pinv(Ja);

%% Jacobian for hammer impacts
% F = Jf' tau
Jf = [0 -1  0  0  0  0;
      0  0 -1 -d  0  0;
      1  0  0  0  0  0;
      0  0 -1  0 -d  0;
      0  1  0  0  0  0;
      0  0 -1  d  0  0;
     -1  0  0  0  0  0;
      0  0 -1  0  d  0;
     -1  0  0  0  0 -d;
     -1  0  0  0  0  d]';
Jf_inv = pinv(Jf);

%% Load raw measurement data
% "Track1" to "Track12" are the 12 accelerometers
% "Track13" is the instrumented hammer force sensor
data = [
    load('ustation_compliance_hammer_1.mat'), ...
    load('ustation_compliance_hammer_2.mat'), ...
    load('ustation_compliance_hammer_3.mat'), ...
    load('ustation_compliance_hammer_4.mat'), ...
    load('ustation_compliance_hammer_5.mat'), ...
    load('ustation_compliance_hammer_6.mat'), ...
    load('ustation_compliance_hammer_7.mat'), ...
    load('ustation_compliance_hammer_8.mat'), ...
    load('ustation_compliance_hammer_9.mat'), ...
    load('ustation_compliance_hammer_10.mat')];

%% Prepare the FRF computation
Ts = 1e-3; % Sampling Time [s]
Nfft = floor(1/Ts); % Number of points for the FFT computation
win = ones(Nfft, 1); % Rectangular window
[~, f] = tfestimate(data(1).Track13, data(1).Track1, win, [], Nfft, 1/Ts);

% Samples taken before and after the hammer "impact"
pre_n = floor(0.1/Ts);
post_n = Nfft - pre_n - 1;

%% Compute the FRFs for the 10 hammer impact locations.
% The FRF from hammer force the 12 accelerometers are computed
% for each of the 10 hammer impacts and averaged
G_raw = zeros(12,10,length(f));

for i = 1:10 % For each impact location
    % Find the 10 impacts
    [~, impacts_i] = find(diff(data(i).Track13 > 50)==1);
    % Only keep the first 10 impacts if there are more than 10 impacts
    if length(impacts_i)>10
        impacts_i(11:end) = [];
    end

    % Average the FRF for the 10 impacts
    for impact_i = impacts_i
        i_start = impact_i - pre_n;
        i_end = impact_i + post_n;
        data_in = data(i).Track13(i_start:i_end); % [N]
        % Remove hammer DC offset
        data_in  = data_in  - mean(data_in(end-pre_n:end));
        % Combine all outputs [m/s^2]
        data_out = [data(i).Track1( i_start:i_end); ...
                    data(i).Track2( i_start:i_end); ...
                    data(i).Track3( i_start:i_end); ...
                    data(i).Track4( i_start:i_end); ...
                    data(i).Track5( i_start:i_end); ...
                    data(i).Track6( i_start:i_end); ...
                    data(i).Track7( i_start:i_end); ...
                    data(i).Track8( i_start:i_end); ...
                    data(i).Track9( i_start:i_end); ...
                    data(i).Track10(i_start:i_end); ...
                    data(i).Track11(i_start:i_end); ...
                    data(i).Track12(i_start:i_end)];

        [frf, ~] = tfestimate(data_in, data_out', win, [], Nfft, 1/Ts);
        G_raw(:,i,:) = squeeze(G_raw(:,i,:)) + 0.1*frf'./(-(2*pi*f').^2);
    end
end

%% Compute transfer function in cartesian frame using Jacobian matrices
% FRF_cartesian = inv(Ja) * FRF * inv(Jf)
FRF_cartesian = pagemtimes(Ja_inv, pagemtimes(G_raw, Jf_inv));

%% Identification of the compliance of the micro-station

% Initialize simulation with default parameters (flexible elements)
initializeGround();
initializeGranite();
initializeTy();
initializeRy();
initializeRz();
initializeMicroHexapod();

initializeReferences();
initializeSimscapeConfiguration();
initializeLoggingConfiguration();

% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/Micro-Station/Micro Hexapod/Flexible/Fm'], 1, 'openinput');       io_i = io_i + 1; % Direct Forces/Torques applied on the micro-hexapod top platform
io(io_i) = linio([mdl, '/Micro-Station/Micro Hexapod/Flexible/Dm'], 1, 'output'); io_i = io_i + 1; % Absolute displacement of the top platform

% Run the linearization
Gm = linearize(mdl, io, 0);
Gm.InputName  = {'Fmx', 'Fmy', 'Fmz', 'Mmx', 'Mmy', 'Mmz'};
Gm.OutputName = {'Dx', 'Dy', 'Dz', 'Drx', 'Dry', 'Drz'};

%% Extracted FRF of the compliance of the micro-station in the Cartesian frame from the Simscape model
figure;
hold on;
plot(f, abs(squeeze(FRF_cartesian(1,1,:))), '-', 'color', [colors(1,:), 0.5], 'linewidth', 2.5, 'DisplayName', '$D_x/F_x$ - Measured')
plot(f, abs(squeeze(FRF_cartesian(2,2,:))), '-', 'color', [colors(2,:), 0.5], 'linewidth', 2.5, 'DisplayName', '$D_y/F_y$ - Measured')
plot(f, abs(squeeze(FRF_cartesian(3,3,:))), '-', 'color', [colors(3,:), 0.5], 'linewidth', 2.5, 'DisplayName', '$D_z/F_z$ - Measured')
plot(f, abs(squeeze(freqresp(Gm(1,1), f, 'Hz'))), '-', 'color', colors(1,:), 'DisplayName', '$D_x/F_x$ - Model')
plot(f, abs(squeeze(freqresp(Gm(2,2), f, 'Hz'))), '-', 'color', colors(2,:), 'DisplayName', '$D_y/F_y$ - Model')
plot(f, abs(squeeze(freqresp(Gm(3,3), f, 'Hz'))), '-', 'color', colors(3,:), 'DisplayName', '$D_z/F_z$ - Model')
hold off;
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 2);
leg.ItemTokenSize(1) = 15;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [Hz]'); ylabel('Magnitude [m/N]');
xlim([20, 500]); ylim([2e-10, 1e-6]);
xticks([20, 50, 100, 200, 500])

%% Extracted FRF of the compliance of the micro-station in the Cartesian frame from the Simscape model
figure;
hold on;
plot(f, abs(squeeze(FRF_cartesian(4,4,:))), '-', 'color', [colors(1,:), 0.5], 'linewidth', 2.5, 'DisplayName', '$R_x/M_x$ - Measured')
plot(f, abs(squeeze(FRF_cartesian(5,5,:))), '-', 'color', [colors(2,:), 0.5], 'linewidth', 2.5, 'DisplayName', '$R_y/M_y$ - Measured')
plot(f, abs(squeeze(FRF_cartesian(6,6,:))), '-', 'color', [colors(3,:), 0.5], 'linewidth', 2.5, 'DisplayName', '$R_z/M_z$ - Measured')
plot(f, abs(squeeze(freqresp(Gm(4,4), f, 'Hz'))), '-', 'color', colors(1,:), 'DisplayName', '$R_x/M_x$ - Model')
plot(f, abs(squeeze(freqresp(Gm(5,5), f, 'Hz'))), '-', 'color', colors(2,:), 'DisplayName', '$R_y/M_y$ - Model')
plot(f, abs(squeeze(freqresp(Gm(6,6), f, 'Hz'))), '-', 'color', colors(3,:), 'DisplayName', '$R_z/M_z$ - Model')
hold off;
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 2);
leg.ItemTokenSize(1) = 15;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
xlabel('Frequency [Hz]'); ylabel('Magnitude [rad/Nm]');
xlim([20, 500]); ylim([2e-7, 1e-4]);
xticks([20, 50, 100, 200, 500])
