%% ustation_3_disturbances.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);

%% Compute Floor Motion Spectral Density
% Load floor motion data
% velocity in [m/s] is measured in X, Y and Z directions
load('ustation_ground_motion.mat', 'Ts', 'Fs', 'vel_x', 'vel_y', 'vel_z', 't');

% Estimate ground displacement from measured velocity
% This is done by integrating the motion
gm_x = lsim(1/(s+0.1*2*pi), vel_x, t);
gm_y = lsim(1/(s+0.1*2*pi), vel_y, t);
gm_z = lsim(1/(s+0.1*2*pi), vel_z, t);

Nfft = floor(2/Ts);
win = hanning(Nfft);
Noverlap = floor(Nfft/2);

[pxx_gm_vx, f_gm] = pwelch(vel_x, win, Noverlap, Nfft, 1/Ts);
[pxx_gm_vy, ~]    = pwelch(vel_y, win, Noverlap, Nfft, 1/Ts);
[pxx_gm_vz, ~]    = pwelch(vel_z, win, Noverlap, Nfft, 1/Ts);

% Convert PSD in velocity to PSD in displacement
pxx_gm_x = pxx_gm_vx./((2*pi*f_gm).^2);
pxx_gm_y = pxx_gm_vy./((2*pi*f_gm).^2);
pxx_gm_z = pxx_gm_vz./((2*pi*f_gm).^2);

%% Measured ground motion
figure;
hold on;
plot(t, 1e6*gm_x, 'DisplayName', '$D_{xf}$')
plot(t, 1e6*gm_y, 'DisplayName', '$D_{yf}$')
plot(t, 1e6*gm_z, 'DisplayName', '$D_{zf}$')
hold off;
xlabel('Time [s]');
ylabel('Ground motion [$\mu$m]')
xlim([0, 5]); ylim([-2, 2])
leg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Ty errors
% Setpoint is in [mm]
% Vertical error is in [um]
ty_errors = load('ustation_errors_ty.mat');

% Compute best straight line to remove it from data
average_error = mean(ty_errors.ty_z')';
straight_line = average_error - detrend(average_error, 1);

%% Measurement of the linear (vertical) deviation of the Translation stage
figure;
hold on;
plot(ty_errors.setpoint, ty_errors.ty_z(:,1), '-', 'color', colors(1,:), 'DisplayName', 'Raw data')
plot(ty_errors.setpoint, ty_errors.ty_z(:,2:end), '-', 'color', colors(1,:), 'HandleVisibility', 'off')
plot(ty_errors.setpoint, straight_line, '--', 'color', colors(2,:), 'DisplayName', 'Linear fit')
hold off;
xlabel('$D_y$ position [mm]'); ylabel('Vertical error [$\mu$m]');
xlim([-5, 5]); ylim([-8, 8]);
legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);

%% Measurement of the linear (vertical) deviation of the Translation stage - Remove best linear fit
figure;
plot(ty_errors.setpoint, ty_errors.ty_z - straight_line, 'k-')
xlabel('$D_y$ position [mm]'); ylabel('Residual vertical error [$\mu$m]');
xlim([-5, 5]); ylim([-0.4, 0.4]);

%% Convert the data to frequency domain
% Suppose a certain constant velocity scan
delta_ty = (ty_errors.setpoint(end) - ty_errors.setpoint(1))/(length(ty_errors.setpoint) - 1); % [mm]
ty_vel = 8*1.125; % [mm/s]
Ts = delta_ty/ty_vel;
Fs = 1/Ts;

% Frequency Analysis
Nfft = floor(length(ty_errors.setpoint)); % Number of frequency points
win = hanning(Nfft); % Windowing
Noverlap = floor(Nfft/2); % Overlap for frequency analysis

% Comnpute the power spectral density
[pxx_dy_dz, f_dy] = pwelch(1e-6*detrend(ty_errors.ty_z, 1), win, Noverlap, Nfft, Fs);
pxx_dy_dz = mean(pxx_dy_dz')';

% Having the PSD of the lateral error equal to the PSD of the vertical error
% is a reasonable assumption (and verified in practice)
pxx_dy_dx = pxx_dy_dz;

%% Spindle Errors
% Errors are already converted to x,y,z,Rx,Ry
% Units are in [m] and [rad]
spindle_errors = load('ustation_errors_spindle.mat');

%% Measured radial errors of the Spindle
figure;
hold on;
plot(1e6*spindle_errors.Dx(1:100:end), 1e6*spindle_errors.Dy(1:100:end), 'DisplayName', 'Raw data')
% Compute best fitting circle
[x0, y0, R] = circlefit(spindle_errors.Dx, spindle_errors.Dy);
theta = linspace(0, 2*pi, 500); % Angle to plot the circle [rad]
plot(1e6*(x0 + R*cos(theta)), 1e6*(y0 + R*sin(theta)), '--', 'DisplayName', 'Best Fit')
hold off;
xlabel('X displacement [$\mu$m]'); ylabel('Y displacement [$\mu$m]');
axis equal
xlim([-1, 1]); ylim([-1, 1]);
xticks([-1, -0.5, 0, 0.5, 1]);
yticks([-1, -0.5, 0, 0.5, 1]);
leg = legend('location', 'none', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Measured axial errors of the Spindle
figure;
plot(spindle_errors.deg(1:100:end)/360, 1e9*spindle_errors.Dz(1:100:end))
xlabel('Rotation [turn]'); ylabel('Z displacement [nm]');
axis square
xlim([0,2]); ylim([-40, 40]);

%% Measured tilt errors of the Spindle
figure;
plot(1e6*spindle_errors.Rx(1:100:end), 1e6*spindle_errors.Ry(1:100:end))
xlabel('X angle [$\mu$rad]'); ylabel('Y angle [$\mu$rad]');
axis equal
xlim([-35, 35]); ylim([-35, 35]);
xticks([-30, -15, 0, 15, 30]);
yticks([-30, -15, 0, 15, 30]);

%% Remove the circular fit from the data
[x0, y0, R] = circlefit(spindle_errors.Dx, spindle_errors.Dy);

% Search the best angular match
fun = @(theta)rms((spindle_errors.Dx - (x0 + R*cos(pi/180*spindle_errors.deg+theta(1)))).^2 + (spindle_errors.Dy - (y0 - R*sin(pi/180*spindle_errors.deg+theta(1)))).^2);
x0 = [0];
delta_theta = fminsearch(fun, 0);

% Compute the remaining error after removing the best circular fit
spindle_errors.Dx_err = spindle_errors.Dx - (x0 + R*cos(pi/180*spindle_errors.deg+delta_theta));
spindle_errors.Dy_err = spindle_errors.Dy - (y0 - R*sin(pi/180*spindle_errors.deg+delta_theta));

%% Convert the data to frequency domain
Ts = (spindle_errors.t(end) - spindle_errors.t(1))/(length(spindle_errors.t)-1); % [s]
Fs = 1/Ts; % [Hz]

% Frequency Analysis
Nfft = floor(2.0*Fs); % Number of frequency points
win = hanning(Nfft); % Windowing
Noverlap = floor(Nfft/2); % Overlap for frequency analysis

% Comnpute the power spectral density
[pxx_rz_dz, f_rz] = pwelch(spindle_errors.Dz,     win, Noverlap, Nfft, Fs);
[pxx_rz_dx, ~   ] = pwelch(spindle_errors.Dx_err, win, Noverlap, Nfft, Fs);
[pxx_rz_dy, ~   ] = pwelch(spindle_errors.Dy_err, win, Noverlap, Nfft, Fs);

%% Extract sensitivity to disturbances from the Simscape model
% Initialize stages
initializeGround();
initializeGranite();
initializeTy();
initializeRy();
initializeRz();
initializeMicroHexapod();

initializeDisturbances('enable', false);
initializeSimscapeConfiguration('gravity', false);
initializeLoggingConfiguration();

% Configure inputs and outputs
clear io; io_i = 1;
io(io_i) = linio([mdl, '/Disturbances/Dwx'], 1, 'openinput'); io_i = io_i + 1; % Vertical Ground Motion
io(io_i) = linio([mdl, '/Disturbances/Dwy'], 1, 'openinput'); io_i = io_i + 1; % Vertical Ground Motion
io(io_i) = linio([mdl, '/Disturbances/Dwz'], 1, 'openinput'); io_i = io_i + 1; % Vertical Ground Motion
io(io_i) = linio([mdl, '/Disturbances/Fdy_x'], 1, 'openinput'); io_i = io_i + 1; % Parasitic force Ty
io(io_i) = linio([mdl, '/Disturbances/Fdy_z'], 1, 'openinput'); io_i = io_i + 1; % Parasitic force Ty
io(io_i) = linio([mdl, '/Disturbances/Frz_x'], 1, 'openinput'); io_i = io_i + 1; % Parasitic force Rz
io(io_i) = linio([mdl, '/Disturbances/Frz_y'], 1, 'openinput'); io_i = io_i + 1; % Parasitic force Rz
io(io_i) = linio([mdl, '/Disturbances/Frz_z'], 1, 'openinput'); io_i = io_i + 1; % Parasitic force Rz
io(io_i) = linio([mdl, '/Micro-Station/metrology_6dof/x'], 1, 'openoutput'); io_i = io_i + 1; % Relative motion - Hexapod/Granite
io(io_i) = linio([mdl, '/Micro-Station/metrology_6dof/y'], 1, 'openoutput'); io_i = io_i + 1; % Relative motion - Hexapod/Granite
io(io_i) = linio([mdl, '/Micro-Station/metrology_6dof/z'], 1, 'openoutput'); io_i = io_i + 1; % Relative motion - Hexapod/Granite

% Run the linearization
Gd = linearize(mdl, io, 0);
Gd.InputName  = {'Dwx', 'Dwy', 'Dwz', 'Fdy_x', 'Fdy_z', 'Frz_x', 'Frz_y', 'Frz_z'};
Gd.OutputName = {'Dx', 'Dy', 'Dz'};

% The identified dynamics are then saved for further use.
save('./mat/ustation_disturbance_sensitivity.mat', 'Gd')

%% Sensitivity to Ground motion
freqs = logspace(log10(10), log10(2e2), 1000);

figure;
hold on;
plot(freqs, abs(squeeze(freqresp(Gd('Dx', 'Dwx'), freqs, 'Hz'))), 'DisplayName', '$D_x/D_{xf}$');
plot(freqs, abs(squeeze(freqresp(Gd('Dy', 'Dwy'), freqs, 'Hz'))), 'DisplayName', '$D_y/D_{yf}$');
plot(freqs, abs(squeeze(freqresp(Gd('Dz', 'Dwz'), freqs, 'Hz'))), 'DisplayName', '$D_z/D_{zf}$');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/m]'); xlabel('Frequency [Hz]');
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Sensitivity to Translation stage disturbance forces
figure;
hold on;
plot(freqs, abs(squeeze(freqresp(Gd('Dx', 'Fdy_x'), freqs, 'Hz'))), 'color', colors(1,:), 'DisplayName', '$D_x/F_{xD_y}$');
plot(freqs, abs(squeeze(freqresp(Gd('Dz', 'Fdy_z'), freqs, 'Hz'))), 'color', colors(3,:), 'DisplayName', '$D_z/F_{zD_y}$');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); xlabel('Frequency [Hz]');
ylim([1e-10, 1e-7]);
leg = legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Sensitivity to Spindle disturbance forces
figure;
hold on;
plot(freqs, abs(squeeze(freqresp(Gd('Dx', 'Frz_x'), freqs, 'Hz'))), 'DisplayName', '$D_x/F_{xR_z}$');
plot(freqs, abs(squeeze(freqresp(Gd('Dy', 'Frz_y'), freqs, 'Hz'))), 'DisplayName', '$D_y/F_{yR_z}$');
plot(freqs, abs(squeeze(freqresp(Gd('Dz', 'Frz_z'), freqs, 'Hz'))), 'DisplayName', '$D_z/F_{zR_z}$');
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); xlabel('Frequency [Hz]');
ylim([1e-10, 1e-7]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Compute the PSD of the equivalent disturbance sources
pxx_rz_fx = pxx_rz_dx./abs(squeeze(freqresp(Gd('Dx', 'Frz_x'), f_rz, 'Hz'))).^2;
pxx_rz_fy = pxx_rz_dy./abs(squeeze(freqresp(Gd('Dy', 'Frz_y'), f_rz, 'Hz'))).^2;
pxx_rz_fz = pxx_rz_dz./abs(squeeze(freqresp(Gd('Dz', 'Frz_z'), f_rz, 'Hz'))).^2;

pxx_dy_fx = pxx_dy_dx./abs(squeeze(freqresp(Gd('Dx', 'Fdy_x'), f_dy, 'Hz'))).^2;
pxx_dy_fz = pxx_dy_dz./abs(squeeze(freqresp(Gd('Dz', 'Fdy_z'), f_dy, 'Hz'))).^2;

%% Save the PSD of the disturbance sources for further used
% in the Simscape model

% Ground motion
min_f = 1; max_f = 100;
gm_dist.f  = f_gm(f_gm < max_f & f_gm > min_f);
gm_dist.pxx_x = pxx_gm_x(f_gm < max_f & f_gm > min_f);
gm_dist.pxx_y = pxx_gm_y(f_gm < max_f & f_gm > min_f);
gm_dist.pxx_z = pxx_gm_z(f_gm < max_f & f_gm > min_f);

% Translation stage
min_f = 0.5; max_f = 500;
dy_dist.f  = f_dy(f_dy < max_f & f_dy > min_f);
dy_dist.pxx_fx = pxx_dy_fx(f_dy < max_f & f_dy > min_f);
dy_dist.pxx_fz = pxx_dy_fz(f_dy < max_f & f_dy > min_f);

% Spindle
min_f = 1; max_f = 500;
rz_dist.f  = f_rz(f_rz < max_f & f_rz > min_f);
rz_dist.pxx_fx = pxx_rz_fx(f_rz < max_f & f_rz > min_f);
rz_dist.pxx_fy = pxx_rz_fy(f_rz < max_f & f_rz > min_f);
rz_dist.pxx_fz = pxx_rz_fz(f_rz < max_f & f_rz > min_f);

% The identified dynamics are then saved for further use.
save('./mat/ustation_disturbance_psd.mat', 'rz_dist', 'dy_dist', 'gm_dist')

%% Ground Motion
figure;
hold on;
plot(f_gm, sqrt(pxx_gm_x), 'DisplayName', '$D_{wx}$');
plot(f_gm, sqrt(pxx_gm_y), 'DisplayName', '$D_{wy}$');
plot(f_gm, sqrt(pxx_gm_z), 'DisplayName', '$D_{wz}$');
set(gca, 'xscale', 'log'); set(gca, 'yscale', 'log');
xlabel('Frequency [Hz]'); ylabel('Spectral Density $\left[\frac{m}{\sqrt{Hz}}\right]$')
xlim([1, 200]);
leg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

figure;
hold on;
plot(f_dy, sqrt(pxx_dy_fx), 'DisplayName', '$F_{xD_y}$');
plot(f_dy, sqrt(pxx_dy_fz), 'DisplayName', '$F_{zD_y}$');
set(gca, 'xscale', 'log'); set(gca, 'yscale', 'log');
xlabel('Frequency [Hz]'); ylabel('Spectral Density $\left[\frac{N}{\sqrt{Hz}}\right]$')
xlim([1, 200]); ylim([1e-3, 1e3]);
leg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

figure;
hold on;
plot(f_rz, sqrt(pxx_rz_fx), 'DisplayName', '$F_{xR_z}$');
plot(f_rz, sqrt(pxx_rz_fy), 'DisplayName', '$F_{yR_z}$');
plot(f_rz, sqrt(pxx_rz_fz), 'DisplayName', '$F_{zR_z}$');
set(gca, 'xscale', 'log'); set(gca, 'yscale', 'log');
xlabel('Frequency [Hz]'); ylabel('Spectral Density $\left[\frac{N}{\sqrt{Hz}}\right]$')
xlim([1, 200]); ylim([1e-3, 1e3]);
leg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

%% Compute time domain disturbance signals
initializeDisturbances();
load('nass_model_disturbances.mat');

figure;
hold on;
plot(Rz.t, Rz.x, 'DisplayName', '$F_{xR_z}$');
plot(Rz.t, Rz.y, 'DisplayName', '$F_{yR_z}$');
plot(Rz.t, Rz.z, 'DisplayName', '$F_{zR_z}$');
xlabel('Time [s]'); ylabel('Amplitude [N]')
xlim([0, 1]); ylim([-100, 100]);
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

figure;
hold on;
plot(Dy.t, Dy.x, 'DisplayName', '$F_{xD_y}$');
plot(Dy.t, Dy.z, 'DisplayName', '$F_{zD_y}$');
xlabel('Time [s]'); ylabel('Amplitude [N]')
xlim([0, 1]); ylim([-60, 60])
leg = legend('location', 'southwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;

figure;
hold on;
plot(Dw.t, 1e6*Dw.x, 'DisplayName', '$D_{xf}$');
plot(Dw.t, 1e6*Dw.y, 'DisplayName', '$D_{yf}$');
plot(Dw.t, 1e6*Dw.z, 'DisplayName', '$D_{zf}$');
xlabel('Time [s]'); ylabel('Amplitude [$\mu$m]')
xlim([0, 1]); ylim([-0.6, 0.6])
leg = legend('location', 'northwest', 'FontSize', 8, 'NumColumns', 1);
leg.ItemTokenSize(1) = 15;