########################################################################
##
## Copyright (C) 1998-2021 The Octave Project Developers
##
## See the file COPYRIGHT.md in the top-level directory of this
## distribution or <https://octave.org/copyright/>.
##
## This file is part of Octave.
##
## Octave is free software: you can redistribute it and/or modify it
## under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## Octave is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with Octave; see the file COPYING.  If not, see
## <https://www.gnu.org/licenses/>.
##
########################################################################

## -*- texinfo -*-
## @deftypefn  {} {@var{s} =} sumsqr (@var{x})
## @deftypefnx {} {[@var{s}, @var{n}] =} sumsqr (@var{x})
##
## Square each element of the array or cell input @var{x} and return the sum of
## the result as a scalar output @var{s}.  @var{x} may consist of numeric, char,
## logical, or complex values, and cells may contain mixed types.  Non-finite
## (Inf, NaN, NA) values are ignored for the calculation.  If requested, the
## optional output argument @var{n} will contain the number of finite elements
## in @var{x} used to calculate @var{s}.
##
## @example
## @group
## a = [1 2 Inf -3];
## [s, n] = sumsqr (@var{a})
## @result{}
##    s =
##        14
##    n =
##         3
## @end group
## @end example
## @seealso{sum, sumsq, sumabs, meansqr, meansq, meanabs}
## @end deftypefn

function [s, n] = sumsqr (x)

  if ((nargin != 1) || (!any (nargout == [0 1 2])))
    print_usage ();
  endif

  if (isnumeric (x) || islogical (x) || ischar (x))

    ## square first to convert any char to numeric, output column vector
    x = x(:).^2;

  elseif strcmp (class (x), "cell")
    ## input validation: check contents after cell flattened

    ## output 1 column cell array with 1 column array cell contents to avoid
    ## problems with mismatched cell content sizes
    x = flatten_to_column_cell (x);

    ##check flattened content types
    if ! all(cellfun (@(C) isnumeric (C) || islogical (C) || ischar (C), x))
      bad_cells = ! cellfun (@(C) isnumeric (C) || islogical (C) || ischar (C), x);
      bad_class = find (bad_cells, 1, "first");
      error ("sumsqr: SUMSQR not defined for cell arrays containing type %s", class (x{bad_class}));
    endif

    ## convert any chars to numeric with @power, then revert to array.
    x = cell2mat (cellfun (@power, x, {2}, "UniformOutput", false));

  else
    error ("sumsqr: SUMSQR not defined for type %s", class (x));
  endif

  ## convert to single column input and exclude non-finite (NaN,Inf, NA)
  x = x(isfinite (x));

  n = numel (x);

  if n == 0
    s = 0;
  else
    s = sum (x);
  endif

endfunction

function C = flatten_to_column_cell (C)
  ## recursively flatten cell of any depth into a n x 1 cell array.
  if iscell (C)
    C = cellfun (@flatten_to_column_cell, C, 'UniformOutput', false);
    C = vertcat (C{:});
  else
    C = {C(:)};
  endif
endfunction

%!assert (sumsqr (2), 4)
%!assert (class (sumsqr (single (2))), "single")
%!assert (class (sumsqr ({single(2)})), "single")
%!assert (sumsqr (2+3i), -5+12i)
%!assert (sumsqr ([1, 2, 3]), 14)
%!assert (sumsqr ([1, 2, 3]'), 14)
%!assert (sumsqr (magic (2)), 30)
%!assert (sumsqr ({magic(2)}), 30)
%!assert (sumsqr ('test'), 50338)
%!assert (sumsqr (cat (3, magic(2), magic(2))), 60)
%!assert (sumsqr ([1+i, 2-i; 3+2i, 4-2i]), 20-6i)
%!assert (sumsqr ({[1+i, 2-i; 3+2i, 4-2i]}), 20-6i)

##test mixed inputs
%!assert (sumsqr ([{[{1}, {2}]},{[3; 4]'}]), 30)
%!assert (sumsqr ({1, 2, [3, 4]', 'test'}), 50368)
%!assert (sumsqr ([{1}, {[{2, {[3, 4], [5, 6]}, 7}]}]), 140)

## test nonfinite inputs
%!assert (sumsqr ([]), 0)
%!assert (sumsqr (Inf), 0)
%!assert (sumsqr ([Inf, NA, NaN]), 0)
%!assert (sumsqr ([Inf, NA, NaN, 1]), 1)

%!test
%! a = [{'test'}, {[{2, 7}]}, Inf];
%! [s, n] = sumsqr (a);
%! assert (s, 50391)
%! assert (n, 6)

## Test input validation
%!error sumsqr ()
%!error sumsqr (1, 2)
%!error [a, b, c] = sumsqr (1, 2)
%!error <SUMSQR not defined for type> sumsqr (struct("a",1))
%!error <SUMSQR not defined for cell arrays containing type> sumsqr ({1, struct("a",1)})