# Library of (mostly crap) functions I wrote that might be useful for brach_attack project
# by David F. Wright

# some functions require the package 'paleotree' by D.W. Bapst


# Function to get posterior probabilities of a taxon being a sampled ancestor
# Note the the "trees" argument must be a post-burn-in sample of the posterior distribution
samp.Anc <- function(trees){
  
  freq <- vector(mode = "numeric", length = Ntip(trees[[1]]))
  names(freq) <- trees[[1]]$tip.label
  
  for (i in 1:length(trees)){
    
    tip.br <- setNames(trees[[i]]$edge.length[sapply(1:length(trees[[i]]$tip.label),function(x,y)   which(y == x),y = trees[[i]]$edge[,2])],trees[[i]]$tip.label)
    
    for (j in 1:length(tip.br)){
      
      if (tip.br[j] == 0){
        freq[j] <- freq[j] + 1
        
      }
      
    }
    
  }
  
  anc.prop <- (freq / length(trees))
  return(anc.prop)
  
}

# FUNCTION TO CHANGE BRANCH LENGTHS BELOW A THRESHOLD TO A CONSTANT  
# Can take either a list of trees or a single tree. Can also be used to scale branches to unit length when NULL (e.g., a parsimony tree estimated in PHANGORN, etc.)

modify.edges <- function(trees, blens = 1, threshold = 1e-01){
  # Check whether object 'trees' is a single tree or list of trees
  # For multi-phylo objects
  
  if (class(trees) == "multiPhylo"){
    
    for(i in 1:length(trees)){  # loop over all trees a list of trees
      
      tree <- trees[[i]]
      
      # check to make sure there are branch lengths; otherwise add small value to get rid of edge.length == NULL warnings()
      if (is.null(tree$edge.length)){
        N.edges <- (Ntip(tree) - 3) + Ntip(tree) # == number of internal branches + number of tips
        
        for (j in 1:N.edges){
          
          tree$edge.length[j] <- 1e-02
          
        }
        
      }
      
      # Main loop to add edge.lengths below threshold for each jth tree
      for (j in 1:length(tree$edge.length)){
        
        edge <- tree$edge.length[j]
        
        if (edge <= threshold){
          tree$edge.length[j] <- blens
          
        }
        
      }
      
      trees[[i]] <- tree
      
    }
    
    return(trees) # returns a list of trees with modified edges
    
  }
  
  # For a single tree
  if (class(trees) == "phylo"){
    tree <- trees
    
    # check to make sure there are branch lengths; otherwise add small value to get rid of edge.length == NULL warnings()
    if (is.null(tree$edge.length)){
      N.edges <- (Ntip(tree) - 3) + Ntip(tree) # == number of internal branches + number of tips
      
      for (j in 1:N.edges){
        
        tree$edge.length[j] <- 1e-02
        
      }
      
    }
    
    # Main loop to add edge.lengths below threshold 
    for (k in 1:length(tree$edge.length)){
      
      edge <- tree$edge.length[k]
      
      if (edge <= threshold){
        tree$edge.length[k] <- blens
        
      }
      
    }
    
    return(tree) # returns a single tree with modified edges
    
  }
  
}


# Convert MrBayes estimates of diversification dynamics & fossilization to speciation, extinction, and sampling rates

calcMu <- function(netDiv, turnover){
  
  mu <- (netDiv * turnover) / (1 - turnover)
  return(mu)
  
}

calcLambda <- function(netDiv, mu){
  
  lambda <- netDiv + mu
  return(lambda)
  
}

calcPsi <- function(mu, relSamp){
  
  Psi <- (mu * relSamp) / (1 - relSamp)
  return(Psi)
  
}


# Function to add length to Zero Length Branches
removeZLB <- function(trees, blens = 1){
  
  for(i in 1:length(trees)){
    
    tree <- trees[[i]]
    
    for (j in 1:length(tree$edge.length)){
      
      edge <- tree$edge.length[j]
      
      if (edge <= 10e-6){
        tree$edge.length[j] <- blens
        
      }
      
    }
    
    trees[[i]] <- tree
    
  }
  
  return(trees)
  
}


# Sample MrBayes trees from the posterior and time scale them in units of relative time (must use anchorTree function to scale to absolute time)
sample_timescaleMrBayes <- function(t.file, p.file, subsample = 100, burnin = 0.35){ 
 
  n <- length(t.file)
  clockrate <- p.file$clockrate
  samples <- sample(seq(from = round(n*burnin), to=n), size = subsample)
  scaled.trees <- list()
  
     for (i in 1:length(samples)) {
  
        tree <- samples[i]
        foo <- t.file[[tree]]
        foo$edge.length <-  foo$edge.length / clockrate[tree] 
        scaled.trees[[i]] <- foo
  
     }
 
  return(scaled.trees)

}



# Scales all trees in the t.file to relative time. Use anchorTree to scale the tree to absolute time after using this function.
timescaleMrBayes <- function(t.file, p.file){
  
  for (i in 1:length(t.file)){
  
    clock.rate <- p.file$clockrate[i]
    t.file[[i]]$edge.length <- t.file[[i]]$edge.length / clock.rate
  
  }
 
  return(t.file)

}

# Find a tree in the posterior distribution t.file. May be useful to get associated info from p.file for a given tree (e.g., MCCT). May be slow if #trees is large.
findTree <- function(tree, t.file){
  
  clades <- vector(length = length(t.file))
 
   for (i in 1:length(t.file)){
    
     clades[i] <- dist.topo(t.file[[i]], tree)
  
   }
  
  return(which(clades == 0)) 

} 



# Get min taxon age to anchor tree
getMinTaxonAge <- function(Taxon){
  if (Taxon == "Hesperorthis"){ 
    age <- 443.8
  }
  if (Taxon == "Mimella"){ 
    age <- 445.2
  }
  if (Taxon == "Oepikina"){ 
    age <- 450.3
  }
  return(age)
}

# get distribution for the age of the age of the clade

getCladeAge <- function(trees){
  
  age <- c()
  
  for (i in 1:length(trees)){
    
    age[i] <- max(dateNodes(trees[[i]]))
    
  }
  
  return(age)
  
}





  # Use anchor taxon to scale the tree to abosulte time. This only makes sense if trees are already
  # scaled to relative time
  anchorTree <- function(trees, anchorTime){
    # for a list of trees
    if (class(trees) == "multiPhylo"){
      for (i in 1:length(trees)){
        oldRootAge <- max(dateNodes(trees[[i]]))
        trees[[i]]$root.time <- oldRootAge + anchorTime 
      }
      return(trees)
    }
    #for a single tree
    if (class(trees) == "phylo"){
      oldRootAge <- max(dateNodes(trees))
      trees$root.time <- oldRootAge + anchorTime
    }
    return(trees)
}