% WELCH Welch spectral estimation method.
%
% [pxx, f, info] = welch(x,type,pl)
% [pxx, f, info] = welch(x,y,type,pl)
%
% INPUTS:      x    - input analysis objects
%              y    - input analysis objects
%              type - type of spectrum:
%                       'psd'      - compute Power Spectral Denstiy (PSD)
%                       'ms'       - compute Mean-square (Power) Spectrum (MS)
%                       'cpsd'     - compute cross-spectral density
%                       'tfe'      - estimate transfer function between inputs
%                       'mscohere' - estimate magnitude-squared cross-coherence
%                       'cohere'   - estimate complex cross-coherence
%              pl   - input parameter list
%
% PARAMETERS: 'Win'   - a specwin window object [default: Kaiser -200dB psll]
%             'Olap' - segment percent overlap [default: taken from window function]
%             'Nfft'  - number of samples in each fft [default: length of input data]
%             'Scale' - one of
%                                'ASD' - amplitude spectral density
%                                'PSD' - power spectral density [default]
%                                'AS'  - amplitude spectrum
%                                'PS'  - power spectrum
%                       * applies only to spectrum 'Type' 'psd'
%             'Order' - order of segment detrending
%                        -1 - no detrending
%                         0 - subtract mean [default]
%                         1 - subtract linear fit
%                         N - subtract fit of polynomial, order N
%
% Copied directly from MATLAB and extended to do segment-wise detrending,
% compute transfer function variance and to take a plist input
%
% M Hewitson 08-05-08
%
% $Id: welch.m,v 1.34 2011/03/14 16:06:05 mauro Exp $
%

%   Author(s): P. Pacheco
%   Copyright 1988-2006 The MathWorks, Inc.
%   $Revision: 1.34 $  $Date: 2011/03/14 16:06:05 $

%   References:
%     [1] Petre Stoica and Randolph Moses, Introduction To Spectral
%         Analysis, Prentice-Hall, 1997, pg. 15
%     [2] Monson Hayes, Statistical Digital Signal Processing and
%         Modeling, John Wiley & Sons, 1996.
%     [3] JS Bendat and AG Piersol, Engineering applications of correlation
%         and spectral analysis, John Wiley & Sons, 1993.


% Compute the periodogram power spectrum of each segment and average always
% compute the whole power spectrum, we force Fs = 1 to get a PS not a PSD.

function varargout = welch(varargin)
  
  if nargin == 3
    a  = varargin{1};
    esttype = varargin{2};
    pl = varargin{3};
    x  = a.data.y;
    inunits = a.data.yunits;
  else
    a  = varargin{1};
    b  = varargin{2};
    esttype = varargin{3};
    pl = varargin{4};
    if a.data.fs ~= b.data.fs
      error('### Two time-series have different sample rates.');
    end
    inunits = b.data.yunits / a.data.yunits;
    x  = {a.data.y, b.data.y};
  end
  
  % Parse inputs
  win     = find(pl, 'Win');
  nfft    = find(pl, 'Nfft');
  olap    = find(pl, 'Olap')/100;
  scale   = find(pl, 'scale');
  Xolap   = round(olap*nfft);
  fs      = a.data.fs;
  order   = find(pl, 'order');
  
  [x,M,isreal_x,y,Ly,win,winName,winParam,noverlap,k,L,options] = ...
    ao.welchparse(x,esttype,win.win, Xolap, nfft, fs);
  
  % Initialize
  Sxx = zeros(options.nfft,1);
  
  % Frequency vector was specified, return and plot two-sided PSD
  freqVectorSpecified = false; nrow = 1;
  if length(options.nfft) > 1,
    freqVectorSpecified = true;
    [ncol,nrow] = size(options.nfft);
  end
  
  % Compute the periodogram power spectrum of each segment and average always
  % compute the whole power spectrum, we force Fs = 1 to get a PS not a PSD.
  
  % Initialize
  Mnxx = 0; Mn2xx = 0;
  if freqVectorSpecified,
    Sxx = zeros(length(options.nfft),1);
  else
    Sxx = zeros(options.nfft,1);
  end
  range = options.range;
  
  LminusOverlap = L-noverlap;
  xStart = 1:LminusOverlap:k*LminusOverlap;
  xEnd   = xStart+L-1;
  switch lower(esttype)
    case {'ms','psd'}
      for ii = 1:k,
        if order < 0
          seg = x(xStart(ii):xEnd(ii));
        else
          [seg,coeffs] = ltpda_polyreg(x(xStart(ii):xEnd(ii)), order);
        end
        [Sxxk,w] = ao.computeperiodogram(seg,win,...
          options.nfft,esttype,options.Fs);
        % Welford's algorithm for updating mean and variance
        if ii == 1
          Mnxx = Sxxk;
        else
          Qxx = Sxxk - Mnxx;
          Mnxx = Mnxx + Qxx/ii;
          Mn2xx = Mn2xx + Qxx.*(Sxxk - Mnxx);
        end
      end
      Sxx = Mnxx;
      if k == 1
        Svxx = [];
      else
        Svxx = Mn2xx/(k-1)/k;
      end
    case 'cpsd'
      for ii = 1:k,
        if order < 0
          xseg = x(xStart(ii):xEnd(ii));
        else
          [xseg,coeffs] = ltpda_polyreg(x(xStart(ii):xEnd(ii)), order);
        end
        if order < 0
          yseg = y(xStart(ii):xEnd(ii));
        else
          [yseg,coeffs] = ltpda_polyreg(y(xStart(ii):xEnd(ii)), order);
        end
        [Sxxk,w] =  ao.computeperiodogram({xseg,...
          yseg},win,options.nfft,esttype,options.Fs);
        % Welford's algorithm to update mean and variance
        Qxx = Sxxk - Mnxx;
        Mnxx = Mnxx +Qxx/ii;
        Mn2xx = Mn2xx + abs(Qxx.*conj(Sxxk - Mnxx));
      end
      Sxx = Mnxx;
      if k ==1
        Svxx = [];
      else
        Svxx = Mn2xx/(k-1)/k;
      end
    case 'tfe'
      % compute transfer function
      Sxy = zeros(options.nfft,1); % Initialize
      Syy = zeros(options.nfft,1); % Initialize
      for ii = 1:k,
        if order < 0
          xseg = x(xStart(ii):xEnd(ii));
        else
          [xseg,coeffs] = ltpda_polyreg(x(xStart(ii):xEnd(ii)), order);
        end
        if order < 0
          yseg = y(xStart(ii):xEnd(ii));
        else
          [yseg,coeffs] = ltpda_polyreg(y(xStart(ii):xEnd(ii)), order);
        end
        [Sxxk,w] = ao.computeperiodogram(xseg,...
          win,options.nfft,esttype,options.Fs);
        [Sxyk,w] = ao.computeperiodogram({yseg,...
          xseg},win,options.nfft,esttype,options.Fs);
        [Syyk,w] =  ao.computeperiodogram(yseg,...
          win,options.nfft,esttype,options.Fs);
        Sxx = Sxx + Sxxk;
        Sxy = Sxy + Sxyk;
        Syy = Syy + Syyk;
        % don't need to be divided by k because only rations are used here
      end
    case {'mscohere','cohere'}
      % Note: (Sxy1+Sxy2)/(Sxx1+Sxx2) ~= (Sxy1/Sxy2) + (Sxx1/Sxx2)
      % ie, we can't push the computation of Cxy into computeperiodogram.
      Sxy = zeros(options.nfft,1); % Initialize
      Syy = zeros(options.nfft,1); % Initialize
      for ii = 1:k,
        if order < 0
          xseg = x(xStart(ii):xEnd(ii));
        else
          [xseg,coeffs] = ltpda_polyreg(x(xStart(ii):xEnd(ii)), order);
        end
        if order < 0
          yseg = y(xStart(ii):xEnd(ii));
        else
          [yseg,coeffs] = ltpda_polyreg(y(xStart(ii):xEnd(ii)), order);
        end
        [Sxxk,w] = ao.computeperiodogram(xseg,...
          win,options.nfft,esttype,options.Fs);
        [Syyk,w] =  ao.computeperiodogram(yseg,...
          win,options.nfft,esttype,options.Fs);
        [Sxyk,w] = ao.computeperiodogram({xseg,...
          yseg},win,options.nfft,esttype,options.Fs);
        Sxx = Sxx + Sxxk;
        Sxy = Sxy + Sxyk;
        Syy = Syy + Syyk;
        % don't need to be divided by k because only rations are used here
      end
  end
  % Generate the freq vector directly in Hz to avoid roundoff errors due to
  % conversions later.
  if ~freqVectorSpecified,
    w = psdfreqvec('npts',options.nfft, 'Fs',options.Fs);
  else
    if strcmpi(options.range,'onesided')
      warning(generatemsgid('InconsistentRangeOption'),...
        'Ignoring ''onesided'' option. When a frequency vector is specified, a ''twosided'' PSD is computed');
    end
    options.range = 'twosided';
  end
  
  switch lower(esttype)
    case {'psd','cpsd'}
      % Compute the 1-sided or 2-sided PSD [Power/freq] or mean-square [Power].
      % Also, corresponding freq vector and freq units.
      % Here we use our own 'computepsd' to scale correctly the variance
      if k == 1
        [Pxx,w,units] = computepsd(Sxx,w,options.range,options.nfft,options.Fs,esttype);
        P = Pxx;
        dP = [];
      else
        [Pxx,Pvxx,w,units] = utils.math.computepsd(Sxx,Svxx,w,options.range,options.nfft,options.Fs,esttype);
        P = Pxx;
        % the 1/k factor should come after welchscale if we don't want to apply sqrt() to it.
        % we correct for that here. It is only needed for 'asd','as' in
        % psd/cpsd, the other cases go always through 'PSD'.
        if (strcmpi(scale,'PSD') || strcmpi(scale,'PS'))
          dP = Pvxx;
        elseif (strcmpi(scale,'ASD') || strcmpi(scale,'AS'))
          dP = Pvxx/k;
        else
          error('### Unknown scale')
        end
      end
    case 'tfe'
      % Compute the 1-sided or 2-sided PSD [Power/freq] or mean-square [Power].
      % Also, corresponding freq vector and freq units.
      % In the Cross PSD, the frequency vector and xunits are not used.
      [Pxx,w,units] = computepsd(Sxx,w,options.range,options.nfft,options.Fs,esttype);
      [Pxy,w,units] = computepsd(Sxy,w,options.range,options.nfft,options.Fs,esttype);
      [Pyy,w,units] = computepsd(Syy,w,options.range,options.nfft,options.Fs,esttype);
      % mean and std
      P = Pxy ./ Pxx; % Txy
      if k == 1
        dP =[];
      else
        dP = (k/(k-1)^2)*(Pyy./Pxx).*(1 - (abs(Pxy).^2)./(Pxx.*Pyy));
      end
    case 'mscohere'
      % Magnitude Square Coherence estimate.
      % Auto PSD for 2nd input vector. The freq vector & xunits are not
      % used.
      [Pxx,w,units] = computepsd(Sxx,w,options.range,options.nfft,options.Fs,esttype);
      [Pxy,w,units] = computepsd(Sxy,w,options.range,options.nfft,options.Fs,esttype);
      [Pyy,w,units] = computepsd(Syy,w,options.range,options.nfft,options.Fs,esttype);
      % mean and std
      P = (abs(Pxy).^2)./(Pxx.*Pyy); % Magnitude-squared coherence
      dP = (2*P/k).*(1-P).^2;
    case 'cohere'
      % Complex Coherence estimate.
      % Auto PSD for 2nd input vector. The freq vector & xunits are not
      % used.
      [Pxx,w,units] = computepsd(Sxx,w,options.range,options.nfft,options.Fs,esttype);
      [Pxy,w,units] = computepsd(Sxy,w,options.range,options.nfft,options.Fs,esttype);
      [Pyy,w,units] = computepsd(Syy,w,options.range,options.nfft,options.Fs,esttype);
      P = Pxy./sqrt(Pxx.*Pyy); % Complex coherence
      dP = (2*abs(P)/k).*(1-abs(P)).^2;
  end
  %   end
  
  % Scale to required units
  [P, dP, info] = ao.welchscale(P, dP, win, fs, scale, inunits);
  info.navs = k;
  
  if k ==1
    dev = [];
  else
    dev = sqrt(dP);
  end
  
  varargout = {P, w, info, dev};
  
end