%% Clear Workspace and Close figures
clear; close all; clc;

%% Intialize Laplace variable
s = zpk('s');

addpath('active_damping/src/');

open('nass_model.slx')

load('./active_damping/mat/undamped_plants.mat', 'G_dvf');
load('./active_damping/mat/plants_variable.mat', 'masses', 'Gm_dvf');

freqs = logspace(0, 3, 1000);

figure;

ax1 = subplot(2, 1, 1);
hold on;
for i=1:length(masses)
  plot(freqs, abs(squeeze(freqresp(-Gm_dvf{i}('Dnlm1', 'Fnl1'), freqs, 'Hz'))));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); set(gca, 'XTickLabel',[]);

ax2 = subplot(2, 1, 2);
hold on;
for i=1:length(masses)
  plot(freqs, 180/pi*angle(squeeze(freqresp(-Gm_dvf{i}('Dnlm1', 'Fnl1'), freqs, 'Hz'))), ...
       'DisplayName', sprintf('$M = %.0f$ [kg]', masses(i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-180, 180]);
yticks([-180, -90, 0, 90, 180]);
legend('location', 'southwest');

linkaxes([ax1,ax2],'x');

K_dvf = s*30000/(1 + s/2/pi/10000);

freqs = logspace(0, 3, 1000);

figure;

ax1 = subplot(2, 1, 1);
hold on;
for i=1:length(masses)
  plot(freqs, abs(squeeze(freqresp(K_dvf*Gm_dvf{i}('Dnlm1', 'Fnl1'), freqs, 'Hz'))));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Loop Gain'); set(gca, 'XTickLabel',[]);

ax2 = subplot(2, 1, 2);
hold on;
for i=1:length(masses)
  plot(freqs, 180/pi*angle(squeeze(freqresp(K_dvf*Gm_dvf{i}('Dnlm1', 'Fnl1'), freqs, 'Hz'))), ...
       'DisplayName', sprintf('$M = %.0f$ [kg]', masses(i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
ylim([-180, 180]);
yticks([-180, -90, 0, 90, 180]);
legend('location', 'southwest');

linkaxes([ax1,ax2],'x');

K_dvf = -K_dvf*eye(6);

save('./active_damping/mat/K_dvf.mat', 'K_dvf');

prepareLinearizeIdentification();

load('./active_damping/mat/K_dvf.mat', 'K_dvf');
initializeController('type', 'dvf', 'K', K_dvf);

%% Options for Linearized
options = linearizeOptions;
options.SampleTime = 0;

%% Name of the Simulink File
mdl = 'nass_model';

%% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/Controller'],     1, 'openinput');  io_i = io_i + 1; % Actuator Inputs
io(io_i) = linio([mdl, '/Tracking Error'], 1, 'output', [], 'En'); io_i = io_i + 1;

load('./active_damping/mat/cart_plants.mat', 'masses');

G_cart_dvf = {zeros(length(masses))};

load('mat/stages.mat', 'nano_hexapod');

for i = 1:length(masses)
  initializeSample('mass', masses(i));

  %% Run the linearization
  G = linearize(mdl, io, 0.3, options);
  G.InputName  = {'Fnl1', 'Fnl2', 'Fnl3', 'Fnl4', 'Fnl5', 'Fnl6'};
  G.OutputName = {'Dnx', 'Dny', 'Dnz', 'Rnx', 'Rny', 'Rnz'};

  G_cart = G*inv(nano_hexapod.J');
  G_cart.InputName = {'Fnx', 'Fny', 'Fnz', 'Mnx', 'Mny', 'Mnz'};

  G_cart_dvf(i) = {G_cart};
end

save('./active_damping/mat/cart_plants.mat', 'G_cart_dvf', '-append');

load('./active_damping/mat/cart_plants.mat', 'masses', 'G_cart', 'G_cart_dvf');

freqs = logspace(0, 3, 1000);

figure;

ax1 = subplot(2, 1, 1);
hold on;
for i = 1:length(masses)
  set(gca,'ColorOrderIndex',i);
  p1 = plot(freqs, abs(squeeze(freqresp(G_cart_dvf{i}('Dnx', 'Fnx'), freqs, 'Hz'))));
  set(gca,'ColorOrderIndex',i);
  p2 = plot(freqs, abs(squeeze(freqresp(G_cart_dvf{i}('Dny', 'Fny'), freqs, 'Hz'))), '--');
  set(gca,'ColorOrderIndex',i);
  p3 = plot(freqs, abs(squeeze(freqresp(G_cart_dvf{i}('Dnz', 'Fnz'), freqs, 'Hz'))), ':');
end
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); xlabel('Frequency [Hz]');
legend([p1,p2,p3], {'Fx/Dx', 'Fy/Dx', 'Fz/Dz'});

ax2 = subplot(2, 1, 2);
hold on;
for i = 1:length(masses)
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_dvf{i}('Dnx', 'Fnx'), freqs, 'Hz')))), ...
       'DisplayName', sprintf('$M = %.0f$ [kg]', masses(i)));
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_dvf{i}('Dny', 'Fny'), freqs, 'Hz')))), '--', 'HandleVisibility', 'off');
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_dvf{i}('Dnz', 'Fnz'), freqs, 'Hz')))), ':', 'HandleVisibility', 'off');
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
yticks([-540:180:540]);
legend('location', 'northeast');

linkaxes([ax1,ax2],'x');

freqs = logspace(0, 3, 1000);

figure;

ax1 = subplot(2, 1, 1);
hold on;
for i = 1:length(masses)
  set(gca,'ColorOrderIndex',i);
  p1 = plot(freqs, abs(squeeze(freqresp(G_cart_dvf{i}('Rnx', 'Mnx'), freqs, 'Hz'))));
  set(gca,'ColorOrderIndex',i);
  p2 = plot(freqs, abs(squeeze(freqresp(G_cart_dvf{i}('Rny', 'Mny'), freqs, 'Hz'))), '--');
  set(gca,'ColorOrderIndex',i);
  p3 = plot(freqs, abs(squeeze(freqresp(G_cart_dvf{i}('Rnz', 'Mnz'), freqs, 'Hz'))), ':');
end
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude [m/N]'); xlabel('Frequency [Hz]');
legend([p1,p2,p3], {'Fx/Dx', 'Fy/Dx', 'Fz/Dz'});

ax2 = subplot(2, 1, 2);
hold on;
for i = 1:length(masses)
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_dvf{i}('Rnx', 'Mnx'), freqs, 'Hz')))), ...
       'DisplayName', sprintf('$M = %.0f$ [kg]', masses(i)));
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_dvf{i}('Rny', 'Mny'), freqs, 'Hz')))), '--', 'HandleVisibility', 'off');
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_dvf{i}('Rnz', 'Mnz'), freqs, 'Hz')))), ':', 'HandleVisibility', 'off');
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
ylabel('Phase [deg]'); xlabel('Frequency [Hz]');
yticks([-540:180:540]);
legend('location', 'northeast');

linkaxes([ax1,ax2],'x');

freqs = logspace(1, 3, 1000);

figure;

for ix = 1:6
  for iy = 1:6
    subplot(6, 6, (ix-1)*6 + iy);
    hold on;
    plot(freqs, abs(squeeze(freqresp(G_cart{1}(ix, iy), freqs, 'Hz'))), 'k-');
    plot(freqs, abs(squeeze(freqresp(G_cart_dvf{1}(ix, iy), freqs, 'Hz'))), 'k--');
    set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
    ylim([1e-13, 1e-4]);
    xticks([])
    yticks([])
  end
end

prepareTomographyExperiment();

load('./active_damping/mat/K_dvf.mat', 'K_dvf');
initializeController('type', 'dvf', 'K', K_dvf);

load('mat/conf_simulink.mat');
set_param(conf_simulink, 'StopTime', '4.5');

sim('nass_model');

En_dvf = En;
Eg_dvf = Eg;
save('./active_damping/mat/tomo_exp.mat', 'En_dvf', 'Eg_dvf', '-append');

load('./active_damping/mat/tomo_exp.mat', 'En', 'En_dvf');
Fs = 1e3; % Sampling Frequency of the Data
t = (1/Fs)*[0:length(En(:,1))-1];

figure;
hold on;
plot(En(:,1),     En(:,2),     'DisplayName', '$\epsilon_{x,y}$ - OL')
plot(En_dvf(:,1), En_dvf(:,2), 'DisplayName', '$\epsilon_{x,y}$ - DVF')
xlabel('X Motion [m]'); ylabel('Y Motion [m]');
legend();

figure;
ax1 = subplot(3, 1, 1);
hold on;
plot(t, En(:,1), 'DisplayName', '$\epsilon_{x}$')
plot(t, En_dvf(:,1), 'DisplayName', '$\epsilon_{x}$ - DVF')
legend();

ax2 = subplot(3, 1, 2);
hold on;
plot(t, En(:,2), 'DisplayName', '$\epsilon_{y}$')
plot(t, En_dvf(:,2), 'DisplayName', '$\epsilon_{y}$ - DVF')
legend();
ylabel('Position Error [m]');

ax3 = subplot(3, 1, 3);
hold on;
plot(t, En(:,3), 'DisplayName', '$\epsilon_{z}$')
plot(t, En_dvf(:,3), 'DisplayName', '$\epsilon_{z}$ - DVF')
legend();
xlabel('Time [s]');

linkaxes([ax1,ax2,ax3],'x');
xlim([0.5,inf]);

figure;
ax1 = subplot(3, 1, 1);
hold on;
plot(t, En(:,4), 'DisplayName', '$\epsilon_{\theta_x}$')
plot(t, En_dvf(:,4), 'DisplayName', '$\epsilon_{\theta_x}$ - DVF')
legend();

ax2 = subplot(3, 1, 2);
hold on;
plot(t, En(:,5), 'DisplayName', '$\epsilon_{\theta_y}$')
plot(t, En_dvf(:,5), 'DisplayName', '$\epsilon_{\theta_y}$ - DVF')
legend();
ylabel('Position Error [rad]');

ax3 = subplot(3, 1, 3);
hold on;
plot(t, En(:,6), 'DisplayName', '$\epsilon_{\theta_z}$')
plot(t, En_dvf(:,6), 'DisplayName', '$\epsilon_{\theta_z}$ - DVF')
legend();
xlabel('Time [s]');

linkaxes([ax1,ax2,ax3],'x');
xlim([0.5,inf]);