function [Hv, f, Puu, Pyy, Puy, coh, Sv] = tfe(u,y,Fs,nfft)
% [Hv, f, Puu, Pyy, Puy, coh, Sv] = tfe(u,y,Fs,nfft)
%
% Transfer function estimate from discrete time data using Welch's averaged 
% periodogram method, in which the vectors u and y are divided into overlapping
% segments, each of which is detrended, then windowed by the Hann window,
% then zero-padded to length NFFT.  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 Puy.  Hv is the
% geometric mean of the H1 and H2 estimates where H1=Puy/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      vector of sampled  input to the system               (1 x points)
%   y      vector of sampled output to the system               (1 x points)
%   Fs     sampling frequency ... 1 / delta_t                   (1 x 1)
%  nfft    the number of points used in the fast Fourier transform (1x1)
%
%  OUTPUT
%   Hv     the complex-valued Hv transfer function estimate     (1 x nfft/2)
%   f      the frequency values  corresponding to H             (1 x nfft/2)
%   Puu    the real-valued auto-power spectrum of u             (1 x nfft/2)
%   Pyy    the real-valued auto-power spectrum of y             (1 x nfft/2)
%   Puy    the complex-valued cross-power spectrum from u to y  (1 x nfft/2)
%   coh    the coherence function                               (1 x nfft/2)
%   Sv     the standard deviation of the transfer fctn estimate (1 x nfft/2)
%
% TFE with no output arguments plots the transfer function estimate. 
%

error(nargchk(2,4,nargin))
if (length(u)~=length(y)),
    error('U and Y must be the same length.')
end

% Defaults 

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

u = u(:);		% Make sure u is a column vector
y = y(:);		% Make sure y is a column vector
n = length(u);		% Number of data points

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(nWindow,1);					% 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 PSDs and CSD from FFTs

Puu  = zeros(nfft,1);		Puu2 = zeros(nfft,1); 
Pyy  = zeros(nfft,1);		Pyy2 = zeros(nfft,1); 
Puy  = zeros(nfft,1);		Puy2 = zeros(nfft,1); 

for k=1:K

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

    Uu = fft(uw);
    Yy = fft(yw);
    Uu2 = abs(Uu).^2;
    Yy2 = abs(Yy).^2;
    Uy2 = conj(Uu).*Yy;

    Puu = Puu + Uu2;		% Puu2 = Puu2 + Uu2.*Uu2;	% for testing Sv
    Pyy = Pyy + Yy2;		% Pyy2 = Pyy2 + Yy2.*Yy2;	% for testing Sv
    Puy = Puy + Uy2;		% Puy2 = Puy2 + Uu2.*Yy2;	% for testing Sv

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

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 );
Puy = Puy / ( KNW2 * Fs );

H1 = Puy ./ Puu;		% H1 transfer function estimate 
H2 = Pyy ./ conj(Puy);		% H2 transfer function estimate
Hv = sqrt(H1) .* sqrt(H2);	% Hv transfer function estimate
coh = real ( H1 ./ H2 );	% coherence of the transfer 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

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

if nargout == 0 || nargout == 7 
  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 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 = (Puy2 - KNW2*Fs * abs(Puy).^2) / ((K-1)*K);
%%  
%%   Sv = 0.5 * abs(Hv) .* sqrt ( Suu2./Puu2 + 2*Suy2./Puy2 + Syy2./Pyy2 );

% approx Sv estimate from Chi-Squared distribution, and 50% window overlap
% from Kay, p. 76, eqn 4.16:
  Sv = abs(Hv) * sqrt( 0.5 * (CI(1) - CI(2)) ) / K;
end

if nargout == 0   % make plots

   figure(11);
   clf
   subplot(311)
    loglog(f,abs([H1 H2 Hv]), f,abs(Hv)+Sv,'-m', f, abs(Hv)-Sv,'-m')
     legend('H1', 'H2', 'Hv')
     ylabel('magnitude');
     axis([ f(3) max(f) min(abs(H1)), max(abs(H2)) ])
     grid on
   subplot(312)
    semilogx(f, 180/pi*(angle([H1 H2 Hv])))
     legend('H1', 'H2', 'Hv')
     ylabel('phase, degrees');
     axis([ f(3) max(f) -190 190 ])
     grid on
   subplot(313)
    semilogx(f,coh)
     xlabel('Frequency'), ylabel('coherence');
     axis([ f(3) max(f) 0 1.05 ])
     grid on

   figure(12);
   clf
   subplot(311)
%   semilogx(f,Puu, f,Puu+sqrt(Suu2),'-k') % for testing Sv
    semilogx(f,Puu, f,CI(1)*Puu,'-k', f,CI(2)*Puu,'-k')
     ylabel('Puu')
     axis([ f(3) 0.2*max(f) min(Puu), max(Puu) ])
   subplot(312)
%   semilogx(f,Pyy, f,Pyy+sqrt(Syy2),'-k') % for testing Sv
    semilogx(f,Pyy, f,CI(1)*Pyy,'-k', f,CI(2)*Pyy,'-k')
     ylabel('Pyy')
     axis([ f(3) 0.2*max(f) min(Pyy), max(Pyy) ])
   subplot(313)
    semilogx(f,real([H1 H2 Hv]), f,imag([H1 H2 Hv]))
     ylabel('real(H) and imag(H)');
     axis([ f(3) 0.2*max(f)])
     grid on

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

end

% ----------------------------------------------------------------------- TFE
