#' Resolving polytomies to non-zero length branches

#usage
#' fix.poly(tree,type=c("collapse","resolve"),node=NULL)

#description
#' The function recognizes instances of clade-associated and/or character-associated diversification occurring on a phylogeny. Clade-associated diversification is meant as significant shift in taxonomic and phenotypic diversfication rates occurring on the same clade. Character-associated diversification is meant as linear relationship between diversification rates and phenotypic values.   
#' tree: a phylogenetic tree. It must be supplied when the rate of phenotypic evolution are to be computed by applying \code{compare.evol.rates} from the package \pkg{geomorph} (see details). The tree needs not to be ultrametric or fully dichotomous, yet polytomies may deeply influence DR values. We suggest to resolve polytomies by using the function \code{fix.poly}.
#' type: either 'collapse' to create or 'resolve' to resolve (fix) a polytomy to a specific node indicated by the argument 'node'
#' node: the node in the tree where a polytomy should be resolved of fixed either

#' The function returns a tree with randomly fixed polytomies or created polytomies (if 'type' is set to 'collapse')


fix.poly<-function(tree,type=c("collapse","resolve"),node=NULL){
  require(ape)
  require(phytools)
  require(RRphylo)
  require(geiger)
  
  tree->treeO
  
  if(is.null(node)==FALSE & length(node)>1){
    sapply(node,function(x) tips(tree,x),simplify=FALSE)->xx
    sapply(node,function(x) length(tips(tree,x)))->lls
    names(xx)<-node
    combn(names(xx),2)->cb
    for(y in 1:ncol(cb)){
      if(lls[match(cb[1,y],names(lls))]>=lls[match(cb[2,y],names(lls))]) cb[,y]->cb[,y] else rev(cb[,y])->cb[,y] 
    } 
    
    outs<-array()
    for(y in 1:ncol(cb)){
      if(length(which(xx[match(cb[2,y],names(xx))][[1]]%in%xx[match(cb[1,y],names(xx))][[1]]))==length(xx[match(cb[2,y],names(xx))][[1]])){
        warning(paste("node", names(xx[match(cb[2,y],names(xx))]), "is nested within", names(xx[match(cb[1,y],names(xx))]), "and will be removed"))
        names(xx[match(cb[2,y],names(xx))])->outs[y]
      } else { 
        outs[y]<-NA
      }
    }
    outs[!is.na(outs)]->outs
    if(length(outs)==0) node->node else node[-which(node==unique(outs))]->node
  }
  
  if(type=="collapse"){
    if(is.null(node)) stop("node must be supplied for type 'collapse' ")
    diag(vcv(tree))->las
    mean(tree$edge.length)/10e6->xme
    for(x in 1:length(node)) tree$edge.length[match(getDescendants(tree,node[x]),tree$edge[,2])]<-xme
    di2multi(tree,xme*1.000001)->treeN
    unique(unlist(sapply(node, function(x) tips(treeO,x))))->focal
    max(diag(vcv(treeN)))-las[focal]->fixer
    suppressWarnings(scaleTree(treeN,tip.ages = fixer)->treeN)
    return(treeN)
    
    
    
  }else{
    if(is.binary(tree))
    {
      if(min(tree$edge.length)==0) di2multi(tree,tol=1e-06)->tree else stop("binary tree provided, no polytomies to resolve")
    } 
    
    tree->treeN
    max(diag(vcv(treeN)))-diag(vcv(treeN))->f2
    if(is.null(node)){
      table(treeN$edge[,1])->tt
      if(any(tt>2)) names(which(tt>2))->nn
    } else {node->nn}
    
    
    which(sapply(nn, function(x) length(which(getDescendants(treeN,x)>Ntip(treeN))))>=2)->checker
    if(length(which(checker==TRUE)>0)) nn[-checker]->nn
    if(length(nn)==0) stop("non polytomous clade selected to resolve")
    
    sapply(nn, function(y) any(getMommy(treeN,as.numeric(y))%in%nn))->nest
    if(length(which(nest==TRUE)>0)) nn[-which(nest==TRUE)]->nn
    
    for(i in 1:length(nn)){
      if(i==1) treeN->xtree else xtree->xtree
      extract.clade(treeN,as.numeric(nn[i]))->tar
      multi2di(tar)->trx
      diag(vcv(trx))->dtar
      max(dtar)-dtar->fixer
      
      names(which(table(tar$edge[,1])==2))->fixnode
      max(dtar)-nodeHeights(tar)[match(fixnode,tar$edge[,2]),2]->nodage
      names(nodage)<-fixnode
      
      trx$edge.length+.01->trx$edge.length
      makeL(trx)->L->Lx
      Lx[,1]<-L[,1]<-0
      xtime<-array(); for(e in 1:Ntip(trx)) (sum(Lx[e,])-length(getMommy(trx,e))*.01)/(length(getMommy(trx,e)))->xtime[e]
      
      for(o in 1:Ntip(trx)){
        Lx[o,][which(Lx[o,]!=0)]<-xtime[o] 
      }
      suppressWarnings(apply(Lx,2,function(x) min(x[-x!=0]))->xtar)
      xtar[-1]->xtar
      
      
      trx$edge.length->egg
      names(egg)<-trx$edge[,2]
      names(egg)[which(as.numeric(names(egg))<=Ntip(trx))]<-trx$tip.label[trx$edge[which(as.numeric(names(egg))<=Ntip(trx)),2]]
      xtar[match(names(egg),names(xtar))]->trx$edge.length 
      geiger::rescale(trx,"depth",max(nodeHeights(tar)))->trx
      
      if(length(nodage)==0)
      {
        suppressWarnings(scaleTree(trx,tip.ages=fixer)->aa)
      }else{
      newn<-array(); for(e in 1:length(nodage)) getMRCA(trx,tips(tar,as.numeric(names(nodage)[e])))->newn[e]
      names(nodage)<-newn
      suppressWarnings(scaleTree(trx,tip.ages=fixer,node.ages = nodage)->aa)
      }
      
      for(e in 1:Nnode(aa)) aa$node.label[e]<-paste("new.node",i,sep = "")
      aa$node.label[1]<-nn[i]
      match(aa$tip.label,tar$tip.label)->m
      n.lab<-array(); for(o in 1:Ntip(aa)) paste("tip",o,sep = "_")->n.lab[o]
      n.lab->aa$tip.label
      
      if(i==1) bind.tree(xtree,aa,as.numeric(nn[i]))->temp else bind.tree(xtree,aa,getMRCA(xtree,tips(treeN,as.numeric(nn[i]))))->temp
      tips(treeN,as.numeric(nn[i]))->nams
      drop.tip(temp,nams)->tree1
      tree1$tip.label[match(aa$tip.label,tree1$tip.label)]<-nams[m]
      geiger::rescale(tree1,"depth",max(nodeHeights(treeN)))->tree1
      suppressWarnings(scaleTree(tree1,tip.ages=f2)->tree1)
      if(is.binary(tree1)==FALSE) multi2di(tree1)->xtree1 else tree1->xtree1
      xtree1->xtree
    }
    return(xtree)
  }
}
