## Copyright (C) 2021-2022 Andreas Bertsatos <abertsatos@biol.uoa.gr>
## Copyright (C) 2022 Andrew Penn <A.C.Penn@sussex.ac.uk>
##
## This file is part of the statistics package for GNU Octave.
##
## This program 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.
##
## This program 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
## this program; if not, see <http://www.gnu.org/licenses/>.

## -*- texinfo -*-
## @deftypefn  {statistics} {@var{p} =} anova1 (@var{x})
## @deftypefnx {statistics} {@var{p} =} anova1 (@var{x}, @var{group})
## @deftypefnx {statistics} {@var{p} =} anova1 (@var{x}, @var{group}, @var{displayopt})
## @deftypefnx {statistics} {@var{p} =} anova1 (@var{x}, @var{group}, @var{displayopt}, @var{vartype})
## @deftypefnx {statistics} {[@var{p}, @var{atab}] =} anova1 (@var{x}, @dots{})
## @deftypefnx {statistics} {[@var{p}, @var{atab}, @var{stats}] =} anova1 (@var{x}, @dots{})
##
## Perform a one-way analysis of variance (ANOVA) for comparing the means of two
## or more groups of data under the null hypothesis that the groups are drawn
## from distributions with the same mean. For planned contrasts and/or
## diagnostic plots, use @qcode{anovan} instead.
##
## anova1 can take up to three input arguments:
##
## @itemize
## @item
## @var{x} contains the data and it can either be a vector or matrix.
## If @var{x} is a matrix, then each column is treated as a separate group.
## If @var{x} is a vector, then the @var{group} argument is mandatory.
##
## @item
## @var{group} contains the names for each group.  If @var{x} is a matrix, then
## @var{group} can either be a cell array of strings of a character array, with
## one row per column of @var{x}.  If you want to omit this argument, enter an
## empty array ([]).  If @var{x} is a vector, then @var{group} must be a vector
## of the same length, or a string array or cell array of strings with one row
## for each element of @var{x}.  @var{x} values corresponding to the same value
## of @var{group} are placed in the same group.
##
## @item
## @var{displayopt} is an optional parameter for displaying the groups contained
## in the data in a boxplot.  If omitted, it is 'on' by default.  If group names
## are defined in @var{group}, these are used to identify the groups in the
## boxplot. Use 'off' to omit displaying this figure.
##
## @item
## @var{vartype} is an optional parameter to used to indicate whether the
## groups can be assumed to come from populations with equal variance. When
## @qcode{vartype} is @qcode{'equal'} the variances are assumed to be equal
## (this is the default). When @qcode{vartype} is @qcode{'unequal'} the
## population variances are not assumed to be equal and Welch's ANOVA test is
## used instead.
##
## @var{vartype} is an Octave extension: MATLAB's @code{anova1} takes no
## fourth argument.  It does not error on one either, it accepts it and
## ignores it, returning the same @var{p} and @var{F} for @qcode{'unequal'}
## as for @qcode{'equal'}.  Code written against this function and then run
## in MATLAB therefore gets the classic ANOVA silently, with no diagnostic
## of any kind.  Note that @code{anova2}'s analogous fourth argument does
## make MATLAB error, so the silence here is particular to @code{anova1}.
## @end itemize
##
## anova1 can return up to three output arguments:
##
## @itemize
## @item
## @var{p} is the p-value of the null hypothesis that all group means are equal.
##
## @item
## @var{atab} is a cell array containing the results in a standard ANOVA table.
##
## @item
## @var{stats} is a structure containing statistics useful for performing
## a multiple comparison of means with the MULTCOMPARE function.
## @end itemize
##
## A categorical @var{group} may declare levels that no observation uses.
## Such a level takes no part in the analysis and is dropped from every field
## of @var{stats}, so @code{n}, @code{means} and @code{gnames} always describe
## the same groups, in the same order, and can be indexed together.
##
## This is a deliberate deviation from MATLAB, which drops an unused level
## from @code{gnames} but keeps it in @code{n} and @code{means} as a count of
## zero and a mean of @code{NaN}.  Those fields then disagree in length and
## the group indices run past @code{gnames}, so MATLAB's own @code{multcompare}
## reports comparisons against a group holding no observations and labels them
## with indices that its @code{gnames} cannot resolve.
##
## If anova1 is called without any output arguments, then it prints the results
## in a one-way ANOVA table to the standard output. It is also printed when
## @var{displayopt} is 'on'.
##
##
## Examples:
##
## @example
## x = meshgrid (1:6);
## x = x + normrnd (0, 1, 6, 6);
## anova1 (x, [], 'off');
## [p, atab] = anova1(x);
## @end example
##
##
## @example
## x = ones (50, 4) .* [-2, 0, 1, 5];
## x = x + normrnd (0, 2, 50, 4);
## groups = @{"A", "B", "C", "D"@};
## anova1 (x, groups);
## @end example
##
## @seealso{anova2, anovan, multcompare}
## @end deftypefn

function [p, anovatab, stats] = anova1 (x, group, displayopt, vartype)

  ## Check for valid number of input arguments
  if (nargin < 1 || nargin > 4)
    error ("anova1: invalid number of input arguments.");
  endif
  ## Add defaults
  if (nargin < 2)
    group = [];
  endif
  if (nargin < 3)
    displayopt = 'on';
  endif
  if (nargin < 4)
    vartype = 'equal';
  endif
  plotdata = ! (strcmp (displayopt, 'off'));

  ## Convert group to cell array from character array, make it a column
  if (! isempty (group) && ischar (group))
    group = cellstr (group);
  endif
  if (size (group, 1) == 1)
    group = group';
  endif

  ## If x is a matrix, convert it to column vector and create a
  ## corresponding column vector for groups
  if (length (x) < prod (size (x)))
    [n, m] = size (x);
    x = x(:);
    gi = reshape (repmat ((1:m), n, 1), n*m, 1);
    if (length (group) == 0)          ## no group names are provided
      group = gi;
    elseif (size (group, 1) == m)     ## group names exist and match columns
      group = group(gi,:);
    else
      error ("anova1: columns in X and GROUP length do not match.");
    endif
  endif

  ## Check that x and group are the same size
  if (numel (x) != prod (size (group)))
    error ("anova1: GROUP must be a vector with the same number of rows as x.");
  endif

  ## Convert group to indices and separate names first
  if (isa (group, "categorical"))
    group_id = double (group(:));
    group_names = categories (group);
  else
    [group_id, group_names] = grp2idx (group);
    group_id = group_id(:);
  endif
  x = x(:);

  ## identify NaN values in x or missing/empty categories in the group and remove them.
  valid_data = ! isnan (x) & ! isnan (group_id);
  x = x(valid_data);
  group_id = group_id(valid_data);
  ## A declared but unused categorical level, or one emptied by NaN removal,
  ## is dropped from the design and from every field of STATS, so the group
  ## counts, means and names stay the same length.  MATLAB drops it from the
  ## names alone; see the deviation noted in the help text above.
  groups = size (group_names, 1);
  xs = accumarray (group_id, 1, [groups, 1], @sum, 0);
  observed_groups = (xs > 0);
  if (any (! observed_groups))
    group_map = cumsum (observed_groups);
    group_id = group_map(group_id);
    group_names = group_names(observed_groups, :);
    xs = xs(observed_groups);
  endif
  named = 1;

  ## Center data to improve accuracy and keep uncentered data for plotting
  xorig = x;
  mu = mean (x);
  x = x - mu;
  xr = x;

  ## Get group size and mean for each group
  groups = size (group_names, 1);
  xsum = accumarray (group_id, xr, [groups, 1], @sum, 0);
  xm = xsum ./ xs;
  xdev = xr - xm(group_id);
  xss = accumarray (group_id, xdev .^ 2, [groups, 1], @sum, 0);
  xv = xss ./ max (xs - 1, 1);
  xv(xs == 0) = NaN;
  xs = xs';
  xm = xm';
  xv = xv';

  ## Calculate statistics
  lx = length (xr);                       ## Number of samples in groups
  gm = mean (xr);                         ## Grand mean of groups
  dfm = length (xm) - 1;                  ## degrees of freedom for model
  dfe = lx - dfm - 1;                     ## degrees of freedom for error
  SSM = xs .* (xm - gm) * (xm - gm)';     ## Sum of Squares for Model
  SST = (xr(:) - gm)' * (xr(:) - gm);     ## Sum of Squares Total
  SSE = SST - SSM;                        ## Sum of Squares Error
  if (dfm > 0)
      MSM = SSM / dfm;                    ## Mean Square for Model
  else
      MSM = NaN;
  endif
  if (dfe > 0)
      MSE = SSE / dfe;                    ## Mean Square for Error
  else
      MSE = NaN;
  endif
  ## Calculate F statistic
  if (dfm <= 0)                     ## No between-group degrees of freedom.
    F = NaN;
    p = NaN;
  elseif (SSE != 0)                 ## Regular Matrix case.
    switch (lower (vartype))
      case 'equal'
        ## Assume equal variances (Fisher's One-way ANOVA)
        F = (SSM / dfm) / MSE;
      case 'unequal'
        ## Accommodate for unequal variances (Welch's One-way ANOVA)
        ## Calculate the sampling variance for each group (i.e. the square of the SEM)
        sv = xv ./ xs;
        ## Calculate weights as the reciprocal of the sampling variance
        w = 1 ./ sv;
        ## Calculate the origin
        ori = sum (w .* xm) ./ sum (w);
        ## Calculate Welch's F statistic
        F = (groups - 1)^-1 * sum (w .* (xm - ori).^2) /...
            (1 + ((2 * (groups - 2)/(groups^2 - 1)) * ...
            sum ((1 - w / sum (w)).^2 .* (xs - 1).^-1)));
        ## Welch's test does not use a pooled error term
        MSE = NaN;
        ## Correct the error degrees of freedom
        dfe = (3 /(groups^2 - 1) * sum ((1 - w / sum (w)).^2 .* (xs-1).^-1))^-1;
      otherwise
        error ("anova1: invalid fourth (vartype) argument to anova1.");
    endswitch
    p = 1 - fcdf (F, dfm, dfe);     ## Probability of F given equal means.
  elseif (SSM == 0)                 ## Constant Matrix case.
    ## Both sums-of-squares vanish, so F is 0/0.  There is no within-group
    ## scale to test against, and the limit depends on the direction of
    ## approach, so the statistic is undefined rather than zero.
    F = NaN;
    p = NaN;
  else                              ## Perfect fit case.
    F = Inf;
    p = 0;
  endif

  ## Create results table (if requested)
  if (nargout > 1)
    switch (lower (vartype))
      case 'equal'
        anovatab = {'Source', 'SS', 'df', 'MS', 'F', 'Prob>F'; ...
                    'Groups', SSM, dfm, MSM, F, p; ...
                    'Error', SSE, dfe, MSE, '', ''; ...
                    'Total', SST, dfm + dfe, '', '', ''};
      case 'unequal'
        anovatab = {'Source', 'F', 'df', 'dfe', 'F', 'Prob>F'; ...
                    'Groups', SSM, dfm, dfe, F, p};
    endswitch
  endif
  ## Create stats structure (if requested) for MULTCOMPARE
  if (nargout > 2)
    if (length (group_names) > 0)
        stats.gnames = group_names;
    else
        stats.gnames = strjust (num2str ((1:length (xm))'), 'left');
    endif
    stats.n = xs;
    stats.source = 'anova1';
    stats.vartype = vartype;
    stats.means = xm + mu;
    stats.vars = xv;
    stats.df = dfe;
    stats.s = sqrt (MSE);
  endif
  ## Print results table on screen if no output argument was requested
  if (nargout == 0 || plotdata)
    switch (lower (vartype))
      case 'equal'
        printf ("\n                      ANOVA Table\n\n");
        printf ("Source        SS      df        MS       F      Prob>F\n");
        printf ("------------------------------------------------------\n");
        printf ("Groups  %10.4f %5.0f %10.4f %8.2f %9.4f\n", SSM, dfm, MSM, F, p);
        printf ("Error   %10.4f %5.0f %10.4f\n", SSE, dfe, MSE);
        printf ("Total   %10.4f %5.0f\n\n", SST, dfm + dfe);
      case 'unequal'
        printf ("\n           Welch's ANOVA Table\n\n");
        printf ("Source        F     df     dfe     Prob>F\n");
        printf ("-----------------------------------------\n");
        printf ("Groups  %8.2f %5.0f %7.2f %10.4f\n\n", F, dfm, dfe, p);
    endswitch
  endif
  ## Plot data using BOXPLOT (unless opted out)
  if (plotdata)
    ## Create a new figure and ensure it becomes active
    f = figure ();
    set (0, 'currentfigure', f);

    ## Create a new axes inside this figure and make it current
    ax = axes ('parent', f);
    set (f, 'currentaxes', ax);

    ## Now boxplot will ALWAYS draw here
    boxplot (x, group_id, 'Notch', 'on', 'Labels', group_names);

  endif

endfunction


%!demo
%! x = meshgrid (1:6);
%! randn ('seed', 15);    # for reproducibility
%! x = x + normrnd (0, 1, 6, 6);
%! anova1 (x, [], 'off');

%!demo
%! x = meshgrid (1:6);
%! randn ('seed', 15);    # for reproducibility
%! x = x + normrnd (0, 1, 6, 6);
%! [p, atab] = anova1 (x);

%!demo
%! x = ones (50, 4) .* [-2, 0, 1, 5];
%! randn ('seed', 13);    # for reproducibility
%! x = x + normrnd (0, 2, 50, 4);
%! groups = {'A', 'B', 'C', 'D'};
%! anova1 (x, groups);

%!demo
%! y = [54 87 45; 23 98 39; 45 64 51; 54 77 49; 45 89 50; 47 NaN 55];
%! g = [1  2  3 ; 1  2  3 ; 1  2  3 ; 1  2  3 ; 1  2  3 ; 1  2  3 ];
%! anova1 (y(:), g(:), 'on', 'unequal');

## testing against GEAR.DAT data file and results for one-factor ANOVA from
## https://www.itl.nist.gov/div898/handbook/eda/section3/eda354.htm
%!test
%! data = [1.006, 0.996, 0.998, 1.000, 0.992, 0.993, 1.002, 0.999, 0.994, 1.000, ...
%!         0.998, 1.006, 1.000, 1.002, 0.997, 0.998, 0.996, 1.000, 1.006, 0.988, ...
%!         0.991, 0.987, 0.997, 0.999, 0.995, 0.994, 1.000, 0.999, 0.996, 0.996, ...
%!         1.005, 1.002, 0.994, 1.000, 0.995, 0.994, 0.998, 0.996, 1.002, 0.996, ...
%!         0.998, 0.998, 0.982, 0.990, 1.002, 0.984, 0.996, 0.993, 0.980, 0.996, ...
%!         1.009, 1.013, 1.009, 0.997, 0.988, 1.002, 0.995, 0.998, 0.981, 0.996, ...
%!         0.990, 1.004, 0.996, 1.001, 0.998, 1.000, 1.018, 1.010, 0.996, 1.002, ...
%!         0.998, 1.000, 1.006, 1.000, 1.002, 0.996, 0.998, 0.996, 1.002, 1.006, ...
%!         1.002, 0.998, 0.996, 0.995, 0.996, 1.004, 1.004, 0.998, 0.999, 0.991, ...
%!         0.991, 0.995, 0.984, 0.994, 0.997, 0.997, 0.991, 0.998, 1.004, 0.997];
%! group = [1:10] .* ones (10,10);
%! group = group(:);
%! [p, tbl] = anova1 (data, group, 'off');
%! assert_equal (p, 0.022661, 1e-6);
%! assert_equal (tbl{2,5}, 2.2969, 1e-4);
%! assert_equal (tbl{2,3}, 9, 0);
%! assert_equal (tbl{4,2}, 0.003903, 1e-6);
%! data = reshape (data, 10, 10);
%! [p, tbl, stats] = anova1 (data, [], 'off');
%! assert_equal (p, 0.022661, 1e-6);
%! assert_equal (tbl{2,5}, 2.2969, 1e-4);
%! assert_equal (tbl{2,3}, 9, 0);
%! assert_equal (tbl{4,2}, 0.003903, 1e-6);
%! means = [0.998, 0.9991, 0.9954, 0.9982, 0.9919, 0.9988, 1.0015, 1.0004, 0.9983, 0.9948];
%! N = 10 * ones (1, 10);
%! assert_equal (stats.means, means, 1e-6);
%! assert_equal (length (stats.gnames), 10, 0);
%! assert_equal (stats.n, N, 0);

## testing against one-way ANOVA example dataset from GraphPad Prism 8
%!test
%! y = [54 87 45; 23 98 39; 45 64 51; 54 77 49; 45 89 50; 47 NaN 55];
%! g = [1  2  3 ; 1  2  3 ; 1  2  3 ; 1  2  3 ; 1  2  3 ; 1  2  3 ];
%! [p, tbl] = anova1 (y(:), g(:), 'off', 'equal');
%! assert_equal (p, 0.00004163, 1e-6);
%! assert_equal (tbl{2,5}, 22.573418, 1e-6);
%! assert_equal (tbl{2,3}, 2, 0);
%! assert_equal (tbl{3,3}, 14, 0);
%! [p, tbl] = anova1 (y(:), g(:), 'off', 'unequal');
%! assert_equal (p, 0.00208877, 1e-8);
%! assert_equal (tbl{2,5}, 15.523192, 1e-6);
%! assert_equal (tbl{2,3}, 2, 0);
%! assert_equal (tbl{2,4}, 7.5786897, 1e-6);

## testing against one-way ANOVA example dataset from GraphPad Prism 8
## using categorical array as a grouping variable
%!test
%! y = [54, 87, 45; 23, 98, 39; 45, 64, 51; ...
%!      54, 77, 49; 45, 89, 50; 47, NaN, 55];
%! g = categorical ([1, 2, 3; 1, 2, 3; 1, 2, 3; 1, 2, 3; 1, 2, 3; 1, 2, 3]);
%! [p, tbl] = anova1 (y(:), g(:), 'off', 'equal');
%! assert_equal (p, 0.00004163, 1e-6);
%! assert_equal (tbl{2,5}, 22.573418, 1e-6);
%! assert_equal (tbl{2,3}, 2, 0);
%! assert_equal (tbl{3,3}, 14, 0);
%! [p, tbl] = anova1 (y(:), g(:), 'off', 'unequal');
%! assert_equal (p, 0.00208877, 1e-8);
%! assert_equal (tbl{2,5}, 15.523192, 1e-6);
%! assert_equal (tbl{2,3}, 2, 0);
%! assert_equal (tbl{2,4}, 7.5786897, 1e-6);

## testing handling of missing values in both data and grouping variables
%!test
%! y = [10; 20; 9999; NaN; 40; 50];
%! g = [1; 1; NaN; 1; 2; 2];
%! [p, tbl, stats] = anova1 (y, g, 'off');
%! assert_equal (p, 0.051317, 1e-6);
%! assert_equal (tbl{2,5}, 18, 1e-6);
%! assert_equal (tbl{2,3}, 1, 0);
%! assert_equal (tbl{3,3}, 2, 0);
%! assert_equal (tbl{4,3}, 3, 0);
%! assert_equal (stats.n, [2, 2], 0);

## Grouped variance remains stable when group means have large offsets.
%!test
%! y = [1e12 + (1:10), -1e12 + (1:10)](:);
%! g = [ones(10, 1); 2 * ones(10, 1)];
%! [~, ~, stats] = anova1 (y, g, 'off', 'unequal');
%! assert_equal (stats.vars, [55 / 6, 55 / 6], 1e-10);

## A single group leaves no between-group degrees of freedom.
%!test
%! [p, tbl] = anova1 ((1:5)', ones (5, 1), 'off');
%! assert_equal (p, NaN);
%! assert_equal (tbl{2, 3}, 0);
%! assert_equal (tbl{3, 2}, 10);
%! [p, tbl] = anova1 (7, 1, 'off');
%! assert_equal (p, NaN);
%! assert_equal (tbl{3, 3}, 0);

## One observation per group leaves no residual, so the fit is exact.
%!test
%! [p, tbl] = anova1 ([1; 2], [1; 2], 'off');
%! assert_equal (p, 0);
%! assert_equal (tbl{2, 2}, 0.5);
%! assert_equal (tbl{2, 3}, 1);
%! assert_equal (tbl{3, 2}, 0);
%! assert_equal (tbl{3, 3}, 0);

## Identical observations leave no scale to test against, so F is 0/0.
%!test
%! [p, tbl] = anova1 (ones (6, 1), [1; 1; 1; 2; 2; 2], 'off');
%! assert_equal (p, NaN);
%! assert_equal (tbl{2, 5}, NaN);
%! assert_equal (tbl{2, 2}, 0);
%! assert_equal (tbl{2, 3}, 1);
%! assert_equal (tbl{3, 2}, 0);
%! assert_equal (tbl{3, 3}, 4);

## An unused categorical level leaves every field of stats, not just
## gnames as in MATLAB, so the three stay indexable together.
%!test
%! y = (1:6)';
%! g = categorical ([1; 1; 3; 3; 3; 1], [1, 2, 3]);
%! [p, tbl, stats] = anova1 (y, g, 'off');
%! [p_ref, tbl_ref] = anova1 (y, [1; 1; 2; 2; 2; 1], 'off');
%! assert_equal (p, p_ref, 1e-12);
%! assert_equal (tbl, tbl_ref);
%! assert_equal (stats.n, [3, 3]);
%! assert_equal (stats.gnames, {'1'; '3'});
%! assert_equal (stats.means, [3, 4]);
%! g = categorical ([3; 3; 1; 1], [3, 2, 1]);
%! [~, ~, stats] = anova1 ((1:4)', g, 'off');
%! assert_equal (stats.gnames, {'3'; '1'});
%! assert_equal (stats.means, [1.5, 3.5]);

## Many factor levels retain the expected degrees of freedom.
%!test
%! g = kron ((1:120)', ones (2, 1));
%! [p, tbl, stats] = anova1 ((1:240)', g, 'off');
%! assert_equal (isfinite (p), true);
%! assert_equal (tbl{2, 3}, 119);
%! assert_equal (tbl{3, 3}, 120);
%! assert_equal (stats.n, 2 * ones (1, 120));

## testing handling of missing values in both data and grouping variables
## using categorical array as a grouping variable
%!test
%! y = [10; 20; 9999; NaN; 40; 50];
%! g = categorical ([1; 1; NaN; 1; 2; 2]);
%! [p, tbl, stats] = anova1 (y, g, 'off');
%! assert_equal (p, 0.051317, 1e-6);
%! assert_equal (tbl{2,5}, 18, 1e-6);
%! assert_equal (tbl{2,3}, 1, 0);
%! assert_equal (tbl{3,3}, 2, 0);
%! assert_equal (tbl{4,3}, 3, 0);
%! assert_equal (stats.n, [2, 2], 0);
