## -*- 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.
##
## FIXME: need explanation of the parameters, their default values,
## and exactly what is computed.
## @end deftypefn

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

  if (nargin < 1)
    print_usage ();
  endif

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

  [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.

  ## Padded data segments.
  ## STPR: Account for multidimensional array.
  x_size = size(x);
  xx = zeros ([nfft, x_size(2:end)]);

  ## Periodogram sums.
  pxx_val = xx;

  n_ffts = 0;
  
  ## STPR: MATLAB pwelch works for 2D matrices as well. 
  ##  However, this can be extended to an nD matrix quite easily. 
  ##  Should we go for that or limit to 2D?
  ##  Since we do not know ndims in advance, I am going to generate a list of
  ##  colons to deal with that issue. 
  
  colons = repmat({':'},1,ndims(x)-1);
  for start_seg = 1:seg_len-noverlap:x_len-seg_len+1
    ## Don't truncate/remove the zero padding in xx.
    end_seg = start_seg+seg_len-1;
    xx(1:seg_len, colons{:}) =window.*x(start_seg:end_seg, colons{:});
    fft_x = fft (xx);

    ## Force Pxx to be real; pgram = periodogram.
    pgram = real (fft_x .* conj (fft_x));
    
    ## STPR: Trace argument was never used. According to MATLAB documentation,
    ##  maxhold and minhold just 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;
        endif
        pxx_val = min(pxx_val, pgram);
      else
        ## STPR: Args must have been validated, but just in case.
        error ("pwelch: internal error");
      endif      
    else
      ## STPR: If no trace agruments, default is mean.
      pxx_val = pxx_val + pgram;         
    endif   
    n_ffts = n_ffts + 1;
  endfor

  if (info.have_confidencelevel_arg)
    if (n_ffts < 2)
      Vxx = zeros (nfft, 1);
    else
      ## FIXME: the following note doesn't really matter for Octave,
      ## does it?
      ## STPR: The comment here has a typo. Functions to compute the confidence 
      ##  intervals are in the statistics package of OCTAVE.
    
      ## STPR: Confidence interval is estimated using Chi-Squared method, NOT
      ##  t-test. However, chi2inv is part of the statistics package. There does
      ##  not appear to be any simple approximation methods. So, I am going to
      ##  check to see if statistics package is available. If it is not found, 
      ##  throw an error.
      ##  Simple Ref: https://rhenanbartels.wordpress.com/2014/04/06/first-post/
      ##  The values appear to agree with MATLAB from some initial tests.
      ##  FIXME: how does this work with the other trace args? Needs testing.
      [~, stat_flag] = pkg("describe", "statistics");
      if (~strcmp(stat_flag, "Loaded"))
        error("Statistics package required to calculate confidence interval. Please install/load the package and try again.")
      endif
      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
  endif
   
  ## STPR: I noticed both old and current versions of Octave do not appear to 
  ##  use the "power" argument at all. 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.
  ##  Ref: 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;
  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 behaviour 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, colons{:}) = pxx_val(2:psd_len-1, colons{:}) + ...
           pxx_val(nfft:-1:psd_len+1, colons{:});
      pxx_val = pxx_val(1:psd_len, colons{:});    

    else                    # one-sided, nfft is odd
      psd_len = (nfft + 1) / 2;
      pxx_val(2:psd_len, colons{:}) = pxx_val(2:psd_len, colons{:}) + ...
           pxx_val(nfft:-1:psd_len+1, colons{:});
      pxx_val = pxx_val(1:psd_len, colons{:});       
    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;
  if (info.have_trace_arg)
    if (strcmp(info.trace, "maxhold") || strcmp(info.trace, "minhold"))
      ## STPR: For this case, there is no averaging since we are just holding 
      ##  a single value.
      scale = seg_len * fs * win_meansq;
    else
      scale = n_ffts * seg_len * fs * win_meansq;
    endif    
  endif  

  spectra = pxx_val / scale;


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

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

  if (strcmp (info.freqrange, "centered"))
    ## STPR: Previous method incorrectly shifted Nyquist to the beginning. Fixing.    
    len2 = ceil(nfft/2)+1;
    ## STPR: Nfft -> nfft. Also account for nD array. 
    spectra = [spectra(len2+1:nfft, colons{:}); spectra(1:len2, colons{:})];
    freq = [freq(len2+1:nfft)-fs; freq(1:len2)];
  endif
  ## STPR: So according to MATLAB documentation, 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)
    colons = repmat({':'},1,ndims(x)-2);
    Vxx_size = size(spectra);    
    Vxx_size(2) = 2*Vxx_size(2);
    Vxx = zeros(Vxx_size);
    Vxx(:, 1:2:end-1, colons{:}) = spectra*conf_scale(1);
    Vxx(:, 2:2:end, colons{:}) = spectra*conf_scale(2);
  endif
  if (nargout >= 1)
    pxx = spectra;
    iargout = 1;
    if (info.have_f_arg || info.have_fs_arg)
      varargout{iargout++} = freq;
    elseif (info.have_w_arg)
      varargout{iargout++} = w;
    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"}
        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;
