function [Hv, f, Puu, Pyy, Pyu, coh, Sv] = mimoTFE(u,y,Fs,nfft,figNo)
% [Hv, f, Puu, Pyy, Pyu, coh, Sv] = mimoTFE(u,y,Fs,nfft,figNo)
%
% Frequency response function of MIMO system from discrete-time i/o data 
% Welch's averaged periodogram method, in which the matrices u and y are
% segmented, detrended, windowed   The magnitude squared of the FFTs of the
% segments of U and Y are averaged to form Puu and Pyy.  The products of the
% FFTs of the segments of U and Y are averaged to form Pyu.  Hv is the
% geometric mean of the H1 and H2 estimates where H1=Pyu/Puu and H2=Pyy/Pyu.
% The coherence is the ratio H1/H2, and indicates the degree to which 
% the output, y, is linearly related to the input, u.
%
%  INPUT   
%   u      matrix of sampled  input to the system               (r x points)
%   y      matrix of sampled output to the system               (m x points)
%   Fs     sampling frequency ... 1 / delta_t                   (1 x 1)
%  nfft    the number of points used in the fast Fourier transform (1x1) 2^n
%  figNo   optional figure number for plotting spectra
%
%  OUTPUT
%   Hv     the complex-valued Hv freq. response fctn estimate   (nfft/2 x m x r)
%   f      the frequency values  corresponding to H             (1 x nfft/2)
%   Puu    the Hermtian matrix of power spectra of u            (nfft/2 x r x r)
%   Pyy    the Hermtian matrix of power spectra of y            (nfft/2 x m x m)
%   Pyu    the complex-valued cross-power spectra from u to y   (nfft/2 x m x r)
%   coh    the coherence function                               (nfft/2 x m x r)
%   Sv     the standard deviation of the freq. resp. fctn est.  (nfft/2 x m x r)

if nargin < 4, error('mimoTFE needs 4 input arguments: u, y, sr, nfft'); end
if nargin < 5, figNo = 0; end
[r,nu] = size(u);		% number of data points
[m,ny] = size(y);		% number of data points
if ( nu ~= ny ), 
  error('U and Y must be the same number of columns.')
else
  n = nu;
end

nfft = 2^(round(log2(nfft)));   % round to the closest power-of-2

% Defaults 

make_Hv_causal = 0;     % 1: make causal frequency response function, 0: don't

nWindow = nfft;		% number of points in a segment of data
  dflag = 'linear';	% straight-line detrending of data
% dflag = 'mean';	% subtract average value from data
% dflag = 'none';	% no detrending
nOverlap = nWindow / 2;	% number of overlapping points 
index = 1:nWindow;	% the points in the first segment of data
Pw = floor(nfft/20);	% number of points in a rectangular taper

% use Sine window for PSD estimates, use Hann window for TFE estimates

% window = ones(1,nWindow);					% Dirichlet 
% window = sin(pi*index/nWindow);				% Sine  
% window = (1 - cos(2*pi*index/nWindow-1))/2;			% Hann 
% window = 0.53836 - 0.46164*cos(2*pi*index/nWindow));		% Hamming 
% window = sinc(2*index/(nWindow-1)-1);				% Lanczos
% window = exp(-0.5*((index-nWindow/2)/(0.2*nWindow).^2);	% Gauss
  window = [ 0.5*(1-cos(pi*[0:Pw]/Pw)) , ones(1,nWindow-2*Pw-2) , 0.5*(1+cos(pi*[0:Pw]/Pw)) ];							% tapered rect.

K = fix((n-nOverlap)/(nWindow-nOverlap)); 	% number of windows

% Normalizing scale factor ==> asymptotically unbiased
  KNW2 = K*norm(window)^2;

% alternative Normalizing scale factor ==> peaks are about right
% KNW2 = K*sum(window)^2;

%% compute PSD matrices from FFTs and store in 3D arrays =====================

UU = zeros(nfft,r,r);   Puu = zeros(nfft,r,r);   % Puu2 = NaN(nfft,r,r);
YY = zeros(nfft,m,m);   Pyy = zeros(nfft,m,m);   % Pyy2 = NaN(nfft,m,m);
YU = zeros(nfft,m,r);   Pyu = zeros(nfft,m,r);   % Pyu2 = NaN(nfft,m,r);

u = u';  y = y';   % for compatibility with detrend
for k = 1:K

    if strcmp(dflag,'linear')
        uw = (ones(r,1)*window).*detrend(u(index,:),1)';
        yw = (ones(m,1)*window).*detrend(y(index,:),1)';
    elseif strcmp(dflag,'mean')
        uw = (ones(r,1)*window).*detrend(u(index,:),0)';
        yw = (ones(m,1)*window).*detrend(y(index,:),0)';
    else
        uw = (ones(r,1)*window).*u(:,index);
        yw = (ones(m,1)*window).*y(:,index);
    end

    U = fft(uw,nfft,2);
    Y = fft(yw,nfft,2);
    for ii = 1:r
      for jj = 1:r
        UU(:,ii,jj) = U(ii,:).*conj(U(jj,:));
      end
    end
    for ii = 1:m
      for jj = 1:m
        YY(:,ii,jj) = Y(ii,:).*conj(Y(jj,:));
      end
    end
    for ii = 1:m
      for jj = 1:r
        YU(:,ii,jj) = Y(ii,:).*conj(U(jj,:));
      end
    end

    Puu = Puu + UU;		% Puu2 = Puu2 + Uu2.*Uu2; % for testing Sv
    Pyy = Pyy + YY;		% Pyy2 = Pyy2 + Yy2.*Yy2; % for testing Sv
    Pyu = Pyu + YU;		% Pyu2 = Pyu2 + Uu2.*Yy2; % for testing Sv

    index = index + (nWindow - nOverlap);        % advance to the next segment

end

%% compute the H1, H2, and Hv frequency response function estimates =======
for kk =1:nfft
    puu = reshape(Puu(kk,:,:),r,r);
    pyy = reshape(Pyy(kk,:,:),m,m);
    pyu = reshape(Pyu(kk,:,:),m,r);
    H1(kk,:,:) = ( pyu * puu' ) / ( puu * puu' );
    H2(kk,:,:) = ( pyy * pyy' ) / ( pyu'* pyy' );
end

Hv = sqrt(H1) .* sqrt(H2);	% Hv frequency response function estimate
coh = real ( H1 ./ H2 );	% coherence of the frequenc response function estimate
if make_Hv_causal
   hv = ifft(Hv);               % impulse resposne function
   hv = real ( [ hv(1:nfft/2+1) ; zeros(nfft/2-1,1) ] ) ;  % make h causal
   Hv = fft(hv);                % frequency response for causal h
end

%% normalize double-sided spectral density to agree with Parseval's theorem
%% ... trapz(t,x.^2)/T  equals  2*trapz(f',Pxx)

Puu = Puu / ( KNW2 * Fs );
Pyy = Pyy / ( KNW2 * Fs );
Pyu = Pyu / ( KNW2 * Fs );

%% for real-valued signals, retain only the non-negative frequencies ========
if ~any(any(imag([u,y])~=0)),   % if u and y are real
    				% return the first half: non-neg've frequencies
    if rem(nfft,2),		% nfft is odd
	nn_idx = [1:(nfft+1)/2];
    else			% nfft is even, include DC, include Nyquist
	nn_idx = [1:nfft/2+1]; 
    end
    Puu  = Puu(nn_idx,:,:);	% Puu2 = Puu(nn_idx);  % for testing Sv
    Pyy  = Pyy(nn_idx,:,:);	% Pyy2 = Pyy(nn_idx);  % for testing Sv
    Pyu  = Pyu(nn_idx,:,:);	% Pyu2 = Pyu(nn_idx);  % for testing Sv
    Hv   = Hv(nn_idx,:,:);		
    coh  = coh(nn_idx,:,:);		
    f    = (nn_idx - 1)' * Fs/nfft;
else
    % return positive and negative frequencies
    f   = [ [0:nfft/2] [-nfft/2+1:-1] ]' * Fs/nfft;
end

%% Confidence Intervals =====================================================
if nargout == 7 
% approx Sv estimate from Chi-Squared distribution, and 50% window overlap
% from Kay, p. 76, eqn 4.16:
  p = 0.864;			% probability level for confidence intervals
  alpha = 1 - p;
  CI = (K-1) ./ [ chi2inv(alpha/2,K-1) , chi2inv(1-alpha/2,K-1) ];
  Sv = abs(Hv) * sqrt( 0.5 * (CI(1) - CI(2)) ) / K;

%  Sv from squared standard error of the mean PSD estimate ... for testing 
%{
   Suu2 = (Puu2 - KNW2*Fs * Puu.^2 ) / ((K-1)*K);
   Syy2 = (Pyy2 - KNW2*Fs * Pyy.^2 ) / ((K-1)*K);
   Suy2 = (Pyu2 - KNW2*Fs * abs(Pyu).^2) / ((K-1)*K);
    
   Sv = 0.5 * abs(Hv) .* sqrt ( Suu2./Puu2 + 2*Suy2./Pyu2 + Syy2./Pyy2 );
%}

end


%% PLOTS  ==================================================================== 
if (figNo > 0) 

  Hp  = Hv;        % which frequency response estimate to plot ... H1, H2, Hv
  nf  = length(f);
  idx = [2:nfft/8];
  ff  = [ f(idx(1)) , f(idx(end)) ];

  figure(figNo); % INPUT auto power spectra and cross spectra

% unwrap the phase for plotting 
  pha        =  NaN(nf,m,r);
  pha2       = -atan2(imag(Hp(2,:,:)),real(Hp(2,:,:)));  % -ve sign!
  pha(1,:,:) =  pha2;
  pha(2,:,:) =  pha2;
  pha2 = reshape( ones(nf-2,1,1)*pha2 , [nf-2,m,r] );
  d_pha = angle( Hp(3:nf,:,:)   ./ Hp(2:nf-1,:,:) );
  dp = pi;
  d_pha_idx = find(d_pha > +dp);   d_pha(d_pha_idx) = d_pha(d_pha_idx)-2*pi;
  d_pha_idx = find(d_pha < -dp);   d_pha(d_pha_idx) = d_pha(d_pha_idx)+2*pi;
  pha(3:nf,:,:) = pha2 - cumtrapz(d_pha);

%pha = unwrap(pha,0.01);
%pha = unwrap(-atan2(imag(Hp),real(Hp)));  % -ve sign!

  pha_min = floor(min(min(min(pha(idx,:,:)*180/pi))/90))*90;
  pha_max =  ceil(max(max(max(pha(idx,:,:)*180/pi))/90))*90;

   clf
   for ii = 1:r
     for jj = 1:r
       subplot(r,r,(ii-1)*r+jj)
        if ii == jj 
          semilogx(f(idx),Puu(idx,ii,jj))%, f,CI(1)*Puu,'-k', f,CI(2)*Puu,'-k')
          axis('tight')
        else
          semilogx(f(idx),real(Puu(idx,ii,jj)), f(idx),imag(Puu(idx,ii,jj)))
          axis('tight')
        end
        if ii == 1,  title(sprintf('u_%d',jj)); end
        if jj == 1,  ylabel(sprintf('u_%d',ii)); end
        if ii == r,  xlabel('freq (Hz)'); end
     end
   end

  figure(figNo+1); % OUTPUT auto power spectra and cross spectra
   clf
   for ii = 1:m
     for jj = 1:m
       subplot(m,m,(ii-1)*m+jj)
        if ii == jj 
          semilogx(f(idx),Pyy(idx,ii,jj))%, f,CI(1)*Pyy,'-k', f,CI(2)*Pyy,'-k')
          axis('tight')
        else
          semilogx(f(idx),real(Pyy(idx,ii,jj)), f(idx),imag(Pyy(idx,ii,jj)))
          axis('tight')
        end
        if ii == 1,  title(sprintf('y_%d',jj)); end
        if jj == 1,  ylabel(sprintf('y_%d',ii)); end
        if ii == m,  xlabel('freq (Hz)'); end
     end
   end
  
  figure(figNo+2); % Frequency Response Functions
   clf
   for k=1:r
    subplot(3,r,k)
      semilogy(f(idx), abs(Hp(idx,:,k)) );
      axis('tight')
     if k == 1, ylabel('magnitude'); end
    subplot(3,r,k+r)
      plot( f(idx), 180/pi*pha(idx,:,k) );
      set (gca, 'ytick', [pha_min : 90 : pha_max ])
      axis([ ff pha_min pha_max ]);
     if k == 1, ylabel('phase (deg)'); end
    subplot(3,r,k+2*r)
      plot( f(idx), coh(idx,:,k));%, f(idx),sqrt(sum(coh(idx,:,k).^2,2)),'-k');
      axis('tight')
     if k == 1, ylabel('coherence'); end
     xlabel('frequency (Hertz)')
     grid on
   end

end

% figure(13)	% for testing Sv
%   plot(f,Sv./abs(Hp), [.1 50], [1 1] * sqrt( 0.5 * (CI(1) - CI(2)) ) / K )

% ------------------------------------------------------------------- mimoTFE
% H.P. Gavin 2017-12-03
