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

%% Colors for the figures
colors = colororder;

%% Simscape model name
mdl = 'rotating_model';

%% Model parameters for first analysis
kn = 1;    % Actuator Stiffness [N/m]
mn = 1;    % Payload Mass [kg]
cn = 0.05; % Actuator Damping [N/(m/s)]

xin = cn/(2*sqrt(kn*mn)); % Modal Damping [-]
w0n = sqrt(kn/mn);       % Natural Frequency [rad/s]

%% Computation of the poles as a function of the rotating velocity
Wzs = linspace(0, 2, 51); % Vector of rotation speeds [rad/s]

p_ws = zeros(4, length(Wzs));

for i = 1:length(Wzs)
    Wz = Wzs(i);

    pole_G = pole(1/(((s^2)/(w0n^2) + 2*xin*s/w0n + 1 - (Wz^2)/(w0n^2))^2 + (2*Wz*s/(w0n^2))^2));
    [~, i_sort] = sort(imag(pole_G));
    p_ws(:, i) = pole_G(i_sort);
end

clear pole_G;

%% Campbell diagram - Real and Imaginary parts of the poles as a function of the rotating velocity
figure;
hold on;
plot(Wzs, real(p_ws(1, :)), '-', 'color', colors(1,:), 'DisplayName', '$p_{+}$')
plot(Wzs, real(p_ws(4, :)), '-', 'color', colors(1,:), 'HandleVisibility', 'off')
plot(Wzs, real(p_ws(2, :)), '-', 'color', colors(2,:), 'DisplayName', '$p_{-}$')
plot(Wzs, real(p_ws(3, :)), '-', 'color', colors(2,:), 'HandleVisibility', 'off')
plot(Wzs, zeros(size(Wzs)), 'k--', 'HandleVisibility', 'off')
hold off;
xlabel('Rotational Speed $\Omega$'); ylabel('Real Part');
leg = legend('location', 'northwest', 'FontSize', 8);
leg.ItemTokenSize(1) = 8;
xlim([0, 2*w0n]);
xticks([0,w0n/2,w0n,3/2*w0n,2*w0n])
xticklabels({'$0$', '', '$\omega_0$', '', '$2 \omega_0$'})
ylim([-3*xin, 3*xin]);
yticks([-3*xin, -2*xin, -xin, 0, xin, 2*xin, 3*xin])
yticklabels({'', '', '$-\xi\omega_0$', '$0$', ''})

figure
hold on;
plot(Wzs, imag(p_ws(1, :)), '-', 'color', colors(1,:))
plot(Wzs, imag(p_ws(4, :)), '-', 'color', colors(1,:))
plot(Wzs, imag(p_ws(2, :)), '-', 'color', colors(2,:))
plot(Wzs, imag(p_ws(3, :)), '-', 'color', colors(2,:))
plot(Wzs, zeros(size(Wzs)), 'k--')
hold off;
xlabel('Rotational Speed $\Omega$'); ylabel('Imaginary Part');
xlim([0, 2*w0n]);
xticks([0,w0n/2,w0n,3/2*w0n,2*w0n])
xticklabels({'$0$', '', '$\omega_0$', '', '$2 \omega_0$'})
ylim([-3*w0n, 3*w0n]);
yticks([-3*w0n, -2*w0n, -w0n, 0, w0n, 2*w0n, 3*w0n])
yticklabels({'', '', '$-\omega_0$', '$0$', '$\omega_0$', '', ''})

%% Identify the dynamics for several rotating velocities
% Sample
ms = 0.5; % Sample mass [kg]

% Tuv Stage
kn = 1; % Stiffness [N/m]
mn = 0.5; % Tuv mass [kg]
cn = 0.01*2*sqrt(kn*(mn+ms)); % Damping [N/(m/s)]

% General Configuration
model_config = struct();
model_config.controller = "open_loop"; % Default: Open-Loop
model_config.Tuv_type   = "normal";    % Default: 2DoF stage

% Input/Output definition
clear io; io_i = 1;
io(io_i) = linio([mdl, '/controller'],        1, 'openinput');  io_i = io_i + 1; % [Fu, Fv]
io(io_i) = linio([mdl, '/fd'],                1, 'openinput');  io_i = io_i + 1; % [Fdu, Fdv]
io(io_i) = linio([mdl, '/xf'],                1, 'openinput');  io_i = io_i + 1; % [Dfx, Dfy]
io(io_i) = linio([mdl, '/translation_stage'], 1, 'openoutput'); io_i = io_i + 1; % [Fmu, Fmv]
io(io_i) = linio([mdl, '/translation_stage'], 2, 'openoutput'); io_i = io_i + 1; % [Du, Dv]
io(io_i) = linio([mdl, '/ext_metrology'],     1, 'openoutput'); io_i = io_i + 1; % [Dx, Dy]

% Tested rotating velocities [rad/s]
Wzs = [0, 0.1, 0.2, 0.7, 1.2]; % Vector of rotation speeds [rad/s]

Gs  = {zeros(2, 2, length(Wzs))}; % Direct terms

for i = 1:length(Wzs)
    Wz = Wzs(i);

    %% Linearize the model
    G = linearize(mdl, io, 0);

    %% Input/Output definition
    G.InputName = {'Fu', 'Fv', 'Fdx', 'Fdy', 'Dfx', 'Dfy'};
    G.OutputName = {'fu', 'fv', 'Du', 'Dv', 'Dx', 'Dy'};

    Gs{:,:,i} = G;
end

% Save All Identified Plants
save('./mat/rotating_generic_plants.mat', 'Gs', 'Wzs');

%% Bode plot of the direct and coupling terms for several rotating velocities
freqs = logspace(-1, 1, 1000);

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

% Magnitude
ax1 = nexttile([2, 1]);
hold on;
for i = 1:length(Wzs)
    plot(freqs, abs(squeeze(freqresp(Gs{i}('du', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(i,:))
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
set(gca, 'XTickLabel',[]); ylabel('Magnitude [m/N]');
ylim([1e-2, 1e2]);
yticks([1e-2,1e-1,1,1e1,1e2])
yticklabels({'$0.01/k$', '$0.1/k$', '$1/k$', '$10/k$', '$100/k$'})

ax2 = nexttile;
hold on;
for i = 1:length(Wzs)
    plot(freqs, 180/pi*angle(squeeze(freqresp(Gs{i}('du', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(i,:))
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [rad/s]'); ylabel('Phase [deg]');
yticks(-180:90:180);
ylim([-180 180]);
xticks([1e-1,1,1e1])
xticklabels({'$0.1 \omega_0$', '$\omega_0$', '$10 \omega_0$'})

linkaxes([ax1,ax2],'x');
xlim([freqs(1), freqs(end)]);

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

ax1 = nexttile([2, 1]);
hold on;
for i = 1:length(Wzs)
    plot(freqs, abs(squeeze(freqresp(Gs{i}('dv', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(i,:), ...
         'DisplayName', sprintf('$\\Omega = %.1f \\omega_0$', Wzs(i)))
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
set(gca, 'XTickLabel',[]); ylabel('Magnitude [m/N]');
ldg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
ldg.ItemTokenSize = [10, 1];
ylim([1e-2, 1e2]);
yticks([1e-2,1e-1,1,1e1,1e2])
yticklabels({'$0.01/k$', '$0.1/k$', '$1/k$', '$10/k$', '$100/k$'})


ax2 = nexttile;
hold on;
for i = 1:length(Wzs)
    plot(freqs, 180/pi*angle(squeeze(freqresp(Gs{i}('dv', 'Fu'), freqs, 'rad/s'))), '-', 'color', colors(i,:));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [rad/s]'); ylabel('Phase [deg]');
yticks(-180:90:180);
ylim([-180 180]);
xticks([1e-1,1,1e1])
xticklabels({'$0.1 \omega_0$', '$\omega_0$', '$10 \omega_0$'})

linkaxes([ax1,ax2],'x');
xlim([freqs(1), freqs(end)]);
