# Copyright (C) Tal Galili
#
# This file is part of dendextend.
#
# dendextend 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 2 of the License, or
# (at your option) any later version.
#
# dendextend 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.
#
#  A copy of the GNU General Public License is available at
#  http://www.r-project.org/Licenses/
#





#' @title Counts the number of leaves in a tree
#' @rdname nleaves
#' @export
#' @description
#' Counts the number of leaves in a tree (dendrogram or hclust).
#'
#' @param x tree object (dendrogram/hclust/phylo,\link{dendlist})
#' @param method a character scalar (default is "members"). If "order"
#' than nleaves is based on length of \link{order.dendrogram}.
#' If "members", than length is trusting what is written in the
#' dendrogram's root \link{attr}.
#' "members" is about 4 times faster than "order".
#'
#' @param ... not used
#' @details
#' The idea for the name is from functions like ncol, and nrow.
#'
#' Also, it is worth noting that the nleaves.dendrogram is based on
#' order.dendrogram instead of labels.dendrogram since the first is
#' MUCH faster than the later.
#'
#' The phylo method is based on turning the phylo to hclust and than to
#' dendrogram. It may not work for complex phylo trees.
#'
#' @return The number of leaves in the tree
#' @seealso \link{nrow}, \link{count_terminal_nodes}
#' @examples
#' hc <- hclust(dist(USArrests[1:5, ]), "ave")
#' dend <- as.dendrogram(hc)
#'
#' nleaves(dend) # 5
#' nleaves(hc) # 5
nleaves <- function(x, ...) {
  UseMethod("nleaves")
}

#' @export
#' @rdname nleaves
nleaves.default <- function(x, ...) {
  stop("object x must be a dendrogram/hclust/phylo object")
}

#' @export
#' @rdname nleaves
nleaves.dendrogram <- function(x, method = c("members", "order"), ...) {
  if (method[1] == "members") {
    return(as.integer(attr(x, "members")))
  } else {
    # "order"
    return(length(order.dendrogram(x)))
  }
}


#' @export
#' @rdname nleaves
nleaves.dendlist <- function(x, ...) {
  sapply(x, nleaves, ...)
}



# library(microbenchmark)
# big_dend <- USArrests %>% dist %>% hclust %>% as.dendrogram
# microbenchmark(
#    nleaves.dendrogram(big_dend, "order"),
#    nleaves.dendrogram(big_dend, "members")
# )


# TODO: add count by number of labels/ or by the number of is.leaf==TRUE


#' @export
#' @rdname nleaves
nleaves.hclust <- function(x, ...) {
  length(x$order)
}


#' @export
#' @rdname nleaves
nleaves.phylo <- function(x, ...) {
  # nleaves(as.dendrogram(x))
  length(labels(x))
}

# TODO: there is probably a better way for getting the tree size for a phylo object.




#' @title Counts the number of nodes (Vertices) in a tree
#' @export
#' @rdname nnodes
#' @description Counts the number of nodes in a tree (dendrogram, hclust, phylo).
#'
#' @param x tree object (dendrogram or hclust)
#' @param ... not used
#' @details
#' The idea for the name is from functions like ncol, and nrow.
#'
#' The phylo method is based on turning the phylo to hclust and than to
#' dendrogram. It may not work for complex phylo trees.
#'
#' @return The number of leaves in the tree
#' @seealso \link{nrow}, \link{count_terminal_nodes}, \link{nleaves}
#' @examples
#' hc <- hclust(dist(USArrests[1:5, ]), "ave")
#' dend <- as.dendrogram(hc)
#'
#' nnodes(dend) # 9
#' nnodes(hc) # 9
nnodes <- function(x, ...) {
  UseMethod("nnodes")
}

#' @export
#' @rdname nnodes
nnodes.default <- function(x, ...) {
  stop("object x must be a dendrogram/hclust/phylo object")
}

#' @export
#' @rdname nnodes
nnodes.dendrogram <- function(x, ...) {
  if (!inherits(x, "dendrogram")) warning("'object' should be a dendrogram.")

  i_node_counter <- 0

  go_through_nodes <- function(dend_node) {
    i_node_counter <<- i_node_counter + 1
  }

  dendrapply(x, go_through_nodes)

  return(i_node_counter)
}
# nnodes(dend)


#' @export
#' @rdname nnodes
nnodes.hclust <- function(x, ...) {
  nnodes(as.dendrogram(x))
}

#' @export
#' @rdname nnodes
nnodes.phylo <- function(x, ...) {
  nnodes(as.dendrogram(x))
}





#' @title Counts the number of terminal nodes (merging 0 nodes!)
#' @export
#' @description This function counts the number of "practical" terminal nodes (nodes which are not leaves, but has 0 height to them are considered "terminal" nodes).
#' If the tree is standard, that would simply be the number of leaves (only the leaves will have height 0).
#' However, in cases where the tree has several nodes (before the leaves) with 0 height,
#' the count_terminal_nodes counts such nodes as terminal nodes
#'
#' The function is recursive in that it either returns 1 if it reached a terminal node (either a leaf or a 0 height node),
#' else: it will count the number of terminal nodes in each of its sub-nodes, sum them up, and return them.
#'
#' @param dend_node a dendrogram object for which to count its number of terminal nodes (leaves or 0 height nodes).
#' @param ... not used
#' @return The number of terminal nodes (excluding the leaves of nodes of height 0)
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:3, ]), "ave")
#' dend <- as.dendrogram(hc)
#'
#' ###
#' # Trivial case
#' count_terminal_nodes(dend) # 3 terminal nodes
#' length(labels(dend)) # 3 - the same number
#' plot(dend,
#'   main = "This is considered a tree \n with THREE terminal nodes (leaves)"
#' )
#'
#' ###
#' # NON-Trivial case
#' str(dend)
#' attr(dend[[2]], "height") <- 0
#' count_terminal_nodes(dend) # 2 terminal nodes, why? see this plot:
#' plot(dend,
#'   main = "This is considered a tree \n with TWO terminal nodes only"
#' )
#' # while we have 3 leaves, in practice we have only 2 terminal nodes
#' # (this is a feature, not a bug.)
count_terminal_nodes <- function(dend_node, ...) {

  # a recursive function for counting the number of first encountered zero nodes
  # if you are a leaf, or a node with height 0, return 1
  if (is.leaf(dend_node) | attr(dend_node, "height") == 0) {
    return(1L)
  }

  # else : # go through all of the childrens and apply count.first.zero.nodes
  terminal_nodes_counter <- 0
  for (i in seq_len(length(dend_node))) {
    terminal_nodes_counter <- terminal_nodes_counter + # let each sub-node add its counts
      count_terminal_nodes(dend_node[[i]], ...)
  }
  return(terminal_nodes_counter)
}




#' @title unclass an entire dendrogram tree
#' @description unclass all the nodes in a dendrogram tree. (Helps in cases when a dendrapply function was used wrongly)
#' @export
#' @param dend a dendrogram object
#' @param ... not used
#' @return The list which was the dendrogram (but without a class)
#' @seealso \link{nleaves}
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:3, ]), "ave")
#' dend <- as.dendrogram(hc)
#'
#' itself <- function(x) x
#' dend <- dendrapply(dend, itself)
#' unclass(dend) # this only returns a list with
#' # two dendrogram objects inside it.
#' str(dend) # this is a great way to show a dendrogram,
#' # but it doesn't help us understand how the R object is built.
#' str(unclass(dend)) # this is a great way to show a dendrogram,
#' # but it doesn't help us understand how the R object is built.
#' unclass_dend(dend) # this only returns a list
#' #  with two dendrogram objects inside it.
#' str(unclass_dend(dend)) # NOW we can more easily understand
#' # how the dendrogram object is structured...
unclass_dend <- function(dend, ...) {
  # I could have also made this into a method
  # unclass_dendrogram
  # But decided not to, so to allow unclass to work as usual in case someone wants to use it only on one branch.
  #    if(!is.leaf(dend))
  #    {
  #       for(i in seq_len(length(dend)))
  #       {
  #          dend[[i]] <- unclass_dend(dend[[i]])
  #       }
  #    }
  #    dend <- unclass(dend)
  dend <- dendrapply(dend, unclass)
  return(dend)
}



# example(count_terminal_nodes)
# example(labels_colors)




#' @export
head.dendrogram <- function(x, n = 3L, ...) {
  # library(utils)
  utils::str(x, max.leve = n, ...)
  cat("etc...", "\n")
}
# hc <- hclust(dist(USArrests), "ave")
# dend <- as.dendrogram(hc)
# str(dend)
# head(dend)
# Some thoughts on imports: http://obeautifulcode.com/R/How-R-Searches-And-Finds-Stuff/
# and: https://stackoverflow.com/questions/8637993/better-explanation-of-when-to-use-imports-depends

# methods(nleaves)




# This is what should have been done for dendrograms
# order.default <- order
# order <- function(...) UseMethod("order")
