## Copyright (C) 1995-2017 Kurt Hornik
## Copyright (C) 2018 John Donoghue
## Copyright (C) 2021 Stefano Guidoni
## Copyright (C) 2022-2025 Andreas Bertsatos <abertsatos@biol.uoa.gr>
## Copyright (C) 2025 Yasin Achengli
##
## 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; see the file COPYING.  If not, see
## <http://www.gnu.org/licenses/>.

## -*- texinfo -*-
## @deftypefn  {statistics} {@var{t} =} crosstab (@var{x1}, @var{x2})
## @deftypefnx {statistics} {@var{t} =} crosstab (@var{x1}, @dots{}, @var{xn})
## @deftypefnx {statistics} {[@var{t}, @var{chisq}, @var{p}, @var{labels}] =} crosstab (@dots{})
##
## Create a cross-tabulation (contingency table) @var{t} from data vectors.
##
## The inputs @var{x1}, @var{x2}, ... @var{xn} must be vectors of equal length
## with a data type of numeric, logical, char array, categorical, strings, or
## cell array of character vectors.
##
## As additional return values @code{crosstab} returns the chi-square statistics
## @var{chisq}, its p-value @var{p} and a cell array @var{labels}, containing
## the labels of each input argument.
##
## @seealso{grp2idx, tabulate}
## @end deftypefn

function [t, chisq, p, labels] = crosstab (varargin)

  ## check input
  if (nargin < 2)
    print_usage ();
  endif

  ## main - begin
  v_length = [];                    # vector of lengths of input vectors
  reshape_format = [];              # vector of the dimensions of t
  X = [];                           # matrix of the indexed input values
  labels_data = {};                 # temporary cell array for unique labels
  coordinates = {};                 # cell array of unique elements

  for i = 1:nargin
    vector = varargin{i};
    ## Convert char, cellstr, or categorical to indexed numeric vector
    if (ischar (vector) || iscellstr (vector) ||
                           iscategorical (vector) || isstring (vector))
      try
        [vector, gnames] = grp2idx (vector);
        labels_data{i} = gnames;
      catch
        error ("crosstab: x1, x2 ... xn must be vectors.");
      end_try_catch
    elseif (isnumeric (vector) || islogical (vector))
      if (! isvector (vector))
        error ("crosstab: x1, x2 ... xn must be vectors.");
      endif
      vector = vector(:);
      unique_vals = unique (vector(! isnan (vector)));
      labels_data{i} = cellstr (num2str (unique_vals));
    else
      error ("crosstab: unsupported type for data vector.");
    endif
    v_length(i) = length (vector);
    if (length (unique (v_length)) != 1)
      error ("crosstab: x1, x2 ... xn must be vectors of the same length.");
    endif
    X = [X, vector];
    reshape_format(i) = length (unique (vector(! isnan (vector))));
    coordinates(i) = unique (vector(! isnan (vector)));
  endfor

  if (nargout > 3)
    max_rows = 0;
    for i = 1:nargin
      max_rows = max (max_rows, numel (labels_data{i}));
    endfor
    labels = cell (max_rows, nargin);
    for i = 1:nargin
      col_labels = labels_data{i};
      labels(1:numel (col_labels), i) = col_labels;
    endfor
  else
    labels = {};
  endif

  t = zeros (reshape_format);

  ## Main logic:
  ## For each combination of x1, x2, ... xn, search in unique elements stored
  ## in coordinates for each dimension and increment the position value in t
  ## multidimensional matrix (always if there is no NaN element in the combination).
  for idx = 1:size (X, 1)
    if (! any (isnan (X(idx,:))))
      location = zeros (1,size (X, 2));
      for jdx = 1:size (X,2)
        location(jdx) = find (cell2mat (coordinates(jdx)) == X(idx, jdx));
      endfor
      t(num2cell (location){:}) += 1;
    endif
  endfor

  if (nargout > 1)
    if (isscalar (t) || isempty (t))
      chisq = NaN;
      p = NaN;
    else
      [p, chisq] = chi2test (t);
    endif
  endif
endfunction

## Test input validation
%!error crosstab ()
%!error crosstab (1)
%!error <crosstab: x1, x2 ... xn must be vectors.> crosstab (ones (2), [1 1])
%!error <crosstab: x1, x2 ... xn must be vectors.> crosstab ([1 1], ones (2))
%!error <x1, x2 .* must be vectors of the same length> crosstab ([1], [1 2])
%!error <x1, x2 .* must be vectors of the same length> crosstab ([1 2], [1])
%!error <crosstab: unsupported type for data vector.> crosstab ([1 2], {1, 2})
%!test
%! load carbig
%! [t, chisq, p, labels] = crosstab (cyl4, when, org);
%! assert_equal (t(2,3,1), 38);
%! assert_equal (labels{3,3}, "Japan");
%!test
%! load carbig
%! [t, chisq, p, labels] = crosstab (cyl4, when, org);
%! assert_equal (t(2,3,2), 17);
%! assert_equal (labels{1,3}, "USA");
%!test
%! x = [1, 1, 2, 3, 1];
%! y = [1, 2, 5, 3, 1];
%! t = crosstab (x, y);
%! assert_equal (t, [2, 1, 0, 0; 0, 0, 0, 1; 0, 0, 1, 0]);
%!test
%! x = [1, 1, 2, 3, 1];
%! y = [1, 2, 3, 5, 1];
%! t = crosstab (x, y);
%! assert_equal (t, [2, 1, 0, 0; 0, 0, 1, 0; 0, 0, 0, 1]);
%!test
%! x1 = [1, 3, 7, 7, 8];
%! x2 = [4, 2, 1, 1, 1];
%! x3 = [6, 2, 6, 2, NaN];
%! T1 = [0, 0, 0; 0, 1, 0; 1, 0, 0; 0, 0, 0];
%! T2 = [0, 0, 1; 0, 0, 0; 1, 0, 0; 0, 0, 0];
%! T = zeros (4, 3, 2);
%! T(:,:,1) = T1;
%! T(:,:,2) = T2;
%! t = crosstab (x1, x2, x3);
%! assert_equal (t, T);
%!test
%! x = [1, 2, NaN, 1];
%! y = [1, 2, 3, NaN];
%! t = crosstab (x, y);
%! assert_equal (t, [1, 0, 0; 0, 1, 0]);

## Test categorical input
%!test
%! x = categorical ({'A', 'B', 'A', 'C', 'B'});
%! y = [1, 2, 1, 3, 2];
%! t = crosstab (x, y);
%! assert_equal (size (t), [3, 3]);
%! assert_equal (t(1, 1), 2);  # A with 1
%! assert_equal (t(2, 2), 2);  # B with 2
%!test
%! x = categorical ({'low', 'med', 'high', 'low', 'med'});
%! y = categorical ({'X', 'Y', 'X', 'Y', 'X'});
%! t = crosstab (x, y);
%! assert_equal (size (t), [3, 2]);
%!test
%! ## Test categorical with numeric
%! x = categorical ([10, 20, 10, 30, 20]);
%! y = [1, 2, 1, 3, 2];
%! t = crosstab (x, y);
%! assert_equal (t, [2, 0, 0; 0, 2, 0; 0, 0, 1]);
%!test
%! smoker = [1 1 0 0 1 0 1 1 0 0 1 0]';
%! gender = [1 0 1 0 1 1 0 0 1 0 0 1]';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p, labels] = crosstab (smoker, gender);
%! warning (w);
%! assert_equal (t, [2 4; 4 2]);
%! assert_equal (chisq, 1.33333333, 1e-8);
%! assert_equal (p, 0.24821308, 1e-8);
%! assert_equal (labels{1,1}, '0');
%! assert_equal (labels{1,2}, '0');
%! assert_equal (labels{2,1}, '1');
%! assert_equal (labels{2,2}, '1');
%!test
%! ## Test for categorical
%! smk_cat = categorical ([0 0 1 1 0 1 0 0 1 1 0 1]');
%! gen_cat = categorical ([0 1 0 1 0 0 1 1 0 1 1 0]');
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (smk_cat, gen_cat);
%! warning (w);
%! assert_equal (t, [2 4; 4 2]);
%! assert_equal (chisq, 1.33333333, 1e-6);
%! assert_equal (p, 0.24821308, 1e-6);
%!test
%! x = [1 1 1 2 2 2 3 3 3]';
%! y = [1 1 1 2 2 2 3 3 3]';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p, labels] = crosstab (x, y);
%! warning (w);
%! assert_equal (t, diag ([3 3 3]));
%! assert_equal (chisq, 18.00000000);
%! assert_equal (p, 0.00123410, 1e-8);
%!test
%! ## Test for Partial NaN giving NaN for chisq/p
%! x7 = [1 2 3 4 NaN NaN]';
%! y7 = [10 20 30 40 50 60]';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p, labels] = crosstab (x7, y7);
%! warning (w);
%! assert_equal (t, [eye(4), zeros(4, 2)]);
%! assert_equal (isnan (chisq), true);
%! assert_equal (isnan (p), true);
%! assert_equal (labels{1,1}, '1');
%! assert_equal (labels{1,2}, '10');
%! assert_equal (labels{2,1}, '2');
%! assert_equal (labels{2,2}, '20');
%! assert_equal (labels{3,1}, '3');
%! assert_equal (labels{3,2}, '30');
%! assert_equal (labels{4,1}, '4');
%! assert_equal (labels{4,2}, '40');
%! assert_equal (isempty (labels{5,1}), true);
%! assert_equal (labels{5,2}, '50');
%! assert_equal (isempty (labels{6,1}), true);
%! assert_equal (labels{6,2}, '60');
%!test
%! a = ones (15,1);
%! b = ones (15,1);
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (a, b);
%! warning (w);
%! assert_equal (t, 15);
%! assert_equal (isnan (chisq), true);
%! assert_equal (isnan (p), true);
%!test
%! ## all NaN → empty table + NaN stats
%! na = NaN (6,1);
%! nb = (1:6)';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (na, nb);
%! warning (w);
%! assert_equal (all (t(:) == 0), true);
%! assert_equal (isnan (chisq), true);
%! assert_equal (isnan (p), true);
%!test
%! ## single observation → 1×1 table, NaN statistic
%! w = warning ();
%! warning ('off');
%! [t, chisq, p, labels] = crosstab (5, 'Z');
%! warning (w);
%! assert_equal (t, 1);
%! assert_equal (isnan (chisq), true);
%! assert_equal (isnan (p), true);
%! assert_equal (labels{1,1}, '5');
%! assert_equal (labels{1,2}, 'Z');
%!test
%! xx = [1 1 1 1 2 2 2 2]';
%! yy = [10 10 10 10 20 20 20 20]';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (xx, yy);
%! warning (w);
%! assert_equal (t, [4 0; 0 4]);
%! assert_equal (chisq, 8.00000000);
%! assert_equal (p, 0.00467773, 1e-8);
%!test
%! set1 = repmat ((1:5)', 20, 1);
%! set2 = repmat ([1; 2], 50, 1);
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (set1, set2);
%! warning (w);
%! assert_equal (t, 10 * ones (5,2));
%! assert_equal (chisq, 0);
%! assert_equal (p, 1);
%!test
%! ## 3-way table with NaN
%! a = [1 1 2 2 3 3 1]';
%! b = [1 2 1 2 1 2 1]';
%! c = [1 1 NaN 2 2 1 2]';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (a, b, c);
%! warning (w);
%! expected(:,:,1) = [1 1; 0 0; 0 1];
%! expected(:,:,2) = [1 0; 0 1; 1 0];
%! assert_equal (t, expected);
%! assert_equal (chisq, 6.00000000);
%! assert_equal (p, 0.53974935, 1e-9);
%!test
%! ## sparse cellstr
%! g = [1 5 7 12 1 5 19]';
%! h = {'A', 'B', 'A', 'C', 'D', 'E', 'A'}';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (g, h);
%! warning (w);
%! assert_equal (sum (t(:)), 7);
%! assert_equal (chisq, 16.33333333, 1e-8);
%! assert_equal (p, 0.42994852, 1e-8);
%!test
%! ## string array
%! str1 = ['low'; 'high'; 'med'; 'low'; 'high'];
%! str2 = ['X'; 'Y'; 'X'; 'Y'; 'X'];
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (str1, str2);
%! warning (w);
%! assert_equal (t, [1 1; 1 1; 1 0]);
%! assert_equal (chisq, 0.83333333, 1e-8);
%! assert_equal (p, 0.659240631, 1e-9);
%!test
%! ## cellstr
%! c1 = {'A','B','A','C','B','A'}';
%! c2 = {'1','2','1','3','2','1'}';
%! w = warning ();
%! warning ('off');
%! [t, chisq, p] = crosstab (c1, c2);
%! warning (w);
%! assert_equal (t, [3 0 0; 0 2 0; 0 0 1]);
%! assert_equal (chisq, 12.00000000, 1e-14);
%! assert_equal (p, 0.01735127, 1e-8);
