## Copyright (C) 2021 Stefano Guidoni
##
## 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{leafOrder} =} optimalleaforder (@var{tree}, @var{D})
## @deftypefnx {statistics} {@var{leafOrder} =} optimalleaforder (@dots{}, @var{Name}, @var{Value})
##
## Compute the optimal leaf ordering of a hierarchical binary cluster tree.
##
## The optimal leaf ordering of a tree is the ordering which minimizes the sum
## of the distances between each leaf and its adjacent leaves, without altering
## the structure of the tree, that is without redefining the clusters of the
## tree.
##
## Required inputs:
## @itemize
## @item
## @var{tree}: a hierarchical cluster tree @var{tree} generated by the
## @code{linkage} function.
##
## @item
## @var{D}: a matrix of distances as computed by @code{pdist}.
## @end itemize
##
## Optional inputs can be the following property/value pairs:
## @itemize
## @item
## property 'Criteria' at the moment can only have the value 'adjacent',
## for minimizing the distances between leaves.
##
## @item
## property 'Transformation' can have one of the values 'linear', 'inverse'
## or a handle to a custom function which computes @var{S} the similarity
## matrix.
## @end itemize
##
## optimalleaforder's output @var{leafOrder} is the optimal leaf ordering.
##
## @strong{Reference}
## Bar-Joseph, Z., Gifford, D.K., and Jaakkola, T.S. Fast optimal leaf ordering
## for hierarchical clustering. Bioinformatics vol. 17 suppl. 1, 2001.
##
## @seealso{dendrogram,linkage,pdist}
##
## @end deftypefn

function leafOrder = optimalleaforder ( varargin )

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

  tree = varargin{1};
  D = varargin{2};
  criterion = 'adjacent';               # default and only value at the moment
  transformation = 'linear';

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

  ## Read the paired arguments.  Only the names have to be character vectors:
  ## 'Transformation' takes a function handle as well, and lower-casing every
  ## argument rejected it before it could be used.
  names = 3:2:(nargin - 1);
  if (! all (cellfun ('ischar', varargin(names))))
    error ("optimalleaforder: character inputs expected for arguments 3 and up");
  endif
  varargin(names) = lower (varargin(names));
  pair_index = 3;
  while (pair_index <= (nargin - 1))
    switch (varargin{pair_index})
      case 'criteria'
        criterion = varargin{pair_index + 1};
        if (! ischar (criterion))
          error ("optimalleaforder: character inputs expected for CRITERIA");
        endif
        criterion = lower (criterion);
        if (strcmp (criterion, 'group'))
          ## MATLAB compatibility:
          ## the 'group' criterion is not implemented
          error ("optimalleaforder: unavailable criterion 'group'");
        elseif (! strcmp (criterion, 'adjacent'))
          error ("optimalleaforder: invalid criterion %s", criterion);
        endif
      case 'transformation'
        transformation = varargin{pair_index + 1};
        if (ischar (transformation))
          transformation = lower (transformation);
        elseif (! is_function_handle (transformation))
          error (strcat ("optimalleaforder: TRANSFORMATION must be a", ...
                         " character vector or a function handle"));
        endif
      otherwise
        error ("optimalleaforder: unknown property %s", varargin{pair_index});
    endswitch

    pair_index += 2;
  endwhile

  ## D can be either a vector or a matrix,
  ## but it is easier to work with a matrix
  if (isvector (D))
    D = squareform (D);
  endif

  n = rows (D);
  m = rows (tree);

  if (n != (m + 1))
    error (strcat ("optimalleaforder: D must be a matrix or vector generated by", " the pdist function"));
  endif


  ## the similarity matrix, basically an inverted distance matrix
  S = zeros (n);

  if (strcmpi (transformation, 'linear'))
    ## linear similarity
    maxD = max (max (D));
    S = maxD - D;
  elseif (strcmpi (transformation, 'inverse'))
    ## similarity as inverted distance
    S = 1 ./ D;
  elseif (is_function_handle (transformation))
    ## custom similarity
    S = feval (transformation, D);
  else
    error ("optimalleaforder: invalid transformation %s", transformation);
  endif


  ## main body

  ## With two leaves the tree has a single node and the only orderings are the
  ## one below and its reverse, both optimal.  The tables below index the root
  ## where a leaf is expected and run off the end of M.
  if (n <= 2)
    leafOrder = 1:n;
    return;
  endif

  ## for each node v we compute the maximum similarity of the subtree M(w,u,v),
  ## where the leftmost leaf is w and the rightmost is u; remember that
  ## M(w,u,v) = M(u,w,v)
  M = zeros (n, n, n + m);

  ## O is a utility matrix: for each node of the tree we store the left and
  ## right leaves of the optimal subtree
  O = [1:( n + m ); 1:( n + m ); (zeros (1, (n + m)))]';

  ## compute M for every node v
  for iter = 1 : m
    v = iter + n;                                     # current node
    l = optimalleaforder_getLeafList (tree(iter, 1)); # the left subtree
    r = optimalleaforder_getLeafList (tree(iter, 2)); # the right subtree

    if (tree(iter,1) > n)
      l_l = optimalleaforder_getLeafList (tree(tree(iter, 1) - n, 1));
      l_r = optimalleaforder_getLeafList (tree(tree(iter, 1) - n, 2));
    else
      l_l = l_r = l;
    endif

    if (tree(iter,2) > n)
      r_l = optimalleaforder_getLeafList (tree(tree(iter, 2) - n, 1));
      r_r = optimalleaforder_getLeafList (tree(tree(iter, 2) - n, 2));
    else
      r_l = r_r = r;
    endif

    ## let's find the maximum value of M(w,u,v) when: w is a leaf of the left
    ## subtree of v and u is a leaf of the right subtree of v
    for i = 1 : length (l)
      if (isempty (find (l(i) == l_l)))
        x = l_l;
      else
        x = l_r;
      endif
      for j = 1 : length (r)
        if (isempty (find (r(j) == r_l)))
          y = r_l;
        else
          y = r_r;
        endif

        ## max(M(w,u,v)) = max(M(w,k,v_l)) + max(M(h,u,v_r)) + S(k,h)
        ## where: v_l is the left child of v and v_r the right child of v
        M_tmp = repmat (M(l(i), x(:), tree(iter, 1)), length (y), 1) + ...
                repmat (M(y(:), r(j), tree(iter, 2)), 1, length (x)) + ...
                S(y(:), x(:));
        M_max = max (max (M_tmp));      # this is M(l(i), r(j), v)
        [h, k] = find (M_tmp == M_max);

        M(l(i), r(j), v) = M_max;
        M(r(j), l(i), v) = M(l(i), r(j), v);

        if (M_max > O(v,3))
          O(v, 1) = l(i);               # this is w
          O(v, 2) = r(j);               # this is u
          O(v, 3) = M_max;              # this is M(w, u, v)
        endif
      endfor
    endfor
  endfor

  ## reordering:
  ## we found the M(w,u,v) corresponding to the optimal leaf order, now we can
  ## compute the optimal leaf order given our M(w,u,v)

  ## the return value
  leafOrder = zeros ( 1, n );
  leafOrder(1) = O(end, 1);
  leafOrder(n) = O(end, 2);

  ## the inverse operation, only easier, to get the leaf order: now we know the
  ## leftmost and rightmost leaves of the best subtree, we may have to flip it
  ## though
  for iter = m : -1 : 1
    v = iter + n;

    extremes = O(v, [1, 2]);

    l_node = tree(iter, 1);
    r_node = tree(iter, 2);

    l = optimalleaforder_getLeafList (l_node);
    r = optimalleaforder_getLeafList (r_node);

    if (l_node > n)
      l_l = optimalleaforder_getLeafList (tree(l_node - n, 1));
      l_r = optimalleaforder_getLeafList (tree(l_node - n, 2));
    else
      l_l = l_r = l;
    endif

    if (r_node > n)
      r_l = optimalleaforder_getLeafList (tree(r_node - n, 1));
      r_r = optimalleaforder_getLeafList (tree(r_node - n, 2));
    else
      r_l = r_r = r;
    endif

    ## this means that we need to flip the subtree
    if (isempty (find (extremes(1) == l)))
      l_tmp = l;
      l_l_tmp = l_l;
      l_r_tmp = l_r;

      l = r;
      l_l = r_l;
      l_r = r_r;

      r = l_tmp;
      r_l = l_l_tmp;
      r_r = l_r_tmp;

      node_tmp = l_node;
      l_node = r_node;
      r_node = node_tmp;
    endif

    if (isempty (find (extremes(1) == l_l)))
      x = l_l;
    else
      x = l_r;
    endif

    if (isempty (find (extremes(2) == r_l)))
      y = r_l;
    else
      y = r_r;
    endif

    M_tmp = repmat (M(extremes(1), x(:), l_node), length (y), 1) + ...
            repmat (M(y(:), extremes(2), r_node), 1, length (x)) + ...
            S(y(:), x(:));
    M_max = max (max (M_tmp));
    [h, k] = find (M_tmp == M_max);

    O(l_node, 1) = extremes(1);
    O(l_node, 2) = x(k);
    O(r_node, 1) = y(h);
    O(r_node, 2) = extremes(2);

    p_1 = find (leafOrder == extremes(1));
    p_2 = find (leafOrder == extremes(2));

    leafOrder(p_1 + (length (l)) - 1) = x(k);
    leafOrder(p_1 + (length (l))) = y(h);
  endfor

  ## function: optimalleaforder_getLeafList
  ## get the list of leaves under a given node
  function vector = optimalleaforder_getLeafList (nodes_to_visit)
    vector = [];
    while (! isempty (nodes_to_visit))
      currentnode = nodes_to_visit(1);
      nodes_to_visit(1) = [];
      if (currentnode > n)
        node = currentnode - n;
        nodes_to_visit = [tree(node, [2 1]) nodes_to_visit];
      endif

      if (currentnode <= n)
        vector = [vector currentnode];
      endif
    endwhile
  endfunction

endfunction

%!demo
%! randn ('seed', 5)  # for reproducibility
%! X = randn (10, 2);
%! D = pdist (X);
%! tree = linkage (D, 'average');
%! optimalleaforder (tree, D, 'Transformation', 'linear')

## Test input validation
%!error optimalleaforder ()
%!error optimalleaforder (1)
%!error <tree must be .*> optimalleaforder (ones (2, 2), 1)
%!error <character inputs expected> optimalleaforder ([1 2 3], [1 2; 3 4], 'criteria', 5)
%!error <D must be .*> optimalleaforder ([1 2 1], [1 2 3])
%!error <unknown property .*> optimalleaforder ([1 2 1], 1, 'xxx', 'xxx')
%!error optimalleaforder ([1 2 1], 1, 'Transformation', 'xxx')

## A leaf ordering and its reverse are equally optimal, so the tests below
## compare against MATLAB R2024a up to reversal.
%!shared X, D, Z
%! X = [0, 0; 0.2, 0.1; 3, 3; 3.1, 3.4; 6, 0; 6.2, 0.3; 9, 9; 9.1, 9.2; ...
%!      1, 5; 5, 1];
%! D = pdist (X);
%! Z = linkage (D, "average");

%!test  # a custom transformation is a documented input and now reaches the fit
%! o = optimalleaforder (Z, D, "Transformation", @(x) 1 ./ (1 + x));
%! assert_equal (sort (o), 1:10);
%! assert_equal (o(end:-1:1), [1, 2, 9, 4, 3, 10, 6, 5, 7, 8]);

%!test  # the named transformations are unchanged
%! o = optimalleaforder (Z, D, "Transformation", "linear");
%! assert_equal (o(end:-1:1), [1, 2, 9, 4, 3, 10, 5, 6, 7, 8]);
%! o = optimalleaforder (Z, D, "Transformation", "inverse");
%! assert_equal (o(end:-1:1), [1, 2, 9, 4, 3, 10, 6, 5, 7, 8]);

%!test  # two leaves have a single ordering, where the tables used to run off
%! ## the end of the similarity array
%! for n = 2:5
%!   o = optimalleaforder (linkage (pdist (X(1:n,:))), pdist (X(1:n,:)));
%!   assert_equal (sort (o), 1:n);
%! endfor
%! assert_equal (optimalleaforder (linkage (pdist (X(1:2,:))), ...
%!                                 pdist (X(1:2,:))), [1, 2]);

%!test  # the ordering is optimal: no reordering that keeps the tree beats it
%! o = optimalleaforder (Z, D);
%! Dsq = squareform (D);
%! obj = @(p) sum (Dsq(sub2ind (size (Dsq), p(1:end-1), p(2:end))));
%! assert_equal (obj (o), obj (o(end:-1:1)), 1e-12);
%! assert_equal (obj (o) <= obj (1:10), true);

%!error <character inputs expected> ...
%! optimalleaforder (Z, D, @(x) x, "linear")

%!error <character inputs expected for CRITERIA> ...
%! optimalleaforder ([1 2 1], 1, 'Criteria', 5)
%!error <TRANSFORMATION must be a character vector or a function handle> ...
%! optimalleaforder ([1 2 1], 1, 'Transformation', 5)
