## Copyright (C) 2021 Stefano Guidoni <ilguido@users.sf.net>
##
## 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{T} =} cluster (@var{Z}, "Cutoff", @var{C})
## @deftypefnx {statistics} {@var{T} =} cluster (@var{Z}, "Cutoff", @var{C}, "Depth", @var{D})
## @deftypefnx {statistics} {@var{T} =} cluster (@var{Z}, "Cutoff", @var{C}, "Criterion", @var{criterion})
## @deftypefnx {statistics} {@var{T} =} cluster (@var{Z}, "MaxClust", @var{N})
##
## Define clusters from an agglomerative hierarchical cluster tree.
##
## Given a hierarchical cluster tree @var{Z} generated by the @code{linkage}
## function, @code{cluster} defines clusters, using a threshold value @var{C} to
## identify new clusters ('Cutoff') or according to a maximum number of desired
## clusters @var{N} ('MaxClust').
##
## @var{criterion} is used to choose the criterion for defining clusters, which
## can be either "inconsistent" (default) or "distance". When using
## "inconsistent", @code{cluster} compares the threshold value @var{C} to the
## inconsistency coefficient of each link; when using "distance", @code{cluster}
## compares the threshold value @var{C} to the height of each link.
## @var{D} is the depth used to evaluate the inconsistency coefficient, its
## default value is 2.
##
## @code{cluster} uses "distance" as a criterion for defining new clusters when
## it is used with the 'MaxClust' method.
##
## @seealso{clusterdata, dendrogram, inconsistent, kmeans, linkage, pdist}
## @end deftypefn

function T = cluster (Z, opt, varargin)
  switch (lower (opt))
  ## check the input
    case 'cutoff'
      if (nargin < 3)
        print_usage ();
      else
        C = varargin{1};
        D = 2;
        criterion = 'inconsistent';
        if (nargin > 3)
          pair_index = 2;
          while (pair_index < (nargin - 2))
            switch (lower (varargin{pair_index}))
              case 'depth'
                D = varargin{pair_index + 1};
              case 'criterion'
                criterion = varargin{pair_index + 1};
              otherwise
                error ("cluster: unknown property '%s'.", varargin{pair_index});
            endswitch
            pair_index += 2;
          endwhile
        endif
      endif
      if ((! (isscalar (C) || isvector (C))) || (C < 0))
        error (strcat ("cluster: C must be a positive scalar", ...
                       " or a vector of positive numbers."));
      endif

    case 'maxclust'
      if (nargin != 3)
        print_usage ();
      else
        N = varargin{1};
        C = [];
      endif
      if ((! (isscalar (N) || isvector (N))) || (N < 0))
        error (strcat ("cluster: N must be a positive number", ...
                       " or a vector of positive numbers."));
      endif

    otherwise
      error ("cluster: unknown option '%s'.", opt);
  endswitch

  if (columns (Z) != 3 || ! isnumeric (Z) ||
      ! (max (Z(end, 1:2)) == rows (Z) * 2))
    error ("cluster: Z must be a matrix generated by the linkage function.");
  endif

  ## number of observations
  n = rows (Z) + 1;

  ## vector of values used by the threshold check
  vThresholds = [];

  ## starting number of clusters
  nClusters = 1;

  ## the return value is the matrix T, constituted by one or more vector vT
  T = [];
  vT = zeros (1, n);

  ## main logic
  switch (lower (opt))
    case 'cutoff'
      switch (lower (criterion))
        case 'inconsistent'
          vThresholds = inconsistent (Z, D)(:, 4);
        case 'distance'
          vThresholds = Z(:, 3);
        otherwise
          error ("cluster: unknown criterion '%s'.", criterion);
      endswitch
      ## A node is a cluster only when it and everything below it are under
      ## the cutoff.  Testing the node alone leaves a low node holding a high
      ## descendant undivided, and so returns too few clusters.
      smax = subtree_maximum (Z, vThresholds, n);
      for k = 1:numel (C)
        T = [T; cut_by_cutoff(Z, n, smax, C(k))];
      endfor

    case 'maxclust'
      ## Undoing the highest N-1 merges leaves exactly N groups.  Cutting at
      ## the height of the Nth node from the top cannot: when merges tie there
      ## is no height between them, and on a tree whose merges all sit at one
      ## height -- every observation identical, say -- it separates nothing and
      ## every observation comes back in a cluster of its own.
      for k = 1:numel (N)
        T = [T; cut_by_count(Z, n, N(k))];
      endfor
  endswitch

  T = T';           # return value

endfunction

## Largest criterion value anywhere in the subtree each internal node roots.
## The rows of Z are in merge order, so a node's children are always earlier
## rows and one forward pass suffices.
function s = subtree_maximum (Z, v, n)
  s = v(:);
  for i = 1:rows (Z)
    for j = 1:2
      child = Z(i, j);
      if (child > n)
        s(i) = max (s(i), s(child - n));
      endif
    endfor
  endfor
endfunction

## Put every leaf below an internal node into one cluster.
function t = collect_leaves (Z, n, index, label, t)
  for j = 1:2
    child = Z(index, j);
    if (child > n)
      t = collect_leaves (Z, n, child - n, label, t);
    else
      t(child) = label;
    endif
  endfor
endfunction

## Descend from the root, taking a node as a cluster as soon as its whole
## subtree is under the cutoff and splitting it otherwise.
function t = cut_by_cutoff (Z, n, smax, c)
  t = zeros (1, n);
  k = 0;
  pending = rows (Z);
  while (! isempty (pending))
    index = pending(end);
    pending(end) = [];
    if (smax(index) < c)
      k++;
      t = collect_leaves (Z, n, index, k, t);
    else
      for j = 1:2
        child = Z(index, j);
        if (child > n)
          pending(end+1) = child - n;
        else
          k++;
          t(child) = k;
        endif
      endfor
    endif
  endwhile
endfunction

## Perform the lowest n-N merges and let what remains be the clusters, so that
## exactly N of them come back however the merge heights fall.
function t = cut_by_count (Z, n, N)
  if (N >= n)
    t = 1:n;
    return;
  endif
  keep = n - max (N, 1);
  comp = 1:n;                       # component of each leaf
  nodecomp = zeros (1, rows (Z));   # component of each internal node
  for i = 1:keep
    ca = node_component (Z(i, 1), n, comp, nodecomp);
    cb = node_component (Z(i, 2), n, comp, nodecomp);
    comp(comp == cb) = ca;
    nodecomp(i) = ca;
    nodecomp(nodecomp == cb) = ca;
  endfor
  [~, ~, t] = unique (comp);
  t = t(:)';
endfunction

## Component holding a node, whether it is a leaf or an internal node.
function c = node_component (node, n, comp, nodecomp)
  if (node > n)
    c = nodecomp(node - n);
  else
    c = comp(node);
  endif
endfunction


## Test input validation
%!error cluster ()
%!error <Z must be .* generated by the linkage .*> cluster ([1 1], 'Cutoff', 1)
%!error <unknown option .*> cluster ([1 2 1], 'Bogus', 1)
%!error <C must be a positive scalar .*> cluster ([1 2 1], 'Cutoff', -1)
%!error <unknown property .*> cluster ([1 2 1], 'Cutoff', 1, 'Bogus', 1)

## Test output
## The partitions below are MATLAB R2024a's on the same data.  Cluster numbers
## are arbitrary, so the tests compare groupings rather than labels.

%!shared X, Z
%! X = [0, 0; 0.1, 0.1; 0.2, 0; 5, 5; 5.1, 5.2; 5.2, 5.0; 10, 0; ...
%!      10.1, 0.2; 20, 20; 0.05, 0.3];
%! Z = linkage (pdist (X), "single");

%!test  # MaxClust returns exactly the number of clusters asked for
%! for n = 1:5
%!   assert_equal (numel (unique (cluster (Z, "MaxClust", n))), n);
%! endfor

%!test  # and the groupings are MATLAB's
%! t = cluster (Z, "MaxClust", 3);
%! assert_equal (numel (unique (t([1, 2, 3, 4, 5, 6, 10]))), 1);
%! assert_equal (numel (unique (t([7, 8]))), 1);
%! assert_equal (t(9) != t(1) && t(7) != t(1), true);

%!test  # merges at equal heights still divide: a height threshold cannot
%! ## split a tree whose every merge sits at the same height, which is what
%! ## five identical observations produce.
%! Zflat = linkage (pdist (ones (5, 2)), "single");
%! assert_equal (all (Zflat(:,3) == 0), true);
%! for n = 1:5
%!   assert_equal (numel (unique (cluster (Zflat, "MaxClust", n))), n);
%! endfor

%!test  # a tie at the cut does not cost a cluster either
%! Ztie = [1, 2, 1; 3, 4, 2; 6, 7, 2; 5, 8, 3];
%! for n = 1:5
%!   assert_equal (numel (unique (cluster (Ztie, "MaxClust", n))), n);
%! endfor

%!test  # a cluster is a node whose whole subtree is under the cutoff
%! assert_equal (numel (unique (cluster (Z, "Cutoff", 0.5))), 6);
%! assert_equal (numel (unique (cluster (Z, "Cutoff", 0.8))), 4);
%! assert_equal (numel (unique (cluster (Z, "Cutoff", 1.2))), 1);

%!test  # the depth the inconsistency is measured over reaches the result
%! assert_equal (numel (unique (cluster (Z, "Cutoff", 0.9, "Depth", 2))), 4);
%! assert_equal (numel (unique (cluster (Z, "Cutoff", 0.9, "Depth", 3))), 5);
%! t = cluster (Z, "Cutoff", 0.9, "Depth", 3);
%! assert_equal (numel (unique (t([1, 2, 3]))), 1);
%! assert_equal (numel (unique (t([4, 5, 6]))), 1);
%! assert_equal (t(10) != t(1), true);

%!test  # the distance criterion is unchanged
%! assert_equal (numel (unique (cluster (Z, "Cutoff", 1, ...
%!                                       "Criterion", "distance"))), 4);
%! assert_equal (numel (unique (cluster (Z, "Cutoff", 30, ...
%!                                       "Criterion", "distance"))), 1);

%!test  # a vector of cutoffs gives one column per cutoff
%! T = cluster (Z, "MaxClust", [2, 3, 4]);
%! assert_equal (size (T), [10, 3]);
%! assert_equal (max (T), [2, 3, 4]);

%!test  # well separated groups come out whole
%! Xs = [randn(10, 2) * 0.05 + 1; randn(10, 2) * 0.05 - 1];
%! t = cluster (linkage (pdist (Xs), "ward"), "MaxClust", 2);
%! assert_equal (numel (unique (t)), 2);
%! assert_equal (numel (unique (t(1:10))), 1);
%! assert_equal (numel (unique (t(11:20))), 1);
%! assert_equal (t(1) != t(11), true);

