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

%% Colors for the figures
colors = colororder;

%% Load frequency response matrix
load('frf_matrix.mat', 'freqs', 'frf');

%% Load Accelerometer positions
acc_pos = readtable('mat/acc_pos.txt', 'ReadVariableNames', false);
acc_pos = table2array(acc_pos(:, 1:4));
[~, i] = sort(acc_pos(:, 1));
acc_pos = acc_pos(i, 2:4);

%% Accelerometers ID connected to each solid body
solids = {};
solids.gbot = [17, 18, 19, 20]; % bottom granite
solids.gtop = [13, 14, 15, 16]; % top granite
solids.ty   = [9, 10, 11, 12]; % Ty stage
solids.ry   = [5, 6, 7, 8]; % Ry stage
solids.rz   = [21, 22, 23]; % Rz stage
solids.hexa = [1, 2, 3, 4]; % Hexapod

% Names of the solid bodies
solid_names = fields(solids);

%% Save the accelerometer positions are well as the solid bodies
save('mat/geometry.mat', 'solids', 'solid_names', 'acc_pos');

%% Extract the CoM of considered solid bodies
model_com = reshape(table2array(readtable('mat/model_solidworks_com.txt', 'ReadVariableNames', false)), [3, 6]);

%% Frequency Response Matrix - Response expressed at the CoM of the solid bodies
frfs_CoM = zeros(length(solid_names)*6, 3, 801);

for solid_i = 1:length(solid_names)
  % Number of accelerometers fixed to this solid body
  solids_i = solids.(solid_names{solid_i});

  % "Jacobian" matrix to go from accelerometer frame to CoM frame
  A = zeros(3*length(solids_i), 6);
  for i = 1:length(solids_i)
    acc_i = solids_i(i);

    acc_pos_com = acc_pos(acc_i, :).' - model_com(:, solid_i);

    A(3*(i-1)+1:3*i, 1:3) = eye(3);
    A(3*(i-1)+1:3*i, 4:6) = [ 0               acc_pos_com(3) -acc_pos_com(2) ;
                             -acc_pos_com(3)  0               acc_pos_com(1) ;
                              acc_pos_com(2) -acc_pos_com(1)  0];
  end

  for exc_dir = 1:3
    frfs_CoM((solid_i-1)*6+1:solid_i*6, exc_dir, :) = A\squeeze(frf((solids_i(1)-1)*3+1:solids_i(end)*3, exc_dir, :));
  end
end

%% Save the computed FRF at the CoM
save('mat/frf_com.mat', 'frfs_CoM');

%% Compute the FRF at the accelerometer location from the CoM reponses
frfs_A = zeros(size(frf));

% For each excitation direction
for exc_dir = 1:3
  % For each solid
  for solid_i = 1:length(solid_names)
    v0 = squeeze(frfs_CoM((solid_i-1)*6+1:(solid_i-1)*6+3, exc_dir, :));
    W0 = squeeze(frfs_CoM((solid_i-1)*6+4:(solid_i-1)*6+6, exc_dir, :));

    % For each accelerometer attached to the current solid
    for acc_i = solids.(solid_names{solid_i})
      % We get the position of the accelerometer expressed in frame O
      pos = acc_pos(acc_i, :).' - model_com(:, solid_i);
      % pos = acc_pos(acc_i, :).';
      posX = [0 pos(3) -pos(2); -pos(3) 0 pos(1) ; pos(2) -pos(1) 0];

      frfs_A(3*(acc_i-1)+1:3*(acc_i-1)+3, exc_dir, :) = v0 + posX*W0;
    end
  end
end

%% Comparison of the original accelerometer response and reconstructed response from the solid body response
exc_names = {'$F_x$', '$F_y$', '$F_z$'};
DOFs = {'x', 'y', 'z', '\theta_x', '\theta_y', '\theta_z'};

solid_i = 6; % Considered solid body
exc_dir = 1; % Excited direction

accs_i = solids.(solid_names{solid_i}); % Accelerometers fixed to this solid body

figure;
tiledlayout(2, 2, 'TileSpacing', 'Tight', 'Padding', 'None');

for i = 1:length(accs_i)
  acc_i = accs_i(i);
  nexttile();

  hold on;
  for dir_i = 1:3
    plot(freqs, abs(squeeze(frf(3*(acc_i-1)+dir_i, exc_dir, :))), '-', 'color', [colors(dir_i,:), 0.5], 'linewidth', 2.5, 'DisplayName', sprintf('$a_{%i,%s}$ - meas', acc_i, DOFs{dir_i}));
  end
  for dir_i = 1:3
    plot(freqs, abs(squeeze(frfs_A(3*(acc_i-1)+dir_i, exc_dir, :))), '-', 'color', colors(dir_i, :), 'DisplayName', sprintf('$a_{%i,%s}$ - solid body', acc_i, DOFs{dir_i}));
  end
  hold off;

  if i > 2
    xlabel('Frequency [Hz]');
  else
    set(gca, 'XTickLabel',[]);
  end

  if rem(i, 2) == 1
    ylabel('Amplitude [$\frac{m/s^2}{N}$]');
  else
    set(gca, 'YTickLabel',[]);
  end

  set(gca, 'XScale', 'lin'); set(gca, 'YScale', 'log');
  xlim([0, 200]); ylim([1e-6, 3e-2]);
  xticks([0:20:200]);
  leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 2);
  leg.ItemTokenSize(1) = 15;
end