%% 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_ine');
load('./active_damping/mat/plants_variable.mat', 'masses', 'Gm_ine');

freqs = logspace(0, 3, 1000);

figure;

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

ax2 = subplot(2, 1, 2);
hold on;
for i=1:length(masses)
  plot(freqs, 180/pi*angle(squeeze(freqresp(Gm_ine{i}('Vnlm1', '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_ine = 2.5e4;

freqs = logspace(0, 3, 1000);

figure;

ax1 = subplot(2, 1, 1);
hold on;
for i=1:length(masses)
  plot(freqs, abs(squeeze(freqresp(K_ine*Gm_ine{i}('Vnlm1', '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_ine*Gm_ine{i}('Vnlm1', '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_ine = -K_ine*eye(6);

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

prepareLinearizeIdentification();

load('./active_damping/mat/K_ine.mat', 'K_ine');
initializeController('type', 'ine', 'K', K_ine);

%% 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_ine = {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_ine(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_ine');

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_ine{i}('Dnx', 'Fnx'), freqs, 'Hz'))));
  set(gca,'ColorOrderIndex',i);
  p2 = plot(freqs, abs(squeeze(freqresp(G_cart_ine{i}('Dny', 'Fny'), freqs, 'Hz'))), '--');
  set(gca,'ColorOrderIndex',i);
  p3 = plot(freqs, abs(squeeze(freqresp(G_cart_ine{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_ine{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_ine{i}('Dny', 'Fny'), freqs, 'Hz')))), '--', 'HandleVisibility', 'off');
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_ine{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_ine{i}('Rnx', 'Mnx'), freqs, 'Hz'))));
  set(gca,'ColorOrderIndex',i);
  p2 = plot(freqs, abs(squeeze(freqresp(G_cart_ine{i}('Rny', 'Mny'), freqs, 'Hz'))), '--');
  set(gca,'ColorOrderIndex',i);
  p3 = plot(freqs, abs(squeeze(freqresp(G_cart_ine{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_ine{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_ine{i}('Rny', 'Mny'), freqs, 'Hz')))), '--', 'HandleVisibility', 'off');
  set(gca,'ColorOrderIndex',i);
  plot(freqs, 180/pi*unwrap(angle(squeeze(freqresp(G_cart_ine{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_ine{1}(ix, iy), freqs, 'Hz'))), 'k--');
    set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
    ylim([1e-13, 1e-4]);
    xticks([])
    yticks([])
  end
end