# 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/
#


# like is.infinite but acts differently for character
is.infinite2 <- function(x) {
	if(is.character(x)) { return(x == "Inf") } 
	# else { return(
	is.infinite(x)	
}






### # ' @aliases 
### # ' set_leaves_attr
### # ' @usage
### # ' get_leaves_attr(dend, labels = TRUE,...)
### # ' 
### # ' set_leaves_attr(dend, ...) <- value

#' @title Get/set attributes of dendrogram's leaves
#' @description Get/set attributes of dendrogram's leaves
#' @export
#' @param dend a dendrogram object 
#' @param attribute character scalar of the attribute (\code{attr})
#' we wish to get/set from the leaves
#' @param simplify logical. If TRUE (default), then the return vector is 
#' after using \code{unlist} on it.
#' @param ... not used
#' @source Heavily inspired by the code in the 
#' function \code{labels.dendrogram}, 
#' so credit should go to Martin Maechler.
#' @return 
#' A vector (or a list) with the dendrogram's leaves attribute
#' @seealso \link{get_nodes_attr}, \link{nnodes},
#' \link{nleaves}, \link{assign_values_to_leaves_nodePar}
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:3,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' # get_leaves_attr(dend) # error :)
#' get_leaves_attr(dend, "label")
#' labels(dend, "label")
#' get_leaves_attr(dend, "height") # should be 0's
#' get_nodes_attr(dend, "height") 
#' 
#' get_leaves_attr(dend, "nodePar") 
#'  
#' 
#' get_leaves_attr(dend, "leaf") # should be TRUE's
#' get_nodes_attr(dend, "leaf") # conatins NA's
#' 
#' 
#' get_leaves_attr(dend, "members") # should be 1's
#' get_nodes_attr(dend, "members") # 
#' 
#' 
#' get_leaves_attr(dend, "members",simplify = FALSE) # should be 1's
#' 
#' 
get_leaves_attr <- function (dend, attribute, simplify = TRUE, ...) {
   if(!is.dendrogram(dend)) warning("'dend' should be a dendrogram.")   
   if(missing(attribute)) stop("'attribute' parameter is missing.")
   
   get_attr_from_leaf <- function(dend_node) {
      if(is.leaf(dend_node)) attr(dend_node, attribute)
   }   

   ret <- dendrapply(dend, get_attr_from_leaf)   
   if(simplify) ret <- unlist(ret)   
   
   return(ret)   
}









#' @title Get nodePar of dendrogram's leaves
#' @description Get the nodePar attributes of dendrogram's leaves (includes pch, color, and cex)
#' @export
#' @param dend a dendrogram object 
#' @param simplify logical (default is FALSE). If TRUE, then the return vector is 
#' after using \code{unlist} on it.
#' @param ... not used
#' @return 
#' A list (or a vector) with the dendrogram's leaves nodePar attribute
#' @seealso \link{get_nodes_attr}, \link{assign_values_to_leaves_nodePar}, \link{labels_colors}
#' \link{get_leaves_edgePar}
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:3,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' # get_leaves_attr(dend) # error :)
#' get_leaves_nodePar(dend)
#' labels_colors(dend) <- 1:3
#' get_leaves_nodePar(dend)
#' 
#' dend <- assign_values_to_leaves_nodePar(dend, 2, "lab.cex")
#' get_leaves_nodePar(dend)
#' 
#' plot(dend)
#' 
get_leaves_nodePar <- function (dend, simplify = FALSE, ...) {
   if(!is.dendrogram(dend)) warning("'dend' should be a dendrogram.")   
   
   is_node_leaf <- get_nodes_attr(dend, "leaf")
   is_node_leaf[is.na(is_node_leaf)] <- FALSE
   
   ret <- get_nodes_attr(dend, "nodePar", simplify = FALSE)[is_node_leaf]
   if(simplify) ret <- unlist(ret)   
   
   return(ret)   
}








#' @title Get edgePar of dendrogram's leaves
#' @export
#' @description
#' This is helpful to get the attributes of branches of the leaves.
#' For example, after we use \link{color_branches}, to get the colors
#' of the labels to match (since getting the colors of branches to match
#' those of the labels can be tricky).
#' @param dend a dendrogram object 
#' @param simplify logical (default is FALSE). If TRUE, then the return vector is 
#' after using \code{unlist} on it.
#' @param ... not used
#' @return 
#' A list (or a vector) with the dendrogram's leaves edgePar attribute
#' @seealso \link{get_nodes_attr}, \link{assign_values_to_leaves_nodePar}, \link{labels_colors}
#' \link{get_leaves_nodePar}
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:5,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' # get_leaves_edgePar(dend) # error :)
#' get_leaves_edgePar(dend)
#' dend <- color_branches(dend, k = 3)
#' get_leaves_edgePar(dend)
#' get_leaves_edgePar(dend, TRUE)
#' 
#' dend <- dend %>% set("branches_lwd", c(2,1,2))
#' get_leaves_edgePar(dend)
#' 
#' plot(dend)
#' 
get_leaves_edgePar <- function (dend, simplify = FALSE, ...) {
   if(!is.dendrogram(dend)) warning("'dend' should be a dendrogram.")   
   
   is_node_leaf <- get_nodes_attr(dend, "leaf")
   is_node_leaf[is.na(is_node_leaf)] <- FALSE
   
   ret <- get_nodes_attr(dend, "edgePar", simplify = FALSE)[is_node_leaf]
   if(simplify) ret <- unlist(ret)   
   
   return(ret)   
}






#' @title Get an attribute of the branches of a dendrogram's leaves
#' @export
#' @description
#' This is helpful to get the attributes of branches of the leaves.
#' For example, after we use \link{color_branches}, to get the colors
#' of the labels to match (since getting the colors of branches to match
#' those of the labels can be tricky).
#' This is based on \link{get_leaves_edgePar}.
#' 
#' @param dend a dendrogram object 
#' @param attr character, the attr to get. Can be either "col", "lwd", or "lty".
#' @param ... not used
#' @return 
#' A vector with the dendrogram's leaves nodePar attribute
#' @seealso \link{get_nodes_attr}, \link{assign_values_to_leaves_nodePar}, \link{labels_colors}
#' \link{get_leaves_nodePar}, \link{get_leaves_edgePar}
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:5,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' dend <- dend %>% 
#'          color_branches(k = 3) %>%
#'          set("branches_lwd", c(2,1,2)) %>%
#'          set("branches_lty", c(1,2,1))
#'          
#' plot(dend)
#' 
#' get_leaves_branches_attr(dend, "col")
#' get_leaves_branches_attr(dend, "lwd")
#' get_leaves_branches_attr(dend, "lty")
#' 
#' labels_colors(dend)  <- get_leaves_branches_attr(dend, "col")
#' plot(dend)
#' 
#' 
get_leaves_branches_attr <- function (dend, attr = c("col", "lwd", "lty"), ...) {
   if(!is.dendrogram(dend)) warning("'dend' should be a dendrogram.")   
   attr <- match.arg(attr)
   
   dend_leaves_edgePar <- get_leaves_edgePar(dend)
   
   get_attr <- function(element) element[attr]
   
   unname(unlist(lapply(dend_leaves_edgePar, get_attr))   )
}





#' @title Get the colors of the branches of a dendrogram's leaves
#' @export
#' @description
#' It is useful to get the colors of branches of the leaves,
#' after we use \link{color_branches}, so to then match the colors of the labels
#' to that of the branches (since getting the colors of branches to match
#' those of the labels can be tricky).
#' This is based on \link{get_leaves_branches_attr} which is based on 
#' \link{get_leaves_edgePar}.
#' 
#' TODO: The function get_leaves_branches_col may behave oddly when extracting
#' colors with missing col attributes when the lwd attribute is available. 
#' This may resolt in a vector with the wrong length (with omitted NA values).
#' This might need to be fixed in the future, and attention should be given to this case.
#' 
#' @param dend a dendrogram object 
#' @param ... not used
#' @return 
#' A vector with the dendrogram's leaves' branches' colors
#' @seealso \link{get_nodes_attr}, \link{assign_values_to_leaves_nodePar}, \link{labels_colors}
#' \link{get_leaves_nodePar}, \link{get_leaves_edgePar}, \link{get_leaves_branches_attr}
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:5,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' par(mfrow = c(1,2), mar = c(5,2,1,0))
#' dend <- dend %>% 
#'          color_branches(k = 3) %>%
#'          set("branches_lwd", c(2,1,2)) %>%
#'          set("branches_lty", c(1,2,1)) 
#'          
#' plot(dend)
#' 
#' labels_colors(dend)  <- get_leaves_branches_col(dend)
#' plot(dend)
#' 
#' 
get_leaves_branches_col <- function (dend, ...) {
   get_leaves_branches_attr(dend, attr = "col", ...)
}



#' @title Get attributes of dendrogram's nodes
#' @export
#' @description
#' Allows easy access to attributes of branches and/or leaves, with option
#' of returning a vector with/withough NA's (for marking the missing attr value)
#' 
#' @param dend a dendrogram object 
#' @param attribute character scalar of the attribute (\code{attr})
#' we wish to get from the nodes
#' @param id integer vector. If given - only the attr of these nodes id will be returned (via depth first search)
#' @param include_leaves logical. Should leaves attributes be included as well?
#' @param include_branches logical. Should non-leaf (branch node) 
#' attributes be included as well?
#' @param simplify logical (default is TRUE). should the result be simplified 
#' to a vector (using \link{simplify2array} ) if possible? If it is not possible
#' it will return a matrix. When FALSE, a list is returned.
#' @param na.rm logical. Should NA attributes be REMOVED from the resulting vector?
#' @param ... not used
#' @source Heavily inspired by the code in the 
#' function \code{labels.dendrogram}, 
#' so credit should go to Martin Maechler.
#' @return 
#' A vector with the dendrogram's nodes attribute. If an attribute is missing
#' from some nodes, it will return NA in that vector.
#' @seealso \link{get_leaves_attr}, \link{nnodes}, \link{nleaves}
#' @examples
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:3,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' # get_leaves_attr(dend) # error :)
#' get_leaves_attr(dend, "label")
#' labels(dend, "label")
#' get_leaves_attr(dend, "height") # should be 0's
#' get_nodes_attr(dend, "height") 
#' 
#' 
#' get_leaves_attr(dend, "leaf") # should be TRUE's
#' get_nodes_attr(dend, "leaf") # conatins NA's
#' 
#' 
#' get_leaves_attr(dend, "members") # should be 1's
#' get_nodes_attr(dend, "members", include_branches = FALSE, na.rm = TRUE) # 
#' get_nodes_attr(dend, "members") # 
#' get_nodes_attr(dend, "members",  simplify = FALSE)  
#' get_nodes_attr(dend, "members", include_leaves = FALSE, na.rm = TRUE) # 
#' 
#' get_nodes_attr(dend, "members",  id = c(1,3), simplify = FALSE)  
#' get_nodes_attr(dend, "members", id = c(1,3)) # 
#'
#' 
#' hang_dend <- hang.dendrogram(dend)
#' get_leaves_attr(hang_dend, "height") # no longer 0!
#' get_nodes_attr(hang_dend, "height") # does not include any 0s!
#' 
#' # does not include leaves values:
#' get_nodes_attr(hang_dend, "height", include_leaves = FALSE) 
#' # remove leaves values all together:
#' get_nodes_attr(hang_dend, "height", include_leaves = FALSE, na.rm = TRUE) 
#' 
#' 
#' 
#' \dontrun{
#' library(microbenchmark)
#' # get_leaves_attr is twice faster than get_nodes_attr
#' microbenchmark(   get_leaves_attr(dend, "members"), # should be 1's
#'                     get_nodes_attr(dend, "members", include_branches = FALSE, na.rm = TRUE)
#'                )
#' }
#' 
get_nodes_attr <- function (dend, attribute, 
                            id,
                            include_leaves = TRUE,
                            include_branches = TRUE,
                            simplify = TRUE,
                            na.rm = FALSE, ...) {
   if(!is.dendrogram(dend)) warning("'dend' should be a dendrogram.")   
   if(missing(attribute)) stop("'attribute' parameter is missing.")

   #### for some reason, this doesn't work:   
   #    get_attr_from_node <- function(dend_node) {
   #       i_attr <- attr(dend_node, attribute)
   #       ifelse(is.null(i_attr), NA, i_attr) 
   #    }   
   #    
   #    return((dendrapply(dend, get_attr_from_node)))   
   
   
#    dend_attr <- rep(NA, nnodes(dend))
#   empty_list <- vector("list", nnodes(dend))
   empty_list <- as.list(rep(NA, nnodes(dend)))
   dend_attr <- empty_list
   missing_id  <-  missing(id)   


   # this function is used to modify dend_attr. What it returns is not important.
   i_node <- 0
   get_attr_from_node <- function(dend_node) {
      i_node <<- i_node + 1
      
      # if we have id's and this is not it - we can skip it...
      # FALSE & NULL # fails
      # FALSE && NULL # works...
      if(!missing_id && !(i_node %in% id) ) return(invisible())
      
         
      # if we should not include_leaves, then we skip when a leaf is encountered.
      if(!include_leaves && is.leaf(dend_node)) return(NULL)
      if(!include_branches && !is.leaf(dend_node)) return(NULL)      
      
      i_attr <- attr(dend_node, attribute)
      if(!is.null(i_attr)) dend_attr[[i_node]] <<- i_attr
      return(invisible())
   }   
   dendrapply(dend, get_attr_from_node)   

   # as.vector is to remove all classes of the na.omit
   # thank you Prof. Brian Ripley http://tolstoy.newcastle.edu.au/R/e2/devel/07/01/1965.html
   if(simplify) dend_attr <- simplify2array(dend_attr)

   if(na.rm) dend_attr <- as.vector(na.omit(dend_attr)) 

   
   if(dendextend_options("warn") && identical(dend_attr, simplify2array(empty_list))) warning("It seems that the attribute '", attribute, "' does not exist - returning NA.")

   # TODO: this could probably be more optimized - say, by looking only at the above mentioned id's
   # and not create all of the vector and only then take a subset. 
   # But for now, I think this is more maintainable...
   if(!missing_id) {
      dend_attr <- dend_attr[id]
   }

   return(dend_attr)   
}











#   ' @export
#' @title recursivly apply a function on a list
#' @description
#' recursivly apply a function on a list - and returns the output as a list, 
#' following the naming convention in the {plyr} package
#' the big difference between this and rapply is that this will also apply 
#' the function on EACH element of the list, even if it's not a "terminal node"
#' inside the list tree.
#' An attribute is added to indicate if the value returned is 
#' from a branch or a leaf.
#' @param x a list.
#' @param FUN a function to apply on each element of the list
#' @param add_notation logical. Should each node be 
#' added a "position_type" attribute, stating if it is a "Branch" or a "Leaf".
#' @param ... not used.
#' @return a list with ALL of the nodes (from the original "x" list),
#' that FUN was applied on.
#' 
#' @examples
#' \dontrun{
#' x <- list(1)
#' x
#' rllply(x, function(x){x}, add_notation  = TRUE)
#' 
#' x <- list(1, 2, list(31))
#' x
#' rllply(x, function(x){x}, add_notation  = TRUE)
#'                      # the first element is the entire tree
#'                            # after FUN was applied to its root element.
#' 
#' hc <- hclust(dist(USArrests[1:4,]), "ave")
#' dend <- as.dendrogram(hc)
#' rllply(dend, function(x){attr(x, "height")})
#' rllply(dend, function(x){attr(x, "members")})
#' }
rllply <- function(x, FUN,add_notation = FALSE, ...)
{   
   if(is.list(x)) {
      output <- list()
      for(i in seq_len(length(x)))
      {		
         output[[i]] <- list(rllply(x[[i]], FUN,...))
         if(add_notation) attr(output[[i]][[1]], "position_type") <- "Branch"
      }
      output <- list(FUN(x,...), output)
   } else {
      output <- FUN(x,...)
      if(add_notation) attr(output, "position_type") <- "Leaf"	
   }   
   return(output)
}




# ' @aliases 
# ' old_get_branches_heights


#' @title Get height attributes from a dendrogram
#' @description Get height attributes of a dendrogram's branches
#' @export
#' @aliases
#' dendextend_get_branches_heights
#' @param dend a dendrogram.
#' @param sort logical. Should the heights be sorted?
#' @param decreasing logical. Should the sort be increasing or decreasing? Not available for partial sorting.
#' @param include_leaves logical (FALSE). Should the output include the leaves value (0's).
#' @param ... not used.
#' @return 
#' a vector of the dendrogram's nodes heights (excluding leaves).
#' 
#' @examples
#' 
#' hc <- hclust(dist(USArrests[1:4,]), "ave")
#' dend <- as.dendrogram(hc)
#' get_branches_heights(dend)
#' 
#' \dontrun{
#' dat1 <- iris[1:150,-5]
#' dat1 <- rbind(dat1,dat1,dat1,dat1,dat1,dat1,dat1)
#' dend = as.dendrogram(hclust(dist(dat1)))
#' library(microbenchmark)
#' library(dendextendRcpp)
#' microbenchmark(
#'    dendextendRcpp::dendextendRcpp_get_branches_heights(dend),
#'    old_get_branches_heights(dend,sort=F)
#' )
#' }
#' 
get_branches_heights <- function(dend, sort = TRUE, decreasing = FALSE, include_leaves = FALSE, ...) {
   fo <- dendextend_options("get_branches_heights")
   fo(dend, sort = sort, decreasing = decreasing, include_leaves = include_leaves, ...)
}



#' @export
dendextend_get_branches_heights <- function(dend, sort = TRUE, decreasing = FALSE, include_leaves = FALSE, ...)
{
#    height <- unlist(rllply(dend, function(x){attr(x, "height")}))
#    height <- get_nodes_attr(dend, "height") 
#    height <- height[height != 0] # include only the non zero values
   
   height <- get_nodes_attr(dend, "height", include_leaves = include_leaves, na.rm = TRUE)   
   if(sort) height <- sort(height, decreasing=decreasing) 	# sort the height
   
   return(height)
}












#' @title Hang dendrogram leaves
#' @export
#' @description
#' Adjust the height attr in all of the dendrogram leaves so that
#'  the tree will hang. This is similar to as.dendrogram(hclust, hang=0.1)
#'  Only that it now works on other object than hclust turned into a dendrogram.
#'  For example, this allows us to hang non-binary trees.
#'  
#' @param dend a dendrogram object 
#' @param hang The fraction of the plot height by which labels should hang below 
#' the rest of the plot. A negative value will cause the labels to 
#' hang down from 0.
#' @param hang_height is missing, then using "hang". If a number is given,
#' it overrides "hang" (except if "hang" is negative)
#' @param ... not used
#' @return 
#' A dendrogram, after adjusting the height attr in all of its leaves, 
#' so that the tree will hang.
#' @source 
#' Noticing that as.dendrogram has a "hang" parameter was thanks to Enrique Ramos's answer here:: 
#' \url{http://stackoverflow.com/questions/17088136/plot-horizontal-dendrogram-with-hanging-leaves-r}
#' @examples
#' 
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:5,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' par(mfrow = c(1,2))
#' plot(hang.dendrogram(dend))
#' plot(hc)
#' # identical(as.dendrogram(hc, hang = 0.1), hang.dendrogram(dend, hang = 0.1))
#' # TRUE!!
#' 
#' 
#' par(mfrow = c(1,4))
#' 
#' plot(dend)
#' plot(hang.dendrogram(dend, hang = 0.1))
#' plot(hang.dendrogram(dend, hang = 0))
#' plot(hang.dendrogram(dend, hang = -0.1))
#' 
#' par(mfrow = c(1,1))
#' plot(hang.dendrogram(dend), horiz = TRUE)
#'  
#'  
hang.dendrogram <- function(dend, hang = 0.1, hang_height, ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   

   #    get_heights.dendrogram
   if(missing(hang_height)) hang_height <- attr(dend, "height")*hang
   
   fix_height_attr_per_leaf <- function(dend_node)
   {
      if(!is.leaf(dend_node)) {
         dend_node_height <- attr(dend_node, "height")
         
         for(i_nodes in seq_len(length(dend_node))) {
            if(is.leaf(dend_node[[i_nodes]])) {
               if(hang < 0) {
                  attr(dend_node[[i_nodes]], "height") <- 0
               } else {
                  attr(dend_node[[i_nodes]], "height") <- dend_node_height - hang_height
               }
               
               dend_node[[i_nodes]] <- unclass(dend_node[[i_nodes]]) # makes sure we don't inherent any classes...
               
            } else {
               dend_node[[i_nodes]] <- 
                  fix_height_attr_per_leaf(dend_node[[i_nodes]])
            }           
         }
         
      }
      return(unclass(dend_node))
   }
   
   dend <- fix_height_attr_per_leaf(dend)   
   class(dend) <- "dendrogram"
   return(dend)
}





#' @title Get height attributes from a dendrogram's children
#' @description Get height attributes from a dendrogram's children nodes
#' @export
#' @param dend a dendrogram.
#' @param ... not used.
#' @return 
#' a vector of the heights of a dendrogram's current node's 
#' (first level) children.
#' 
#' @seealso
#' \link{get_branches_heights}
#' @examples
#' 
#' hc <- hclust(dist(USArrests[1:4,]), "ave")
#' dend <- as.dendrogram(hc)
#' get_childrens_heights(dend)
#' 
#' 
get_childrens_heights <- function(dend,...) {
   sapply(dend, function(x) {attr(x, "height")})
}



#' @title Rank branches' heights
#' @export
#' @description
#' Adjust the height attr in all of the dendrogram nodes so that
#'  the tree will have a distance of 1 unit between each parent/child nodes.
#'  It can be thought of as ranking the branches between themselves.
#'  
#'  This is intended for easier comparison of the topology of two trees.
#'  
#'  Notice that this function changes the height of all the leaves into 0, 
#'  thus erasing the effect of \link{hang.dendrogram} (which should be run 
#'  again, if that is the visualization you are intereted in).
#'  
#' @param dend a dendrogram object 
#' @param diff_height Numeric scalar (1). Affects the difference in height
#' between two branches.
#' @param ... not used
#' @seealso
#' \link{get_branches_heights}, \link{get_childrens_heights},
#' \link{hang.dendrogram}, \link{tanglegram}
#' 
#' @return 
#' A dendrogram, after adjusting the height attr in all of its branches.
#' 
#' @examples
#' 
#' # define dendrogram object to play with:
#' dend <- USArrests[1:5,] %>% dist %>% hclust %>% as.dendrogram
#' 
#' par(mfrow = c(1,3))
#' 
#' plot(dend)
#' plot(rank_branches(dend))
#' plot(hang.dendrogram(rank_branches(dend)))
#' 
#'  
rank_branches <- function(dend, diff_height =1, ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   
   fix_height_diff_per_branch <- function(dend_node)
   {
      if(is.leaf(dend_node)) { 
         attr(dend_node, "height") <- 0
         return(unclass(dend_node)) 
      }
      # else:

      # go through and fix all of the nodes beneath this one:
      for(i_nodes in seq_len(length(dend_node))) {
         dend_node[[i_nodes]] <- Recall(dend_node[[i_nodes]])
      }
      
      # now we can fix the current height, by finding the height of 
      # all son-branches and increasing the height of the maximum of them.
      dend_children_heights <- get_childrens_heights(dend_node)
      attr(dend_node, "height") <- max(dend_children_heights) + diff_height
      
      
      return(unclass(dend_node)) 
   }
   
   dend <- fix_height_diff_per_branch(dend)   
   class(dend) <- "dendrogram"
   return(dend)
}
















#' @title Assign values to nodePar of dendrogram's leaves
#' @export
#' @description
#' Go through the dendrogram leaves and updates the values inside its nodePar
#' 
#' If the value has Inf then the value in edgePar will not be changed. 
#' @param dend a dendrogram object 
#' @param value a new value vector for the nodePar attribute. It should be 
#' the same length as the number of leaves in the tree. If not, it will recycle
#' the value and issue a warning.
#' @param nodePar the value inside nodePar to adjust.
#' @param warn logical (default from dendextend_options("warn") is FALSE).
#' Set if warning are to be issued, it is safer to keep this at TRUE,
#' but for keeping the noise down, the default is FALSE.
#' @param ... not used
#' @return 
#' A dendrogram, after adjusting the nodePar attribute in all of its leaves, 
#' @seealso \link{get_leaves_attr}
#' @examples
#' 
#' \dontrun{
#' 
#' dend <- USArrests[1:5,] %>% dist %>% hclust("ave") %>% as.dendrogram
#' 
#' # reproduces "labels_colors<-" 
#' # although it does force us to run through the tree twice, 
#' # hence "labels_colors<-" is better...
#' plot(dend)
#' dend <- assign_values_to_leaves_nodePar(dend=dend, value = c(3,2), nodePar = "lab.col")
#' plot(dend)
#' 
#' dend <- assign_values_to_leaves_nodePar(dend, 1, "pch")
#' plot(dend)
#' # fix the annoying pch=1:
#' dend <- assign_values_to_leaves_nodePar(dend, NA, "pch")
#' plot(dend)
#' # adjust the cex:
#' dend <- assign_values_to_leaves_nodePar(dend, 19, "pch")
#' dend <- assign_values_to_leaves_nodePar(dend, 2, "lab.cex")
#' plot(dend)
#' 
#' str(unclass(dend))
#' 
#' get_leaves_attr(dend, "nodePar", simplify=FALSE)
#' 
#' }
#' 
assign_values_to_leaves_nodePar <- function(dend, value, nodePar, warn = dendextend_options("warn"), ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   if(missing(value)) {
      if(warn) warning("value is missing, returning the dendrogram as is.")
      return(dend)
   }
   
   leaves_length <- nleaves(dend) # length(labels(dend)) # it will be faster to use order.dendrogram than labels...   
   if(leaves_length > length(value)) {
      if(warn) warning("Length of value vector was shorter than the number of leaves - vector value recycled")
      value <- rep(value, length.out = leaves_length)
   }       
   
   set_value_to_leaf <- function(dend_node) {
      if(is.leaf(dend_node)) {			
         i_leaf_number <<- i_leaf_number + 1

		 to_update_attr <- !is.infinite2(value[i_leaf_number])
		 if(to_update_attr) {
         # if(!is.infinite2(value[i_leaf_number])) {
            attr(dend_node, "nodePar")[nodePar] <- list(value[i_leaf_number]) # this way it doesn't erase other nodePar values (if they exist)
         }      
         
         
         if(length(attr(dend_node, "nodePar")) == 0) {
            attr(dend_node, "nodePar") <- NULL # remove nodePar if it is empty
         } else {
            # if we have some nodePar, and we don't have pch - let's make 
            # sure it is NA - so that we don't see that annoying dot.
            if(! ("pch"  %in%  names(attr(dend_node, "nodePar"))) ) {
               attr(dend_node, "nodePar")["pch"] <- list(NA)
            }               
         }
      }
      return(unclass(dend_node))
   }   
   i_leaf_number <- 0
   new_dend <- dendrapply(dend, set_value_to_leaf)
   class(new_dend) <- "dendrogram"
   return(new_dend)
}


# 
#  tanglegram(dend1 , dend2, lab.cex = 2, edge.lwd = 3,margin_inner= 5, type = "t", center = TRUE)
# attr(dend[[1]], "nodePar")[["pch"]]















#' @title Assign values to edgePar of dendrogram's leaves
#' @export
#' @description
#' Go through the dendrogram leaves and updates the values inside its edgePar
#' 
#' If the value has Inf then the value in edgePar will not be changed. 
#' @param dend a dendrogram object 
#' @param value a new value vector for the edgePar attribute. It should be 
#' the same length as the number of leaves in the tree. If not, it will recycle
#' the value and issue a warning.
#' @param edgePar the value inside edgePar to adjust.
#' @param warn logical (default from dendextend_options("warn") is FALSE).
#' Set if warning are to be issued, it is safer to keep this at TRUE,
#' but for keeping the noise down, the default is FALSE.
#' @param ... not used
#' @return 
#' A dendrogram, after adjusting the edgePar attribute in all of its leaves, 
#' @seealso \link{get_leaves_attr}, link{assign_values_to_leaves_nodePar}
#' @examples
#' 
#' \dontrun{
#' 
#' dend <- USArrests[1:5,] %>% dist %>% hclust("ave") %>% as.dendrogram
#' 
#' plot(dend)
#' dend <- assign_values_to_leaves_edgePar(dend=dend, value = c(3,2), edgePar = "col")
#' plot(dend)
#' dend <- assign_values_to_leaves_edgePar(dend=dend, value = c(3,2), edgePar = "lwd")
#' plot(dend)
#' dend <- assign_values_to_leaves_edgePar(dend=dend, value = c(3,2), edgePar = "lty")
#' plot(dend)
#' 
#' get_leaves_attr(dend, "edgePar", simplify=FALSE)
#' 
#' }
#' 
assign_values_to_leaves_edgePar <- function(dend, value, edgePar, warn = dendextend_options("warn"), ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   if(missing(value)) {
      if(warn) warning("value is missing, returning the dendrogram as is.")
      return(dend)
   }
   
   leaves_length <- nleaves(dend) # length(labels(dend)) # it will be faster to use order.dendrogram than labels...   
   if(leaves_length > length(value)) {
      if(warn) warning("Length of value vector was shorter than the number of leaves - vector value recycled")
      value <- rep(value, length.out = leaves_length)
   }       
   
   set_value_to_leaf <- function(dend_node) {
      if(is.leaf(dend_node)) {   		
         i_leaf_number <<- i_leaf_number + 1
         
         to_update_attr <- !is.infinite2(value[i_leaf_number])
         if(to_update_attr) {
            # if(!is.infinite2(value[i_leaf_number])) {
            attr(dend_node, "edgePar")[edgePar] <- list(value[i_leaf_number]) # this way it doesn't erase other edgePar values (if they exist)
         }      
         
         
         if(length(attr(dend_node, "edgePar")) == 0) {
            attr(dend_node, "edgePar") <- NULL # remove edgePar if it is empty
         } else {
            # if we have some edgePar, and we don't have pch - let's make 
            # sure it is NA - so that we don't see that annoying dot.
            if(! ("pch"  %in%  names(attr(dend_node, "edgePar"))) ) {
               attr(dend_node, "edgePar")["pch"] <- list(NA)
            }               
         }
      }
      return(unclass(dend_node))
   }   
   i_leaf_number <- 0
   new_dend <- dendrapply(dend, set_value_to_leaf)
   class(new_dend) <- "dendrogram"
   return(new_dend)
}








#' @title Assign values to nodePar of dendrogram's nodes
#' @export
#' @description
#' Go through the dendrogram nodes and updates the values inside its nodePar
#' 
#' If the value has Inf then the value in edgePar will not be changed. 
#' 
#' @param dend a dendrogram object 
#' @param value a new value vector for the nodePar attribute. It should be 
#' the same length as the number of nodes in the tree. If not, it will recycle
#' the value and issue a warning.
#' @param nodePar the value inside nodePar to adjust. 
#' This may contain components named pch, cex, col, xpd, and/or bg.
#' @param warn logical (default from dendextend_options("warn") is FALSE).
#' Set if warning are to be issued, it is safer to keep this at TRUE,
#' but for keeping the noise down, the default is FALSE.
#' @param ... not used
#' @return 
#' A dendrogram, after adjusting the nodePar attribute in all of its nodes, 
#' @seealso \link{get_leaves_attr}, \link{assign_values_to_leaves_nodePar}
#' @examples
#' 
#' \dontrun{
#' 
#' dend <- USArrests[1:5,] %>% dist %>% hclust("ave") %>% as.dendrogram
#' 
#' # reproduces "labels_colors<-" 
#' # although it does force us to run through the tree twice, 
#' # hence "labels_colors<-" is better...
#' plot(dend)
#' dend2 <- dend %>% 
#'    assign_values_to_nodes_nodePar(value = 19, nodePar = "pch") %>% 
#'    assign_values_to_nodes_nodePar(value = c(1,2), nodePar = "cex")  %>% 
#'    assign_values_to_nodes_nodePar(value = c(2,1), nodePar = "col")
#' plot(dend2)
#' 
#' 
#' ### Making sure this works for NA with character.
#' dend %>% 
#' assign_values_to_nodes_nodePar(value = 19, nodePar = "pch") %>% 
#' assign_values_to_nodes_nodePar(value = c("red", NA), nodePar = "col") -> dend2
#' plot(dend2)
#' 
#' }
#' 
assign_values_to_nodes_nodePar <- function(dend, value, nodePar = c("pch", "cex", "col", "xpd", "bg"), warn = dendextend_options("warn"), ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   if(missing(value)) {
      warning("value is missing, returning the dendrogram as is.")
      return(dend)
   }
   
   nodePar <- match.arg(nodePar)
   
   
   nodes_length <- nnodes(dend) # length(labels(dend)) # it will be faster to use order.dendrogram than labels...   
   if(nodes_length > length(value)) {
      if(warn) warning("Length of value vector was shorter than the number of nodes - vector value recycled")
      value <- rep(value, length.out = nodes_length)
   }       
   
   set_value_to_node <- function(dend_node) {
      i_node_number <<- i_node_number + 1
      
	  to_update_attr <- !is.infinite2(value[i_node_number]) | is.na(value[i_node_number])
	  if(to_update_attr) {
      # if(!is.infinite2(value[i_node_number])) {
         attr(dend_node, "nodePar")[nodePar] <- list(value[i_node_number]) # this way it doesn't erase other nodePar values (if they exist)
      }      
      
      if(length(attr(dend_node, "nodePar")) == 0) {
         attr(dend_node, "nodePar") <- NULL # remove nodePar if it is empty
      } else {
         # if we have some nodePar, and we don't have pch - let's make 
         # sure it is NA - so that we don't see that annoying dot.
         if(! ("pch"  %in%  names(attr(dend_node, "nodePar"))) ) {
            attr(dend_node, "nodePar")["pch"] <- list(NA)
         }               
      }
      return(unclass(dend_node))
   }   
   i_node_number <- 0
   new_dend <- dendrapply(dend, set_value_to_node)
   class(new_dend) <- "dendrogram"
   return(new_dend)
}

















#' @title Assign values to edgePar of dendrogram's branches
#' @export
#' @description
#' Go through the dendrogram branches and updates the values inside its edgePar
#' 
#' If the value has Inf then the value in edgePar will not be changed. 
#' 
#' @param dend a dendrogram object 
#' @param value a new value scalar for the edgePar attribute. 
#' @param edgePar a character indicating the value inside edgePar to adjust.
#' Can be either "col", "lty", or "lwd".
#' @param skip_leaves logical (FALSE) - should the leaves be skipped/ignored?
#' @param warn logical (default from dendextend_options("warn") is FALSE).
#' Set if warning are to be issued, it is safer to keep this at TRUE,
#' but for keeping the noise down, the default is FALSE.
#' @param ... not used
#' @return 
#' A dendrogram, after adjusting the edgePar attribute in all of its branches, 
#' @seealso \link{get_root_branches_attr}
#' @examples
#' 
#' # This failed before - now it works fine. (thanks to Martin Maechler)
#' dend <- 1:2 %>% dist %>% hclust %>% as.dendrogram
#' dend %>% set("branches_lty", 1:2) %>%
#'    set("branches_col", c("topbranch_never_plots", "black", "orange")) %>% 
#'    plot
#' 
#' 
#' \dontrun{
#' 
#' dend <- USArrests[1:5,] %>% dist %>% hclust %>% as.dendrogram
#' plot(dend)
#' dend <- assign_values_to_branches_edgePar(dend=dend, value = 2, edgePar = "lwd")
#' plot(dend)
#' dend <- assign_values_to_branches_edgePar(dend=dend, value = 2, edgePar = "col")
#' plot(dend)
#' dend <- assign_values_to_branches_edgePar(dend=dend, value = "orange", edgePar = "col")
#' plot(dend)
#' dend2 <- assign_values_to_branches_edgePar(dend=dend, value = 2, edgePar = "lty")
#' plot(dend2)
#' 
#' dend2 %>% unclass %>% str
#' 
#' 
#' 
#' 
#' }
#' 
#' 
assign_values_to_branches_edgePar <- function(dend, value, edgePar, skip_leaves = FALSE, warn = dendextend_options("warn"), ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
  
   if(missing(value)) {
      warning("value is missing, returning the dendrogram as is.")
      return(dend)
   }
   
   # if we are skipping leaves, than the number of branches should not include the terminal nodes/leaves!
   n_branches <- nnodes(dend) - ifelse(skip_leaves, nleaves(dend), 0) # length(labels(dend)) # it will be faster to use order.dendrogram than labels...   
   if(n_branches > length(value)) {
      if(warn) warning("Length of value vector was shorter than the number of leaves - vector value recycled")
      value <- rep(value, length.out = n_branches)
   }       
   
   set_value_to_branch <- function(dend_node) {
      # if we ignore leaves, jump the function for leaves:
      if(skip_leaves & is.leaf(dend_node)) return(unclass(dend_node))         
      # else - keep as usual.   
      
      i_node <<- i_node + 1
	  to_update_attr <- !is.infinite2(value[i_node])
      if(to_update_attr) {
         attr(dend_node, "edgePar")[edgePar] <- list(value[i_node]) # [i_leaf_number] # this way it doesn't erase other edgePar values (if they exist)
      }
      if(length(attr(dend_node, "edgePar")) == 0) attr(dend_node, "edgePar") <- NULL # remove edgePar if it is empty
      return(unclass(dend_node))
   }   
   
   i_node <- 0 
   new_dend <- dendrapply(dend, set_value_to_branch)
   class(new_dend) <- "dendrogram"
   return(new_dend)
}






#' @title Remove all edgePar values from a dendrogram's branches
#' @export
#' @description
#' Go through the dendrogram branches and remove its edgePar.
#' @param dend a dendrogram object 
#' @param ... not used
#' @return 
#' A dendrogram, after removing the edgePar attribute in all of its branches, 
#' @seealso \link{get_root_branches_attr}, \link{assign_values_to_branches_edgePar}
#' @examples
#' 
#' \dontrun{
#' 
#' dend <- USArrests[1:5,] %>% dist %>% hclust %>% as.dendrogram
#' dend <- color_branches(dend, 3)
#' par(mfrow = c(1,2))
#' plot(dend)
#' plot(remove_branches_edgePar(dend))
#' 
#' }
#' 
remove_branches_edgePar <- function(dend, ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   remove_edgePar_from_branch <- function(dend_node) {
      attr(dend_node, "edgePar") <- NULL # remove edgePar
      return(unclass(dend_node))
   }   
   new_dend <- dendrapply(dend, remove_edgePar_from_branch)
   class(new_dend) <- "dendrogram"
   return(new_dend)
}




#' @title Remove all nodePar values from a dendrogram's nodes
#' @export
#' @description
#' Go through the dendrogram nodes and remove its nodePar
#' @param dend a dendrogram object 
#' @param ... not used
#' @return 
#' A dendrogram, after removing the nodePar attribute in all of its nodes, 
#' @seealso \link{get_root_branches_attr}, \link{assign_values_to_branches_edgePar}
#' @examples
#' 
#' \dontrun{
#' 
#' dend <- USArrests[1:5,] %>% dist %>% hclust %>% as.dendrogram
#' dend <- color_branches(dend, 3)
#' par(mfrow = c(1,2))
#' plot(dend)
#' plot(remove_branches_edgePar(dend))
#' 
#' }
#' 
remove_nodes_nodePar <- function(dend, ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   remove_nodePar_from_node <- function(dend_node) {
      attr(dend_node, "nodePar") <- NULL # remove nodePar
      return(unclass(dend_node))
   }   
   new_dend <- dendrapply(dend, remove_nodePar_from_node)
   class(new_dend) <- "dendrogram"
   return(new_dend)
}




#' @title Remove all nodePar values from a dendrogram's leaves
#' @export
#' @description
#' Go through the dendrogram leaves and remove its nodePar.
#' @param dend a dendrogram object 
#' @param ... not used
#' @return 
#' A dendrogram, after removing the nodePar attribute in all of its leaves, 
#' @seealso \link{get_leaves_attr}, \link{assign_values_to_leaves_nodePar}
#' @examples
#' 
#' \dontrun{
#' 
#' dend <- USArrests[1:5,] %>% dist %>% hclust %>% as.dendrogram
#' 
#' dend <- color_labels(dend, 3)
#' par(mfrow = c(1,2))
#' plot(dend)
#' plot(remove_leaves_nodePar(dend))
#' 
#' 
#' get_leaves_attr(dend, "nodePar")
#' get_leaves_attr(remove_leaves_nodePar(dend), "nodePar")
#' 
#' 
#' }
#' 
remove_leaves_nodePar <- function(dend, ...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   remove_nodePar_from_leaf <- function(dend_node) {
      if(is.leaf(dend_node)) {   		
         attr(dend_node, "nodePar") <- NULL # remove nodePar if it is empty
      }
      return(unclass(dend_node))
   }   
   new_dend <- dendrapply(dend, remove_nodePar_from_leaf)
   class(new_dend) <- "dendrogram"
   return(new_dend)
}












# nleaves(dend)
# nleaves(dend[[1]])
# nleaves(dend[[2]])


#' @title Fix members attr in a dendrogram
#' @export
#' @description
#' Fix members attr in a dendrogram after (for example), the tree was pruned 
#' or manipulated.
#' @param dend a dendrogram object 
#' @param ... not used
#' @return 
#' A dendrogram, after adjusting the members attr in all of its nodes.
#' @examples
#' 
#' # define dendrogram object to play with:
#' hc <- hclust(dist(USArrests[1:3,]), "ave")
#' dend <- as.dendrogram(hc)
#' # plot(dend)
#' # prune one leaf
#' dend[[2]] <- dend[[2]][[1]]
#' # plot(dend)
#' dend # but it is NO LONGER true that it has 3 members total!
#' fix_members_attr.dendrogram(dend) # it now knows it has only 2 members :)
#' 
#' hc <- hclust(dist(USArrests[1:3,]), "ave")
#' dend <- as.dendrogram(hc)
#' 
#' identical(prune_leaf(dend , "Alaska"), fix_members_attr.dendrogram(prune_leaf(dend , "Alaska")))
#' str(unclass(prune_leaf(dend , "Alaska")))
#' str(unclass(fix_members_attr.dendrogram(prune_leaf(dend , "Alaska"))))
fix_members_attr.dendrogram <- function(dend,...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   
   fix_members_attr_per_node <- function(dend_node)
   {
      if(!is.leaf(dend_node)) attr(dend_node, "members") <- nleaves(dend_node, method="order")
      return(unclass(dend_node))
   }
   # mtrace(".change.label.by.mat")
   dend <- dendrapply(dend, fix_members_attr_per_node)
   class(dend) <- "dendrogram"
   return(dend)
}
#



#' @title Fix rank of leaves order values in a dendrogram
#' @export
#' @description
#' Generally, leaves order value should be a sequence of integer values.
#' From 1 to nleaves(dend). 
#' This function fixes trees by using \link{rank} on existing leaves order
#' values.
#' @param dend a dendrogram object 
#' @param ... not used
#' @return 
#' A dendrogram, after fixing its leaves order values.
#' @seealso
#' \link{prune}
#' @examples
#' 
#' # define dendrogram object to play with:
#' dend <- USArrests[1:4,] %>% dist %>% hclust(method = "ave") %>% as.dendrogram
#' # plot(dend)
#' order.dendrogram(dend)
#' dend2 <- prune(dend, "Alaska")
#' order.dendrogram(dend2)
#' order.dendrogram(rank_order.dendrogram(dend2))
#' 
rank_order.dendrogram <- function(dend,...) {
   if(!is.dendrogram(dend)) stop("'dend' should be a dendrogram.")   
   
   order.dendrogram(dend) <- rank(order.dendrogram(dend), ties.method = "first") 
   
   return(dend)
}
#



