clear all
close all
L = @(x) strrep(x, '_', '\_');  % hack to get the correct channel name in plots 

%% parameters
gpsb = 1096480216;  % starting time
dt = 600;          % duration

fs_out = 512;           % sampling frequency for all computations 
bands = [50, 150];   
lpfreq = 0.5;             % and the low pass filter frequency

T = 20;                 % how many seconds of data to throw away to avoid filter transients 

%% List of channels, the first is the main one, the other are the auxiliary channels
chans={'H1:IMC-IM4_TRANS_SUM_OUT_DQ', ...
       'H1:IMC-MC2_TRANS_PIT_OUT_DQ', 'H1:IMC-MC2_TRANS_YAW_OUT_DQ', ...
       'H1:IMC-IM4_TRANS_PIT_OUT_DQ', 'H1:IMC-IM4_TRANS_YAW_OUT_DQ', ...
       'H1:IMC-WFS_A_I_PIT_OUT_DQ', 'H1:IMC-WFS_A_I_YAW_OUT_DQ', ...
       'H1:IMC-WFS_B_I_PIT_OUT_DQ', 'H1:IMC-WFS_B_I_YAW_OUT_DQ'};
   
 chans={'H1:IMC-IM4_TRANS_SUM_OUT_DQ', ...
        'H1:IMC-DOF_1_P_IN1_DQ', 'H1:IMC-DOF_1_Y_IN1_DQ', ...
        'H1:IMC-DOF_2_P_IN1_DQ', 'H1:IMC-DOF_2_Y_IN1_DQ'};
 
nx = 4; % number of rows and columns in scatter plot
ny = 1;


%% get all data
data=get_data(chans,'raw',gpsb,dt);
df = decimate(data(1).data, data(1).rate/fs_out);

aux = [];
for i=2:length(data)
    aux(:,i-1) = decimate(data(i).data, data(i).rate/fs_out);
end

%% compute the BRMS

% define low pass filter
Nyquist_f=fs_out/2; 
[bb, aa] = butter(4, lpfreq/Nyquist_f, 'low');
    
pd = [];
for i=1:size(bands,1)
    
    % create filter with a more numerically stable algorithm
    [z,p,k] = butter(4, [bands(i,1)/Nyquist_f, bands(i,2)/Nyquist_f], 'bandpass');
    [sos,g] = zp2sos(z, p, k);
    fl = dfilt.df2sos(sos, g);
    
    % filter main signal in the band (double pass to avoid phase lag)
    x = fl.filter(df);
    x = fl.filter(x(end:-1:1));
    x = x(end:-1:1);
    % square it to get the BLRMS
    x = x.^2;

    % low pass BLRMS
    x = filtfilt(bb,aa, x);
    % normalize to units of amplitude spectral density
    pd(:,i) = sqrt(abs(x/diff(bands(i,:))));
end

% define time for x axis 
t = (1:size(pd,1))/fs_out;

% low pass all aux channels
for i=1:size(aux,2)
    aux(:,i) = filtfilt(bb,aa,aux(:,i));
end


% throw away first and last seconds to avoid filter transients
pd = pd(T*fs_out:end-T*fs_out,:);
t = t(T*fs_out:end-T*fs_out);
aux = aux(T*fs_out:end-T*fs_out,:);

%% plot time series of BLRMS
figure(1)
plot(t,pd); 
xlabel('Time [s]')
ylabel('Noise amplitude relative varation')
legend(arrayfun(@(i) sprintf('Band %d - %d Hz', bands(i,1), bands(i,2)), 1:size(bands,1), 'UniformOutput', false))

%% plot the scatter histograms (main signal vs aux signals)
for k=1:size(pd,2)
    figure()
    for i=1:size(aux,2)
        subplot(ny, nx, i)
        % define the histogram ranges
        xe = linspace(min(aux(:,i)), max(aux(:,i)), 100);
        ye = linspace(min(pd(:,k)), max(pd(:,k)), 100);
        % make 2d histogram
        h1 = hist2(aux(:,i), pd(:,k), xe, ye);
        % and plot it
        pcolor(xe, ye, h1'); shading flat
        xlabel(L(chans{i+1}))
    end
end

%% Least square fit with channel ranking
j = 1;  % this selects which band you want to analyze

% build the list of all auxiliary channels
X = zeros(size(aux,1), 2*size(aux,2)+1);
for i=1:size(aux,2)
    % add each channel together with its squared values
    X(:,2*i-1) = aux(:,i);
    X(:,2*i) = aux(:,i).^2;
    % and save the names for future reference
    cname{2*i-1} = chans{i+1};
    cname{2*i} = [chans{i+1}, '^2'];
end
% at the end, add also a constant
X(:,end) = 1;
cname{size(X,2)} = '1';

nsig = size(X,2);       % number of auxiliary channels
idx = zeros(nsig,1);    % this will contain the index of the best channels
p = zeros(nsig,1);      % and this will contain tehe corresponding parameter
for k=1:nsig
    % compute the signal to be fitted
    if k>1
        % it is the residual so far
        res = pd(:,j) - X(:, idx(1:k-1)) * p(1:k-1);
    else
        % first iteration, just fit everything
        res = pd(:,j);
    end
    % find the best additional channel to minimize the residual
    for i=1:nsig
        fit = @(x, t) X(:,i)*x;
        % store the coefficient and the residual error
        c(i) = lsqcurvefit(fit, 0, t, res, [], [], optimset('Display', 'off'));
        err(i) = sum( (res - fit(c(i),t)).^2 );
    end
    % the best is the one with minimum residual error, sae the index and
    % the coefficient
    [minerr(k), idx(k)] = min(err);
    p(k) = c(idx(k));
    fprintf(2, '%d) Best channel = %d (%s) - new residual = %d\n', k, idx(k), cname{idx(k)}, minerr(k));
end

% plot the fit result
figure()
subplot(211)
plot(t, X(:,idx)*p)
hold all
plot(t, pd(:,j), 'r', 'Linewidth', 2)
plot(t, pd(:,j) - X(:,idx)*p, 'k', 'Linewidth', 1)
title(sprintf('Band %d - %d Hz', bands(j,1), bands(j,2)))
xlabel('Time [s]')
legend('Reconstructed', 'Measured', 'Residual')
subplot(212)
hist(pd(:,j) - X(:,idx)*p, 100)
xlabel('Residual')

% and plot the ranking 
figure()
ierr = sum(pd(:,j).^2);
plot([ierr, minerr], 0:length(minerr), 'o-', 'LineWidth', 2)
title(sprintf('Band %d - %d Hz', bands(j,1), bands(j,2)))
xlim([min(minerr), max(minerr)])
set(gca, 'YTick', 0:length(minerr), 'YTickLabel', ['Initial (out of scale)', cname(idx)])
xlabel('Cumulative residual error')

% compute the incremental error reduction
ierr = sum(pd(:,j).^2);
derr = -diff([ierr, minerr]);   % this computes the reduction in the residual obtained when adding each channel
[derr, ix] = sort(derr);        % sort them in ascending order
idxs = idx(ix);

% and plot it
figure()
barh(1:length(minerr), derr)
xlim([0, 1.3*max(derr(idxs~=nsig))])
set(gca, 'YTick', 1:length(minerr), 'YTickLabel',  cname(idxs))
xlabel('Residual error reduction')
