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

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

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

%% Compute Floor Motion Spectral Density
% Load floor motion data
% t: time in [s]
% V: measured voltage genrated by the geophone and amplified by a 60dB gain voltage amplifier [V]
load('ground_motion_measurement.mat', 't', 'V');

% Geophone Transfer Function
Tg = 88; % Sensitivity [V/(m/s)]
w0 = 2*2*pi; % Cut-off frequency [rad/s]
xi = 0.7; % Damping ratio

G_geo = Tg*s*s^2/(s^2 + 2*xi*w0*s + w0^2); % Geophone's transfer function [V/m]

% Voltage amplifier transfer function
g0 = 10^(60/20); % [abs]

% Compute measured voltage PSD
Ts = (t(2)-t(1)); % Sampling Time [s]
Nfft = floor(2/Ts);
win = hanning(Nfft);
Noverlap = floor(Nfft/2);

[psd_V, f] = pwelch(V, win, Noverlap, Nfft, 1/Ts); % [V^2/Hz]

% Ground Motion ASD
psd_xf = psd_V./abs(squeeze(freqresp(G_geo*g0, f, 'Hz'))).^2; % [m^2/Hz]

%% Amplitude Spectral Density of the measured Floor motion on ID31
figure;
plot(f, sqrt(psd_xf), 'DisplayName', '$\Gamma_{x_{f}}$');
set(gca, 'xscale', 'log'); set(gca, 'yscale', 'log');
xlabel('Frequency [Hz]'); ylabel('Ampl. Spectral Density $\left[\frac{m}{\sqrt{Hz}}\right]$')
legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
xlim([1, 500]);
xticks([1e0, 1e1, 1e2]);

%% Estimation of the Spectral density of the stage vibrations

% Measured velocity of granite and hexapod during spindle rotation
% t: time in [s]
% vg: measured granite velocity [m/s]
% vg: measured micro-hexapod's top platform velocity [m/s]
load('meas_spindle_on.mat', 't', 'vg', 'vh');
spindle_off = load('meas_spindle_off.mat', 't', 'vg', 'vh'); % No Rotation

% Compute Power Spectral Density of the relative velocity between granite and hexapod during spindle rotation
Ts = t(2)-t(1); % Sampling Time [s]
Nfft = floor(2/Ts);
win = hanning(Nfft);

[psd_vft, f] = pwelch(vh-vg, win, Noverlap, Nfft, 1/Ts); % [(m/s)^2/Hz]
[psd_off, ~] = pwelch(spindle_off.vh-spindle_off.vg, win, Noverlap, Nfft, 1/Ts); % [(m/s)^2/Hz]

% Disable the Nano-Hexpod for now
model_config = struct();
model_config.nhexa = "none";
model_config.controller = "open_loop";

% Identify the transfer function from u to taum
clear io; io_i = 1;
io(io_i) = linio([mdl, '/micro_station/ft'], 1, 'openinput');  io_i = io_i + 1; % Stage Disturbance Force
io(io_i) = linio([mdl, '/micro_station/xg'], 1, 'openoutput'); io_i = io_i + 1; % Absolute motion of Granite
io(io_i) = linio([mdl, '/micro_station/xh'], 1, 'openoutput'); io_i = io_i + 1; % Absolute motion of Hexapod

% Perform the model extraction
G = linearize(mdl, io, 0.0);
G.InputName = {'ft'};
G.OutputName = {'Dg', 'Dh'};

% Power Spectral Density of the equivalent force ft [N/Hz^2]
psd_ft = (psd_vft./(2*pi*f).^2)./abs(squeeze(freqresp(G('Dh', 'ft') - G('Dg', 'ft'), f, 'Hz'))).^2;

%% Amplitude Spectral Density of the relative motion measured between the granite and the micro-hexapod's top platform during Spindle rotating
figure;
hold on;
plot(f, sqrt(psd_vft)./(2*pi*f), 'DisplayName', '6rpm');
plot(f, sqrt(psd_off)./(2*pi*f), 'DisplayName', '0rpm');
hold off;
set(gca, 'xscale', 'log'); set(gca, 'yscale', 'log');
xlabel('Frequency [Hz]'); ylabel('Ampl. Spectral Density $\left[\frac{m}{\sqrt{Hz}}\right]$')
legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
xlim([1, 500]); ylim([1e-12, 1e-7])

%% Estimated disturbance force ft from measurement and uniaxial model
figure;
hold on;
plot(f, sqrt(psd_ft), 'DisplayName', '$\Gamma_{f_{t}}$');
hold off;
set(gca, 'xscale', 'log'); set(gca, 'yscale', 'log');
xlabel('Frequency [Hz]'); ylabel('Ampl. Spectral Density $\left[\frac{N}{\sqrt{Hz}}\right]$')
xlim([1, 500]);
legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);

%% Save PSD of disturbances
save('./mat/uniaxial_disturbance_psd.mat', 'f', 'psd_ft', 'psd_xf');