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

%% Uniaxial Simscape model name
mdl = 'nass_uniaxial_model';

%% Load the micro-station parameters
load('uniaxial_micro_station_parameters.mat')

%% Load the PSD of disturbances
load('uniaxial_disturbance_psd.mat', 'f', 'psd_ft', 'psd_xf');

%% Load Active Damping Controller
load('uniaxial_active_damping_controllers.mat', 'K_iff_vc', 'K_iff_md', 'K_iff_pz', ...
                                                'K_rdc_vc', 'K_rdc_md', 'K_rdc_pz', ...
                                                'K_dvf_vc', 'K_dvf_md', 'K_dvf_pz');

%% Load High Authority Controllers
load('uniaxial_high_authority_controllers.mat', 'K_hac_vc', 'K_hac_md', 'K_hac_pz');

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

%% Soft Nano-Hexapod
% Light payload mass
mn = 15; % Nano-Hexapod mass [kg]
ms = 1; % Sample Mass [kg]
kn = 1e4; % Nano-Hexapod (soft) Stiffness [N/m]
cn = 2*0.01*sqrt((ms + mn)*kn); % Nano-Hexapod Damping [N/(m/s)]

% Rigid sample
G_vc_rigid_light = 1/((mn + ms)*s^2 + cn*s + kn);

% Soft Sample
ws = 2*pi*20;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_vc_soft_light = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

% Stiff Sample
ws = 2*pi*200;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_vc_stiff_light = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

% Heavy payload mass
mn = 15; % Nano-Hexapod mass [kg]
ms = 50; % Sample Mass [kg]
kn = 1e4; % Nano-Hexapod (soft) Stiffness [N/m]
cn = 2*0.01*sqrt((ms + mn)*kn); % Nano-Hexapod Damping [N/(m/s)]

% Rigid sample
G_vc_rigid_heavy = 1/((mn + ms)*s^2 + cn*s + kn);

% Soft Sample
ws = 2*pi*20;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_vc_soft_heavy = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

% Stiff Sample
ws = 2*pi*200;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_vc_stiff_heavy = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

%% Effect of the payload dynamics on the soft Nano-Hexapod. Light sample on the right, and heavy sample on the left
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G_vc_rigid_light, freqs, 'Hz'))), '-',  'color', colors(1,:), 'DisplayName', 'Rigid sample');
plot(freqs, abs(squeeze(freqresp(G_vc_stiff_light, freqs, 'Hz'))), '-', 'color', colors(2,:), 'DisplayName', '$\omega_s = 200\,Hz$');
plot(freqs, abs(squeeze(freqresp(G_vc_soft_light, freqs, 'Hz'))), '-', 'color', colors(3,:), 'DisplayName', '$\omega_s = 20\,Hz$');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-10, 1e-2])

ax1b = nexttile();
hold on;
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_vc_rigid_light, freqs, 'Hz')))), '-',  'color', colors(1,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_vc_stiff_light, freqs, 'Hz')))), '-', 'color', colors(2,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_vc_soft_light, freqs, 'Hz')))), '-', 'color', colors(3,:));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
xticks([1e0, 1e1, 1e2]);
yticks(-360:90:360);
ylim([-200, 20]);

linkaxes([ax1,ax1b],'x');
xlim([1, 1000]);

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

ax2 = nexttile([2,1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G_vc_rigid_heavy, freqs, 'Hz'))), '-',  'color', colors(1,:), 'DisplayName', 'Rigid sample');
plot(freqs, abs(squeeze(freqresp(G_vc_stiff_heavy, freqs, 'Hz'))), '-', 'color', colors(2,:), 'DisplayName', '$\omega_s = 200\,Hz$');
plot(freqs, abs(squeeze(freqresp(G_vc_soft_heavy, freqs, 'Hz'))), '-', 'color', colors(3,:), 'DisplayName', '$\omega_s = 20\,Hz$');
plot(freqs, abs(squeeze(freqresp(1/(mn*s^2), freqs, 'Hz'))), '-', 'color', [0,0,0,0.5], 'DisplayName', '$\frac{1}{m_n s^2}$');
plot(freqs, abs(squeeze(freqresp(1/((mn + ms)*s^2), freqs, 'Hz'))), '--', 'color', [0,0,0,0.5], 'DisplayName', '$\frac{1}{(m_n + m_s) s^2}$');
text(2.2, 2e-3, '$\omega_n = \sqrt{\frac{k_n}{m_n + m_s}}$', 'horizontalalignment', 'left');
text(20, 1e-8, '$\omega_s = \sqrt{\frac{k_s}{m_s}}$', 'horizontalalignment', 'center');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
ldg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
ldg.ItemTokenSize = [20, 1];
ylim([1e-10, 1e-2])

ax2b = nexttile();
hold on;
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_vc_rigid_heavy, freqs, 'Hz')))), '-',  'color', colors(1,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_vc_stiff_heavy, freqs, 'Hz')))), '-', 'color', colors(2,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_vc_soft_heavy, freqs, 'Hz')))), '-', 'color', colors(3,:));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
xticks([1e0, 1e1, 1e2]);
yticks(-360:90:360);
ylim([-200, 20]);

linkaxes([ax2,ax2b],'x');
xlim([1, 1000]);

%% Stiff Nano-Hexapod
% Light payload mass
mn = 15; % Nano-Hexapod mass [kg]
ms = 1; % Sample Mass [kg]
kn = 1e8; % Nano-Hexapod (soft) Stiffness [N/m]
cn = 2*0.01*sqrt((ms + mn)*kn); % Nano-Hexapod Damping [N/(m/s)]

% Rigid sample
G_pz_rigid_light = 1/((mn + ms)*s^2 + cn*s + kn);

% Soft Sample
ws = 2*pi*20;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_pz_soft_light = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

% Stiff Sample
ws = 2*pi*200;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_pz_stiff_light = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

% Heavy payload mass
mn = 15; % Nano-Hexapod mass [kg]
ms = 50; % Sample Mass [kg]
kn = 1e8; % Nano-Hexapod (soft) Stiffness [N/m]
cn = 2*0.01*sqrt((ms + mn)*kn); % Nano-Hexapod Damping [N/(m/s)]

% Rigid sample
G_pz_rigid_heavy = 1/((mn + ms)*s^2 + cn*s + kn);

% Soft Sample
ws = 2*pi*20;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_pz_soft_heavy = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

% Stiff Sample
ws = 2*pi*200;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);
G_pz_stiff_heavy = (ms*s^2 + cs*s + ks)/((mn*s^2 + cn*s + kn)*(ms*s^2 + cs*s + ks) + ms*s^2*(cs*s + ks));

%% Effect of the payload dynamics on the stiff Nano-Hexapod. Light sample on the right, and heavy sample on the left
figure;
tiledlayout(3, 1, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G_pz_rigid_light, freqs, 'Hz'))), '-',  'color', colors(1,:));
plot(freqs, abs(squeeze(freqresp(G_pz_stiff_light, freqs, 'Hz'))), '-', 'color', colors(2,:));
plot(freqs, abs(squeeze(freqresp(G_pz_soft_light, freqs, 'Hz'))), '-', 'color', colors(3,:));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-10, 1e-6])

ax1b = nexttile();
hold on;
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_pz_rigid_light, freqs, 'Hz')))), '-',  'color', colors(1,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_pz_stiff_light, freqs, 'Hz')))), '-', 'color', colors(2,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_pz_soft_light, freqs, 'Hz')))), '-', 'color', colors(3,:));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
xticks([1e0, 1e1, 1e2]);
yticks(-360:90:360);
ylim([-200, 20]);

linkaxes([ax1,ax1b],'x');
xlim([1, 1000]);

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

ax2 = nexttile([2,1]);
hold on;
plot(freqs, abs(squeeze(freqresp(G_pz_rigid_heavy, freqs, 'Hz'))), '-',  'color', colors(1,:), 'DisplayName', 'Rigid sample');
plot(freqs, abs(squeeze(freqresp(G_pz_stiff_heavy, freqs, 'Hz'))), '-', 'color', colors(2,:), 'DisplayName', 'Stiff sample: $\omega_s = 200\,Hz$');
plot(freqs, abs(squeeze(freqresp(G_pz_soft_heavy, freqs, 'Hz'))), '-', 'color', colors(3,:), 'DisplayName', 'Soft sample: $\omega_s = 20\,Hz$');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);
ylim([1e-10, 1e-6])
ldg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
ldg.ItemTokenSize = [20, 1];

ax2b = nexttile();
hold on;
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_pz_rigid_heavy, freqs, 'Hz')))), '-',  'color', colors(1,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_pz_stiff_heavy, freqs, 'Hz')))), '-', 'color', colors(2,:));
plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_pz_soft_heavy, freqs, 'Hz')))), '-', 'color', colors(3,:));
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
xticks([1e0, 1e1, 1e2]);
yticks(-360:90:360);
ylim([-200, 20]);

linkaxes([ax2,ax2b],'x');
xlim([1, 1000]);

%% Nano-Hexpod model
model_config = struct();
model_config.controller = "open_loop";
mn = 15; % Nano-Hexapod mass [kg]
ms = 1; % Sample Mass [kg]

%% Identification
clear io; io_i = 1;
io(io_i) = linio([mdl, '/controller'],       1, 'openinput');  io_i = io_i + 1; % Actuator Force
io(io_i) = linio([mdl, '/micro_station/xf'], 1, 'openinput');  io_i = io_i + 1; % Floor Motion
io(io_i) = linio([mdl, '/micro_station/ft'], 1, 'openinput');  io_i = io_i + 1; % Stage vibrations
io(io_i) = linio([mdl, '/fs'],               1, 'openinput');  io_i = io_i + 1; % Direct sample forces
io(io_i) = linio([mdl, '/dL'],               1, 'openoutput'); io_i = io_i + 1; % Relative Motion Sensor
io(io_i) = linio([mdl, '/fm'],               1, 'openoutput'); io_i = io_i + 1; % Force Sensor
io(io_i) = linio([mdl, '/vn'] ,              1, 'openoutput'); io_i = io_i + 1; % Geophone
io(io_i) = linio([mdl, '/d'] ,               1, 'openoutput'); io_i = io_i + 1; % Metrology Output
io(io_i) = linio([mdl, '/y'] ,               1, 'openoutput'); io_i = io_i + 1; % Sample's position

%% Soft Nano-Hexapod
% Light payload mass
kn = 1e4; % Nano-Hexapod (soft) Stiffness [N/m]
cn = 2*0.01*sqrt((ms + mn)*kn); % Nano-Hexapod Damping [N/(m/s)]

% Rigid Sample
model_config.nhexa = "1dof";
G_vc_light_rigid = linearize(mdl, io, 0.0);
G_vc_light_rigid.InputName = {'f', 'xf', 'ft', 'fs'};
G_vc_light_rigid.OutputName = {'dL', 'fm', 'vn', 'd', 'y'};

% Soft Sample
model_config.nhexa = "2dof";
ws = 2*pi*20;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);

G_vc_light_soft = linearize(mdl, io, 0.0);
G_vc_light_soft.InputName = {'f', 'xf', 'ft', 'fs'};
G_vc_light_soft.OutputName = {'dL', 'fm', 'vn', 'd', 'y'};

% Rigid Sample
model_config.nhexa = "2dof";
ws = 2*pi*200;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);

G_vc_light_stiff = linearize(mdl, io, 0.0);
G_vc_light_stiff.InputName = {'f', 'xf', 'ft', 'fs'};
G_vc_light_stiff.OutputName = {'dL', 'fm', 'vn', 'd', 'y'};

%% Stiff Nano-Hexapod
% Light payload mass
kn = 1e8; % Nano-Hexapod (soft) Stiffness [N/m]
cn = 2*0.01*sqrt((ms + mn)*kn); % Nano-Hexapod Damping [N/(m/s)]

% Rigid Sample
model_config.nhexa = "1dof";
G_pz_light_rigid = linearize(mdl, io, 0.0);
G_pz_light_rigid.InputName = {'f', 'xf', 'ft', 'fs'};
G_pz_light_rigid.OutputName = {'dL', 'fm', 'vn', 'd', 'y'};

% Soft Sample
model_config.nhexa = "2dof";
ws = 2*pi*20;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);

G_pz_light_soft = linearize(mdl, io, 0.0);
G_pz_light_soft.InputName = {'f', 'xf', 'ft', 'fs'};
G_pz_light_soft.OutputName = {'dL', 'fm', 'vn', 'd', 'y'};

% Rigid Sample
model_config.nhexa = "2dof";
ws = 2*pi*200;
ks = ms * ws^2;
cs = 2*0.01*sqrt(ms*ks);

G_pz_light_stiff = linearize(mdl, io, 0.0);
G_pz_light_stiff.InputName = {'f', 'xf', 'ft', 'fs'};
G_pz_light_stiff.OutputName = {'dL', 'fm', 'vn', 'd', 'y'};

%% Apply IFF and verify stability
% Soft Nano-Hexapod
G_iff_vc_light_rigid = feedback(G_vc_light_rigid, K_iff_vc, 'name', +1);
G_iff_vc_light_soft  = feedback(G_vc_light_soft , K_iff_vc, 'name', +1);
G_iff_vc_light_stiff = feedback(G_vc_light_stiff, K_iff_vc, 'name', +1);

if not(isstable(G_iff_vc_light_rigid) && isstable(G_iff_vc_light_soft) &&isstable(G_iff_vc_light_stiff))
    warning("One of the damped plant with VC and decentralized IFF is not stable.");
end

% Stiff Nano-Hexapod
G_iff_pz_light_rigid = feedback(G_pz_light_rigid, K_iff_pz, 'name', +1);
G_iff_pz_light_soft  = feedback(G_pz_light_soft , K_iff_pz, 'name', +1);
G_iff_pz_light_stiff = feedback(G_pz_light_stiff, K_iff_pz, 'name', +1);

if not(isstable(G_iff_pz_light_rigid) && isstable(G_iff_pz_light_soft) && isstable(G_iff_pz_light_stiff))
    warning("One of the damped plant with PZ and decentralized IFF is not stable.");
end

%% Compute closed-loop plants and verify stability
% Soft Nano-Hexapod
G_hac_iff_vc_light_rigid = feedback(G_iff_vc_light_rigid, K_hac_vc, 'name', -1);
G_hac_iff_vc_light_soft  = feedback(G_iff_vc_light_soft , K_hac_vc, 'name', -1);
G_hac_iff_vc_light_stiff = feedback(G_iff_vc_light_stiff, K_hac_vc, 'name', -1);

if not(isstable(G_hac_iff_vc_light_rigid) && isstable(G_hac_iff_vc_light_soft) && isstable(G_hac_iff_vc_light_stiff))
    warning("One of the damped plant with VC and decentralized IFF is not stable.");
end

% Stiff Nano-Hexapod
G_hac_iff_pz_light_rigid = feedback(G_iff_pz_light_rigid, K_hac_pz, 'name', -1);
G_hac_iff_pz_light_soft  = feedback(G_iff_pz_light_soft , K_hac_pz, 'name', -1);
G_hac_iff_pz_light_stiff = feedback(G_iff_pz_light_stiff, K_hac_pz, 'name', -1);

if not(isstable(G_hac_iff_pz_light_rigid) && isstable(G_hac_iff_pz_light_soft) && isstable(G_hac_iff_pz_light_stiff))
    warning("One of the damped plant with PZ and decentralized IFF is not stable.");
end

%% Cumulative Amplitude Spectrum of d - Effect of Sample's flexibility
figure;
hold on;
plot(f, sqrt(flip(-cumtrapz(flip(f), flip(psd_ft.*abs(squeeze(freqresp(G_hac_iff_pz_light_rigid('d', 'ft'), f, 'Hz'))).^2 + ...
                                     psd_xf.*abs(squeeze(freqresp(G_hac_iff_pz_light_rigid('d', 'xf'), f, 'Hz'))).^2)))), '-', ...
    'DisplayName', 'Rigid sample');
plot(f, sqrt(flip(-cumtrapz(flip(f), flip(psd_ft.*abs(squeeze(freqresp(G_hac_iff_pz_light_stiff('d', 'ft'), f, 'Hz'))).^2 + ...
                                     psd_xf.*abs(squeeze(freqresp(G_hac_iff_pz_light_stiff('d', 'xf'), f, 'Hz'))).^2)))), '-', ...
    'DisplayName', '$\omega_s = 200\,$Hz');
plot(f, sqrt(flip(-cumtrapz(flip(f), flip(psd_ft.*abs(squeeze(freqresp(G_hac_iff_pz_light_soft('d', 'ft'), f, 'Hz'))).^2 + ...
                                     psd_xf.*abs(squeeze(freqresp(G_hac_iff_pz_light_soft('d', 'xf'), f, 'Hz'))).^2)))), '-', ...
    'DisplayName', '$\omega_s = 20\,$Hz');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('CAS of $d$ [m]'); xlabel('Frequency [Hz]');
legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
xlim([1, 500]);
xticks([1e0, 1e1, 1e2]);
ylim([2e-10, 2e-7])

%% Cumulative Amplitude Spectrum - Effect of Sample's flexibility
figure;
hold on;
plot(f, sqrt(flip(-cumtrapz(flip(f), flip(psd_ft.*abs(squeeze(freqresp(G_hac_iff_pz_light_rigid('y', 'ft'), f, 'Hz'))).^2 + ...
                                     psd_xf.*abs(squeeze(freqresp(G_hac_iff_pz_light_rigid('y', 'xf'), f, 'Hz'))).^2)))), '-', ...
    'DisplayName', 'Rigid sample');
plot(f, sqrt(flip(-cumtrapz(flip(f), flip(psd_ft.*abs(squeeze(freqresp(G_hac_iff_pz_light_stiff('y', 'ft'), f, 'Hz'))).^2 + ...
                                     psd_xf.*abs(squeeze(freqresp(G_hac_iff_pz_light_stiff('y', 'xf'), f, 'Hz'))).^2)))), '-', ...
    'DisplayName', '$\omega_s = 200\,$Hz');
plot(f, sqrt(flip(-cumtrapz(flip(f), flip(psd_ft.*abs(squeeze(freqresp(G_hac_iff_pz_light_soft('y', 'ft'), f, 'Hz'))).^2 + ...
                                     psd_xf.*abs(squeeze(freqresp(G_hac_iff_pz_light_soft('y', 'xf'), f, 'Hz'))).^2)))), '-', ...
    'DisplayName', '$\omega_s = 20\,$Hz');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('CAS of $y$ [m]'); xlabel('Frequency [Hz]');
legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
xlim([1, 500]);
xticks([1e0, 1e1, 1e2]);
ylim([2e-10, 2e-7])