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

% System Poles: Campbell Diagram
% The poles of $\mathbf{G}_d$ are the complex solutions $p$ of equation eqref:eq:poles.

% #+name: eq:poles
% \begin{equation}
%   \left( \frac{p^2}{{\omega_0}^2} + 2 \xi \frac{p}{\omega_0} + 1 - \frac{{\Omega}^2}{{\omega_0}^2} \right)^2 + \left( 2 \frac{\Omega}{\omega_0} \frac{p}{\omega_0} \right)^2 = 0
% \end{equation}

% Supposing small damping ($\xi \ll 1$), two pairs of complex conjugate poles are obtained as shown in equation eqref:eq:pole_values.

% #+name: eq:pole_values
% \begin{subequations}
%   \begin{align}
%     p_{+} &= - \xi \omega_0 \left( 1 + \frac{\Omega}{\omega_0} \right) \pm j \omega_0 \left( 1 + \frac{\Omega}{\omega_0} \right) \\
%     p_{-} &= - \xi \omega_0 \left( 1 - \frac{\Omega}{\omega_0} \right) \pm j \omega_0 \left( 1 - \frac{\Omega}{\omega_0} \right)
%   \end{align}
% \end{subequations}


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



% The real and complex parts of these two pairs of complex conjugate poles are represented in Figure ref:fig:rotating_campbell_diagram as a function of the rotational speed $\Omega$.
% As the rotational speed increases, $p_{+}$ goes to higher frequencies and $p_{-}$ goes to lower frequencies.
% The system becomes unstable for $\Omega > \omega_0$ as the real part of $p_{-}$ is positive.
% Physically, the negative stiffness term $-m\Omega^2$ induced by centrifugal forces exceeds the spring stiffness $k$.


%% Campbell diagram - Real and Imaginary parts of the poles as a function of the rotating velocity
figure;
tiledlayout(1, 2, 'TileSpacing', 'Compact', 'Padding', 'None');

ax1 = nexttile();
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$', ''})

ax2 = nexttile();
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 Generic Dynamics                                       :noexport:

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

% System Dynamics: Effect of rotation
% The system dynamics from actuator forces $[F_u, F_v]$ to the relative motion $[d_u, d_v]$ is identified for several rotating velocities.

% Looking at the transfer function matrix $\mathbf{G}_d$ in equation eqref:eq:Gd_w0_xi_k, one can see that the two diagonal (direct) terms are equal and that the two off-diagonal (coupling) terms are opposite.
% The bode plot of these two terms are shown in Figure ref:fig:rotating_direct_coupling_bode_plot for several rotational speeds $\Omega$.
% These plots confirm the expected behavior: the frequency of the two pairs of complex conjugate poles are further separated as $\Omega$ increases.
% For $\Omega > \omega_0$, the low frequency pair of complex conjugate poles $p_{-}$ becomes unstable.


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

figure;
tiledlayout(3, 2, '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]');
title('Direct terms: $d_u/F_u$, $d_v/F_v$');
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([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',[]); set(gca, 'YTickLabel',[]);
title('Coupling Terms: $d_u/F_v$, $d_v/F_u$');
ldg = legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 1);
ldg.ItemTokenSize = [10, 1];
ylim([1e-2, 1e2]);

ax3 = 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$'})

ax4 = 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]'); set(gca, 'YTickLabel',[]);
xticks([1e-1,1,1e1])
xticklabels({'$0.1 \omega_0$', '$\omega_0$', '$10 \omega_0$'})

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

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