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

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

colors = colororder;

addpath('./mat/');
addpath('./src/');

%% Numnbers of the measured legs
leg_nums = [1 2 3 4 5];

%% First identification (low frequency noise)
leg_noise = {};
for i = 1:length(leg_nums)
    leg_noise(i) = {load(sprintf('frf_data_leg_coder_%i_noise.mat', leg_nums(i)), 't', 'Va', 'Vs', 'de', 'da')};
end

%% Second identification (high frequency noise)
leg_noise_hf = {};
for i = 1:length(leg_nums)
    leg_noise_hf(i) = {load(sprintf('frf_data_leg_coder_%i_noise_hf.mat', leg_nums(i)), 't', 'Va', 'Vs', 'de', 'da')};
end

%% Time vector
t = leg_noise{1}.t - leg_noise{1}.t(1) ; % Time vector [s]

%% Sampling
Ts = (t(end) - t(1))/(length(t)-1); % Sampling Time [s]
Fs = 1/Ts; % Sampling Frequency [Hz]

win = hanning(ceil(0.5*Fs)); % Hannning Windows

% Only used to have the frequency vector "f"
[~, f] = tfestimate(leg_noise{1}.Va, leg_noise{1}.de, win, [], [], 1/Ts);
i_lf = f <= 350;
i_hf = f >  350;

%% Coherence computation
coh_enc = zeros(length(f), length(leg_nums));
for i = 1:length(leg_nums)
    [coh_lf, ~] = mscohere(leg_noise{i}.Va,    leg_noise{i}.de,    win, [], [], 1/Ts);
    [coh_hf, ~] = mscohere(leg_noise_hf{i}.Va, leg_noise_hf{i}.de, win, [], [], 1/Ts);
    coh_enc(:, i) = [coh_lf(i_lf); coh_hf(i_hf)];
end

%% Plot the coherence
figure;
hold on;
for i = 1:length(leg_nums)
    plot(f, coh_enc(:, i));
end;
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Coherence [-]');
xlim([10, 2e3]); ylim([0, 1]);

%% Transfer function estimation
enc_frf = zeros(length(f), length(leg_nums));

for i = 1:length(leg_nums)
    [frf_lf, ~] = tfestimate(leg_noise{i}.Va,    leg_noise{i}.de,    win, [], [], 1/Ts);
    [frf_hf, ~] = tfestimate(leg_noise_hf{i}.Va, leg_noise_hf{i}.de, win, [], [], 1/Ts);
    enc_frf(:, i) = [frf_lf(i_lf); frf_hf(i_hf)];
end

%% Bode plot of the FRF from Va to de
figure;
tiledlayout(3, 1, 'TileSpacing', 'None', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:length(leg_nums)
    plot(f, abs(enc_frf(:, i)), ...
         'DisplayName', sprintf('Leg %i', leg_nums(i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude $d_e/V_a$ [m/V]'); set(gca, 'XTickLabel',[]);
hold off;
legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 2);
ylim([1e-8, 1e-3]);

ax2 = nexttile;
hold on;
for i = 1:length(leg_nums)
    plot(f, 180/pi*angle(enc_frf(:, i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
hold off;
yticks(-360:90:360); ylim([-180, 180]);

linkaxes([ax1,ax2],'x');
xlim([10, 2e3]);

%% Coherence computation
coh_int = zeros(length(f), length(leg_nums));
for i = 1:length(leg_nums)
    [coh_lf, ~] = mscohere(leg_noise{i}.Va,    leg_noise{i}.da,    win, [], [], 1/Ts);
    [coh_hf, ~] = mscohere(leg_noise_hf{i}.Va, leg_noise_hf{i}.da, win, [], [], 1/Ts);
    coh_int(:, i) = [coh_lf(i_lf); coh_hf(i_hf)];
end

%% Plot coherence
figure;
hold on;
for i = 1:length(leg_nums)
    plot(f, coh_int(:, i));
end;
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Coherence [-]');
xlim([10, 2e3]); ylim([0, 1]);

%% Transfer function estimation
int_frf = zeros(length(f), length(leg_nums));
for i = 1:length(leg_nums)
    [frf_lf, ~] = tfestimate(leg_noise{i}.Va,    leg_noise{i}.da,    win, [], [], 1/Ts);
    [frf_hf, ~] = tfestimate(leg_noise_hf{i}.Va, leg_noise_hf{i}.da, win, [], [], 1/Ts);
    int_frf(:, i) = [frf_lf(i_lf); frf_hf(i_hf)];
end

%% Plot the FRF from Va to de (interferometer)
figure;
tiledlayout(3, 1, 'TileSpacing', 'None', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:length(leg_nums)
    plot(f, abs(int_frf(:, i)), ...
         'DisplayName', sprintf('Leg %i', leg_nums(i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude $d_a/V_a$ [m/V]'); set(gca, 'XTickLabel',[]);
hold off;
legend('location', 'northeast', 'FontSize', 8, 'NumColumns', 2);
ylim([1e-9, 1e-3]);

ax2 = nexttile;
hold on;
for i = 1:length(leg_nums)
    plot(f, 180/pi*angle(int_frf(:, i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
hold off;
yticks(-360:90:360); ylim([-180 180]);

linkaxes([ax1,ax2],'x');
xlim([10, 2e3]);

%% Coherence
coh_iff = zeros(length(f), length(leg_nums));
for i = 1:length(leg_nums)
    [coh_lf, ~] = mscohere(leg_noise{i}.Va,    leg_noise{i}.Vs,    win, [], [], 1/Ts);
    [coh_hf, ~] = mscohere(leg_noise_hf{i}.Va, leg_noise_hf{i}.Vs, win, [], [], 1/Ts);
    coh_iff(:, i) = [coh_lf(i_lf); coh_hf(i_hf)];
end

%% Plot the coherence
figure;
hold on;
for i = 1:length(leg_nums)
    plot(f, coh_iff(:, i));
end;
hold off;
xlabel('Frequency [Hz]'); ylabel('Coherence [-]');
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlim([10, 2e3]); ylim([0, 1]);

%% FRF estimation of the transfer function from Va to Vs
iff_frf = zeros(length(f), length(leg_nums));
for i = 1:length(leg_nums)
    [frf_lf, ~] = tfestimate(leg_noise{i}.Va,    leg_noise{i}.Vs,    win, [], [], 1/Ts);
    [frf_hf, ~] = tfestimate(leg_noise_hf{i}.Va, leg_noise_hf{i}.Vs, win, [], [], 1/Ts);
    iff_frf(:, i) = [frf_lf(i_lf); frf_hf(i_hf)];
end

%% Plot the FRF from Va to Vs
figure;
tiledlayout(3, 1, 'TileSpacing', 'None', 'Padding', 'None');

ax1 = nexttile([2,1]);
hold on;
for i = 1:length(leg_nums)
    plot(f, abs(iff_frf(:, i)), ...
         'DisplayName', sprintf('Leg %i', leg_nums(i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'log');
ylabel('Amplitude $V_s/V_a$ [V/V]'); set(gca, 'XTickLabel',[]);
hold off;
ylim([1e-2, 1e2]);
legend('location', 'southeast', 'FontSize', 8, 'NumColumns', 2);

ax2 = nexttile;
hold on;
for i = 1:length(leg_nums)
    plot(f, 180/pi*angle(iff_frf(:, i)));
end
hold off;
set(gca, 'XScale', 'log'); set(gca, 'YScale', 'lin');
xlabel('Frequency [Hz]'); ylabel('Phase [deg]');
hold off;
yticks(-360:90:360); ylim([-180 180]);

linkaxes([ax1,ax2],'x');
xlim([10, 2e3]);

%% Save the estimated FRF for further analysis
save('mat/meas_struts_frf.mat', 'f', 'Ts', 'enc_frf', 'int_frf', 'iff_frf', 'leg_nums');