## -*- texinfo -*-
## @deftypefn {} {@var{pxx} =} pwelch (@var{x})
## @deftypefnx {} {@var{pxx} =} pwelch (@var{x}, @var{window})
## @deftypefnx {} {@var{pxx} =} pwelch (@var{x}, @var{window}, @var{noverlap})
## @deftypefnx {} {@var{pxx} =} pwelch (@var{x}, @var{window}, @var{noverlap}, @var{nfft})
## @deftypefnx {} {[@var{pxx}, @var{w}] =} pwelch (@dots{})
## @deftypefnx {} {[@var{pxx}, @var{f}] =} pwelch (@dots{}, @var{fs})
## @deftypefnx {} {[@var{pxx}, @var{w}] =} pwelch (@var{x}, @var{window}, @var{noverlap}, @var{w})
## @deftypefnx {} {[@var{pxx}, @var{f}] =} pwelch (@var{x}, @var{window}, @var{noverlap}, @var{f}, @var{fs})
## @deftypefnx {} {[@dots{}] =} pwelch (@dots{}, @var{freqrange})
## @deftypefnx {} {[@dots{}] =} pwelch (@dots{}, @var{spectrumtype})
## @deftypefnx {} {[@dots{}] =} pwelch (@dots{}, @var{trace})
## @deftypefnx {} {[@dots{}, @var{pxxc}] =} pwelch (@dots{}, "confidencelevel", @var{level})
## @deftypefnx {} {} pwelch (@dots{})
## Estimate power spectral density of @var{x} by the Welch (1967)
## periodogram/FFT method.
##
## All arguments except @var{x} are optional.  The data is divided into
## segments.  If @var{window} is a vector, each segment has the same
## length as @var{window} and is multiplied by @var{window} before
## (optional) zero-padding and calculation of its periodogram.  If
## @var{window} is a scalar, each segment has a length of @var{window}
## and a Hamming window is used.
##
## If @var{noverlap} is an integer scalar, it is the number of
## elements of overlap for each segment.  If it is missing or empty,
## then an overlap of 50% is used for each segment.
##
## If both @var{window} and @var{noverlap} are missing or empty, then
## @var{x} is divided into 8 segments that overlap by 50%.
##
## The spectral density is the mean of the periodograms, scaled so that
## area under the spectrum is the same as the mean square of the
## data.  This equivalence is supposed to be exact, but in practice there
## is a mismatch of up to 0.5% when comparing area under a periodogram
## with the mean square of the data.
##
## The optional argument @var{nfft} must be a non-negative integer
## value specifying the number of points to use when computing the FFT
## of each segment.  By default, @var{nfft} is computed to be the
## @code{max (256, 2 ^ nextpow2 (seg_len))}.
##
## If the fifth argument is a positive scalar it is taken to be the
## sampling frequency.  Otherwise a value of @math{2\pi} is used.
##
## FIXME: Need explanation of optional input (and output) arguments
## @var{w} and @var{f}.
##
## The optional argument @var{freqrange} may be one of the following:
##
## @table @qcode
## @item "onesided"
## The frequency range of the spectrum is zero up to but not including
## @code{fs/2}.  Power from negative frequencies is added to the
## positive side of the spectrum, but not at zero or Nyquist (@code{fs/2})
## frequencies.  This keeps power equal in time and spectral domains.
## See reference [2].
##
## @item "twosided"
## The frequency range of the spectrum is @code{-fs/2} to @code{fs/2},
## with negative frequencies stored in "wrap around" order after the positive
## frequencies; e.g. frequencies for a 10-point 'twosided'
## spectrum are 0 0.1 0.2 0.3 0.4 0.5 -0.4 -0.3 -0.2 -0.1
##
## @item "centered"
## The same as "twosided" but with the first half of the spectrum
## swapped with second half to put the zero-frequency value in the
## middle.  @seealso{fftshift}.
## @end table
##
## If @var{x} is real, the default for @var{freqrange} is
## @qcode{"onesided"}, otherwise the default @qcode{"twosided"}.
##
## The optional argument @var{spectrumtype} may be one of the
## following
##
## @table @qcode
## @item "psd"
## @item "power"
## @end @table
##
## The optional argument @var{trace} may be one of the following
##
## @table @qcode
## @item "minhold"
## @item "maxhold"
## @item "mean"
## @end @table
##
## If the optional name/value pair @code{"confidencelevel"} and a
## corresponding scalar value between 0 and 1 is specified, compute an
## estimate of the confidence interval of the spectral density of
## @var{x}.  The first column is the lower bound.  Second column is
## the upper bound.
##
## If no return values are requested, the results are plotted and
## no values are returned.
##
## NOTES
##
## @table @asis
## @item Empty arguments
## If you don't want to use an optional argument you can leave it empty
## by writing its value as [].
##
## @item For beginners
## The profusion of arguments may make pwelch difficult to use, and an
## unskilled user can easily produce a meaningless result or can easily
## mis-interpret the result.  With real data @var{x} and sampling frequency
## @var{Fs}, the easiest and best way for a beginner to use pwelch is
## probably @code{pwelch (x, [], [], [], fs)}.  Use the @var{window}
## argument to control the length of the spectrum vector.  For real
## data and integer scalar @var{m}, @code{pwelch (x, 2*m, [], [], fs)}
## gives an @code{m+1} point spectrum.  Run @qcode{"demo pwelch"}.
##
## @item Windowing functions
## Without a window function, sharp spectral peaks can have strong
## sidelobes because the FFT of a data in a segment is in effect convolved
## with a rectangular window.  A window function which tapers off
## (gradually) at the ends produces much weaker sidelobes in the FFT.
## Hann (hanning), hamming, bartlett, blackman, flattopwin etc are
## available as separate Matlab/sigproc or Octave functions.  The sidelobes
## of the Hann window have a roll-off rate of 60dB/decade of frequency.
## The first sidelobe of the Hamming window is suppressed and is about 12dB
## lower than the first Hann sidelobe, but the roll-off rate is only
## 20dB/decade.  You can inspect the FFT of a Hann window by plotting
## @code{abs (fft (postpad (hanning (256), 4096, 0)))}.
## The default window is Hamming.
##
## @item Zero padding
## Zero-padding reduces the frequency step in the
## spectrum, and produces an artificially smoothed spectrum.  For
## example, @code{nfft = 2 * length (window)} gives twice as many
## frequency values, but adjacent PSD (power spectral density) values
## are not independent; adjacent PSD values are independent if
## @code{nfft = length (window)}.
##
## @item Removing mean from signal
## If the mean is not removed from the signal there is a large spectral
## peak at zero frequency and the sidelobes of this peak are likely to
## swamp the rest of the spectrum.  For this reason, the default behavior
## is to remove the mean.  However, the matlab pwelch does not do this.
##
## @item Warning on confidence intervals
## Confidence intervals are obtained by measuring the sample variance of
## the periodograms and assuming that the periodograms have a Gaussian
## probability distribution.  This assumption is not accurate.  If, for
## example, the data @var{x} is Gaussian, the periodogram has a Rayleigh
## distribution.  However, the confidence intervals may still be  useful.
##
## REFERENCES
##
## [1] Peter D. Welch (June 1967):  "The use of fast Fourier transform
## for the estimation of power spectra: a method based on time
## averaging over short, modified periodograms."  IEEE Transactions on
## Audio Electroacoustics, Vol AU-15(6), pp 70-73.
##
## [2] William H. Press and Saul A. Teukolsky and William
## T. Vetterling and Brian P. Flannery", "Numerical recipes in C, The
## art of scientific computing", 2nd edition, Cambridge University
## Press, 2002, Section 13.7.
## @end deftypefn

function [pxx, varargout] = pwelch (x, varargin)

  if (nargin < 1)
    print_usage ();
  endif

  ## Force x to be column vector.
  if (isvector (x) && rows (x) == 1)
    x = x(:);
  endif

  x_dims = size (x);

  [window, noverlap, nfft, fs, f, w, info] = parse_args (x, varargin{:});

  x_len = info.x_len;
  seg_len = info.seg_len;

  ## Calculate and accumulate periodograms.

  ## Periodogram sums.
  pxx_dims = x_dims;
  pxx_val = zeros ([nfft, prod(pxx_dims(2:end))]);
  vxx_val = pxx_val;

  n_ffts = 0;

  for start_seg = 1:seg_len-noverlap:x_len-seg_len+1
    end_seg = start_seg+seg_len-1;

    xx = window .* x(start_seg:end_seg,:);

    fft_x = fft (xx, nfft);

    ## Force Pxx to be real; pgram = periodogram.
    pgram = real (fft_x .* conj (fft_x));

    ## maxhold and minhold preserve the maxmimum and minimum values of
    ## each bin from all periodograms.

    if (info.have_trace_arg)
      if (strcmp (info.trace, "maxhold"))
        pxx_val = max (pxx_val, pgram);
      elseif (strcmp (info.trace, "minhold"))
        if (start_seg == 1)
          ## pxx_val is initialised to zeros, so for the first iteration,
          ## should ignore it.
          pxx_val = pgram;
        else
          pxx_val = min (pxx_val, pgram);
        endif
      elseif (strcmp (info.trace, "mean"))
        pxx_val = pxx_val + pgram;
      else
        error ("pwelch: internal error");
      endif
    else
      ## If no trace agruments, default is mean.
      pxx_val = pxx_val + pgram;
    endif

    ## Need the average periodogram regardless of the trace.  Tests
    ## seems to indicate that the confidence levels are same for all
    ## traces in Matlab.
    vxx_val = vxx_val + pgram;

    n_ffts = n_ffts + 1;
  endfor

  if (info.have_confidencelevel_arg)
    ## Confidence interval is estimated using Chi-Squared method,
    ## using chi2inv from the statistics package.  We will just fail
    ## if it is not available.

    ## It seems like the confidence values are the same regardless of
    ## the trace argument.  So we need to store the mean pgram after all.
    ## Also, previous condition returned confidence as 0 for n_ffts = 1,
    ## not Matlab behavior.  Fixing that here.

    dof = 2 * n_ffts;
    signific_level = 1 - info.confidencelevel;
    conf_scale = chi2inv ([1-signific_level/2, signific_level/2], dof);
    conf_scale = dof ./ conf_scale;
  endif

  ## When "power" is specified as the spectrum type, the PSD has to be
  ## scaled with Equivalent Noise BandWidth (ENBW).  When fs is given,
  ## the ENBW is multiplied by the frequency resolution.  Reference:
  ##
  ##   https://www.gaussianwaves.com/2011/02/window-functions-an-analysis/
  ##
  ## Follow that with "The Fundamentals of FFT-Based Signal Analysis and
  ## Measurement".

  if (info.have_spectrumtype_arg && strcmp (info.spectrumtype, "power"))
    enbw = mean (window .^ 2) ./ mean (window) .^ 2 * fs / seg_len;
    pxx_val = pxx_val .* enbw;
    vxx_val = vxx_val .* enbw;
  endif

  ## Convert two-sided spectra to one-sided spectra (if range == 0).
  ## For one-sided spectra, contributions from negative frequencies
  ## are added to the positive side of the spectrum -- but not at zero
  ## or Nyquist (half sampling) frequencies.  This keeps power equal
  ## in time and spectral domains, as required by Parseval theorem.

  if (strcmp (info.freqrange, "onesided"))
    if (! rem (nfft, 2))    # one-sided, nfft is even
      psd_len = nfft / 2 + 1;
      ## STPR: Changed behavior here. Concatenation with multi-dimensional
      ## array can get messy, so I am going to split it into two steps.
      ## First add the respective frequencies, then index out the positive
      ## frequencies.

      pxx_val(2:psd_len-1,:) = (pxx_val(2:psd_len-1,:)
                                + pxx_val(nfft:-1:psd_len+1,:));
      pxx_val = pxx_val(1:psd_len,:);

      vxx_val(2:psd_len-1,:) = (vxx_val(2:psd_len-1,:)
                                + vxx_val(nfft:-1:psd_len+1,:));
      vxx_val = vxx_val(1:psd_len,:);
    else                    # one-sided, nfft is odd
      psd_len = (nfft + 1) / 2;
      pxx_val(2:psd_len,:) = (pxx_val(2:psd_len,:)
                              + pxx_val(nfft:-1:psd_len+1,:));
      pxx_val = pxx_val(1:psd_len,:);

      vxx_val(2:psd_len,:) = (vxx_val(2:psd_len,:)
                              + vxx_val(nfft:-1:psd_len+1,:));
      vxx_val = vxx_val(1:psd_len,:);
    endif
  else                      # two-sided (and shifted)
    psd_len = nfft;
  endif

  ## Scaling and output.

  ## Mean square of window is required for normalizing PSD amplitude.
  win_meansq = (window.' * window) / seg_len;
  scale = n_ffts * seg_len * fs * win_meansq;
  spectra = pxx_val / scale;
  vxx_val = vxx_val / scale;
  if (info.have_trace_arg
      && (strcmp (info.trace, "maxhold") || strcmp (info.trace, "minhold")))
    ## For this case, there is no averaging since we are just holding
    ## a single value.  Remove the effect of n_ffts;
    spectra = spectra * n_ffts;
  endif

  final_pxx_dims = x_dims;
  final_pxx_dims(1) = psd_len;

  freq = (0:psd_len-1).' * (fs / nfft);

  ## range='shift': Shift zero-frequency to the middle

  if (strcmp (info.freqrange, "centered"))
    len2 = ceil (nfft / 2) + 1;
    spectra = [spectra(len2+1:nfft,:); spectra(1:len2,:)];
    vxx_val = [vxx_val(len2+1:nfft,:); vxx_val(1:len2,:)];
    freq = [freq(len2+1:nfft)-fs; freq(1:len2)];
  endif
  ## There are twice as many columns containing the confidence levels.
  ## Following that with the reference above, we end up with the
  ## following.
  if (info.have_confidencelevel_arg)
    Vxx_dims = size (spectra);
    Vxx_dims(2) = 2 * Vxx_dims(2);
    Vxx = zeros (Vxx_dims);
    Vxx(:,1:2:end-1) = spectra * conf_scale(1);
    Vxx(:,2:2:end) = spectra * conf_scale(2);
    final_Vxx_dims = final_pxx_dims;
    final_Vxx_dims(2) = 2 * final_Vxx_dims(2);
    Vxx = reshape (Vxx, final_Vxx_dims);
  endif
  if (nargout >= 1)
    pxx = reshape (spectra, final_pxx_dims);
    iargout = 1;
    if (info.have_f_arg || info.have_fs_arg)
      varargout{iargout++} = freq;
    elseif (info.have_w_arg)
      varargout{iargout++} = w;
    elseif (nargout >= 2)
    info
      if (strcmp (info.freqrange, "onesided"))
        nfreq = pi * freq / max (freq);
      else
        nfreq = 2 * pi * freq / max (freq);
      endif
      varargout{iargout++} = nfreq;
    endif
    if (info.have_confidencelevel_arg)
      varargout{iargout} = Vxx;
    endif
  else
    y1 = 10 * log10 (abs (spectra));
    y2 = [];
    if (info.have_confidencelevel_arg)
      y2 = 10 * log10 (abs (Vxx));
    endif
    nfreq = freq / max (freq);
    plot (nfreq, [y1, y2]);
    title ("Welch PSD Estimate");
    ylabel ("amplitude (dB)");
    xlabel ("normalized frequency (\\pi rad/sample");
    grid ("on");
  endif

endfunction

## The rules for determining argument meanings and default parameter
## and option values for PWELCH are as complicated as just about any
## function I've ever encountered.

function [window, noverlap, nfft, fs, f, w, info] = parse_args (x, varargin)

  x_len = length (x);

  ## Find the first option argument in varargin.

  nvargs = numel (varargin);
  non_option_args = nargin;
  for i = 1:nvargs
    if (ischar (varargin{i}));
      ## Including X.
      non_option_args = i;
      break;
    endif
  endfor

  if (non_option_args > 5)
    evalin ("caller", "print_usage");
  endif

  info.have_fs_arg = false;
  info.have_f_arg = false;
  info.have_w_arg = false;

  ## FS must be the fifth argument.
  if (non_option_args == 5)
    fs = varargin{4};
    if (! (isscalar (fs) && fs > 0))
      error ("pwelch: FS must be a scalar greater than 0");
    endif
    info.have_fs_arg = true;
  else
    fs = 2 * pi;
  endif

  if (non_option_args > 1)
    window = varargin{1};
  else
    window = [];
  endif

  if (! (isempty (window)
         || (isscalar (window) && window > 0 && round (window) == window)
         || isvector (window)))
    error ("pwelch: WINDOW must be a vector or or an integer greater than 0");
  endif

  if (non_option_args > 2)
    noverlap = varargin{2};
  else
    noverlap = [];
  endif

  if (! (isempty (noverlap)
         || (isscalar (noverlap) && noverlap >= 0
             && round (noverlap) == noverlap)))
    error ("pwelch: NOVERLAP must be an integer greater than or equal to 0");
  endif

  ## At this point, WINDOW is empty, a positive scalar, or a vector,
  ## and NOVERLAP is empty, or a non-negative scalar.

  if (isempty (window))
    if (isempty (noverlap))
      ## 8 segments, 50% overlap.
      ## STPR: Why was x_len reduced (subtracted) by 3 in the previous code?
      ##  Also note: spent wayy too much time on this, according to
      ##  https://edoras.sdsu.edu/doc/matlab/toolbox/signal/pwelch.html,
      ##  the segment length will be the following:
      window = fix ((x_len) * 2 / 9);
    else
      ## Choose window such that there are 8 segments that overlap by
      ## noverlap elements.
      window = fix ((x_len + 7 * noverlap) / 8);
    endif
  endif

  if (isscalar (window))
    seg_len = window;
  elseif (isvector (window))
    seg_len = numel (window);
  else
    error ("pwelch: internal error");
  endif

  if (isempty (noverlap))
    noverlap = fix (seg_len / 2);
  endif

  if (noverlap >= seg_len)
    error ("pwelch: NOVERLAP must be less than window length");
  endif

  if (isscalar (window))
    ## Make Hamming window.
    xx = seg_len - 1;
    window = 0.54 - 0.46 * cos (2 * pi * (0 : xx)' / xx);
  endif

  nfft = [];
  f = NaN (2, 1);
  w = NaN (2, 1);
  if (non_option_args > 3)
    argval = varargin{3};
    if (isempty (argval))
      ## Use default value for NFFT.
    elseif (isscalar (argval))
      ## Scalar value must be NFFT.
      nfft = argval;
      if (! (nfft > 0 && round (nfft) == nfft))
        error ("pwelch: NFFT must be an integer greater than 0");
      endif
    elseif (isvector (argval))
      if (info.have_fs_arg)
        f = argval;
        info.have_f_arg = true;
      else
        w = argval;
        info.have_w_arg = true;
      endif
    else
      error ("pwelch: invalid value for NFFT, F, or W")
    endif
  endif

  if (non_option_args < 4 || isempty (nfft))
    nfft = max (256, 2 ^ nextpow2 (seg_len));
  endif

  info.have_freqrange_arg = false;
  info.have_spectrumtype_arg = false;
  info.have_trace_arg = false;
  info.have_confidencelevel_arg = false;

  if (iscomplex (x))
    info.freqrange = "twosided";
  else
    info.freqrange = "onesided";
  endif
  info.spectrumtype = "psd";
  info.trace = "mean";
  info.confidencelevel = 0.95;

  vargidx = non_option_args;
  while (vargidx <= nvargs)
    argval = varargin{vargidx};
    switch (tolower (argval))
      case {"onesided", "twosided", "centered"}
        if (info.have_freqrange_arg)
          error ("pwelch: multiple freqrange options specified");
        endif
        info.have_freqrange_arg = true;
        info.freqrange = argval;

      case {"psd", "power"}
        if (info.have_spectrumtype_arg)
          error ("pwelch: multiple spectrumtype options specified");
        endif
        info.have_spectrumtype_arg = true;
        info.spectrumtype = argval;

      case {"minhold", "maxhold", "mean"}
        if (info.have_trace_arg)
          error ("pwelch: multiple trace options specified");
        endif
        info.have_trace_arg = true;
        info.trace = argval;

      case "confidencelevel"
        if (info.have_confidencelevel_arg)
          error ("pwelch: multiple confidencelevel options specified");
        endif
        info.have_confidencelevel_arg = true;
        vargidx++;
        if (vargidx <= nvargs)
          info.confidencelevel = varargin{vargidx};
          if (! (isscalar (info.confidencelevel)
                 && info.confidencelevel > 0
                 && info.confidencelevel < 1))
            error ("pwelch: confidencelevel must be in the range (0,1)");
          endif
        else
          error ("pwelch: expecting value to follow 'confidencelevel'");
        endif

      otherwise
        error ("pwelch: unrecognized option argument");
    endswitch
    vargidx++;
  endwhile

  info.x_len = x_len;
  info.seg_len = seg_len;

endfunction

%!demo
%! a = [ 1.0 -1.6216505 1.1102795 -0.4621741 0.2075552 -0.018756746 ];
%! white = rand(1,16384);
%! signal = detrend(filter(0.70181,a,white));
%! % frequency shift by modulating with exp(j.omega.t)
%! skewed = signal.*exp(2*pi*i*2/25*[1:16384]);
%! hold on;
%! pwelch(signal);
%! pwelch(skewed);
%! hold off;

%!demo
%! Fs = 25;
%! a = [ 1.0 -1.6216505 1.1102795 -0.4621741 0.2075552 -0.018756746 ];
%! white = rand(1,16384);
%! signal = detrend(filter(0.70181,a,white));
%! % frequency shift by modulating with exp(j.omega.t)
%! skewed = signal.*exp(2*pi*i*2/25*[1:16384]);
%! pwelch(skewed,[],[],[],Fs);
%! pwelch(skewed,[],[],[],Fs,'confidencelevel',0.95);
%! pwelch(signal);

%!demo
%! a = [ 1.0 -1.6216505 1.1102795 -0.4621741 0.2075552 -0.018756746 ];
%! white = rand(1,16384);
%! signal = detrend(filter(0.70181,a,white));
%! pwelch(signal,3640,[],4096,2*pi);

%!demo
%! a = [ 1.0 -1.6216505 1.1102795 -0.4621741 0.2075552 -0.018756746 ];
%! white = rand(1,16384);
%! signal = detrend(filter(0.70181,a,white));
%! hold on;
%! pwelch(signal,[],[],[],2*pi,'confidencelevel',0.95);
%! pwelch(signal,64,[],[],2*pi);
%! pwelch(signal,64,[],256,2*pi);
%! hold off;
