# R code from www.graemetlloyd.com
#
# The following implements the patristic dissimilarity measure of Wagner (1997, 
# Paleobiology, 23, 115-150) and dates phylogenetic trees using the standard 
# palaeontological, Ruta et al. (2006, PRSB, 273, 2107-2111) or the Brusatte 
# et al. (2008, Science, 321, 1485-1488) method.
# 
# The functions can be fed a phylogenetic tree ("tree"), a vector of the number of missing 
# states ("comp"), the total number of characters ("nchar"), a vector of first appearance
# dates ("ages"), a root length in millions of years ("rlen"), a method choice ("method") 
# and a patristic dissimilarity tree ("ptree").

pat.dist.phylo<-function(tree, comp, nchar)
{
  require(ape)
  nt<-Ntip(tree)
  cchar<-nchar-comp[tree$tip.label,1]
  for (i in 1:length(tree$edge[,1])) {
    if (tree$edge[i,2] <= nt)  tree$edge.length[i]<-tree$edge.length[i]/cchar[tree$edge[i,2]]
    if (tree$edge[i,2] > nt)  tree$edge.length[i]<-tree$edge.length[i]/nchar
  }
  tree
}

date.phylo<-function(tree, ages, rlen, method="equal", ptree)
{
  require(ape)
  mm<-mrca(tree)
  mm<-mm[order(colnames(mm)),]
  mm<-mm[,order(colnames(mm))]
  nt<-1:Ntip(tree)
  am<-matrix(0,nrow=Nnode(tree),ncol=max(nt))
  rownames(am)<-(max(nt)+1):(Nnode(tree)+max(nt))
  colnames(am)<-colnames(mm)
  for (i in 1:length(am[1,])) am[match(sort(unique(mm[i,]))[-1],rownames(am)),i]<-1
  for (i in 1:length(ages[,1])) am[grep(1,am[,i]),i]<-ages[colnames(am)[i],1]
  tagf<-taf<-nda<-vector(mode="numeric")
  for (i in 1:length(am[,1])) nda[i]<-max(am[i,])
  names(nda)<-rownames(am)
  for (i in nt) taf[i]<-ages[match(tree$tip.label[i],rownames(ages)),1]
  names(taf)<-nt
  for (i in nt) {
    tagf[i]<-nda[as.character(tree$edge[grep(TRUE,(tree$edge[,2] <= max(nt)))[i],1])]
    names(tagf)[i]<-tree$tip.label[tree$edge[grep(TRUE,(tree$edge[,2] <= max(nt)))[i],2]]
  }
  anda<-c(taf,nda)
  fa<-anda[tree$edge[,1]]
  ta<-anda[tree$edge[,2]]
  anda[max(nt)+1]<-anda[max(nt)+1]+rlen
  bls<-anda[tree$edge[,1]]-anda[tree$edge[,2]]
  ttree<-tree
  ttree$edge.length<-bls
  if (method != "standard") {
    bro<-cbind(ttree$edge[,2],node.depth(ttree)[ttree$edge[,2]])
    rownames(bro)<-1:length(bro[,1])
    bro<-bro[order(bro[,2]),]
    bro<-as.numeric(rownames(bro))
    for (i in 1:length(bro)) {
      i<-bro[i]
      if (ttree$edge.length[i] == 0) {
        brs<-vector(mode="numeric")
        brs[length(brs)+1]<-ttree$edge[i,2]
        while (ttree$edge.length[match(ttree$edge[i,1],ttree$edge[,2])] == 0) {
          i<-match(ttree$edge[i,1],ttree$edge[,2])
          brs[length(brs)+1]<-ttree$edge[i,2]
        }
        brs[length(brs)+1]<-ttree$edge[match(brs[length(brs)],ttree$edge[,2]),1]
        tt<-sum(ttree$edge.length[match(brs,ttree$edge[,2])])
        if (method == "equal") {
          ppt<-1:length(brs)
          for (j in 1:length(ppt)) ppt[j]<-1/length(brs)
        }
        if (method == "ruta") ppt<-ptree$edge.length[match(brs,ttree$edge[,2])]/sum(ptree$edge.length[match(brs,ttree$edge[,2])])
        tc<-ppt*tt
        for (j in 2:length(tc)) tc[j]<-tc[j]+tc[j-1]
        anda[ttree$edge[match(brs,ttree$edge[,2]),1]][1:length(tc)-1]<-anda[ttree$edge[match(brs,ttree$edge[,2]),1]][1:length(tc)-1]+tc[1:length(tc)-1]
        ttree$edge.length<-anda[tree$edge[,1]]-anda[tree$edge[,2]]
      }
    }
  }
  ttree
}