library(ape)

###################################################################
###################################################################
###################################################################
###
### This function calculates the disparity (measured as either the 
### range or the variance) of a continuous trait through time.
### Note if you are using multiple traits e.g. principal components
### the function will have to be run on each trait separately
### and the variances/ranges may be summed.
### The function allows the data from the tips, nodes and ghost 
### lineages present in each time bin rather than just the tips
### to calculate disparity in the manner described in Brocklehurst
### (2017).
### Arguments:
### tree - a time calibrated tree
### trait - a vector containing the trait values and names identical
###		to the tip labels of the tree
### ages - a two column matrix giving the age ranges of the tips.
###		column 1 should be first appearence, column 2 should be
###		last appearence. rownames should be the tip labels of the
###		tree. Only needed if term.sep=T
### time - a two column matrix giving the start time and end time 
###		of the time bins
### mode - either "BM" or "punctuated". BM assumes gradual evolution
### 		of the trait, with no variation in rate. "punctuated" 
###		assumes the morphological change occurs at the 
###		speciation event
### term.sep - True or False. True - the trait value of each taxon  
###			will be treated as the morphology at the taxon's 
###			first appearence. No morphological change will be
###			inferred to have occured during the taxon's range.
###			False - the trait value will be treated as the 
###			morphology at the taxon's last appearence. Only
###			makes a difference if mode="BM"
### anc - True or False. False - the morphologies at each node will
###		be inferred automatically using likelihood assuming 
###		evolution by Brownian Motion. True - the user can input
###		their own ancestral morphologies. These should be 
###		included in the "trait" vector after the tip values, in
###		the node order of the tree
### metric - measure of disparity: "variance" or "range"
###
### The function will output a vector giving the variance or range
### of the trait in each time bin.
###
###################################################################
###################################################################
###################################################################

phylo.disp<-function(tree,trait,ages,time,mode,term.sep,anc,metric)
{
	if(mode=="BM")
	{
		if(anc!=T)
		{
			anc.trait<-ace(trait,tree)$ace
			trait<-c(trait[tree$tip.label],anc.trait)
		}
		else(trait[1:length(tree$tip.label)]<-trait[tree$tip.label])
		internal.edges<-tree$edge
		rownames(internal.edges)<-1:nrow(internal.edges)
		if(term.sep==T)
		{
			species.edges<-matrix(nrow=length(tree$tip.label),ncol=2)
			rownames(species.edges)<-tree$tip.label
			species.edges[,1]<-1:length(tree$tip.label)
			species.edges[,2]<-1:length(tree$tip.label)
			edges<-rbind(internal.edges,species.edges)
		}
		else(edges<-internal.edges)
		node.ages<-tree$edge
		nodes<-dateNodes(tree)
		for(i in 1:nrow(internal.edges))
		{
			node.ages[i,1]<-nodes[node.ages[i,1]]
			node.ages[i,2]<-nodes[node.ages[i,2]]
		}
		rownames(node.ages)<-1:nrow(node.ages)
		if(term.sep==T)
		{
			edge.ages<-rbind(node.ages,ages)
		}
		else(edge.ages<-node.ages)
		trait.start<-array(dim=nrow(edges))
		rownames(trait.start)<-rownames(edges)
		for(i in 1:length(trait.start))
		{
			trait.start[i]<-trait[edges[i,1]]
		}
		trait.end<-array(dim=nrow(edges))
		rownames(trait.end)<-rownames(edges)
		for(i in 1:length(trait.end))
		{
			trait.end[i]<-trait[edges[i,2]]
		}
	}
	if(mode=="punctuated")
	{
		if(anc==F)
		{
			punc.tree<-tree
			punc.tree$edge.length<-rep(1,length(tree$edge.length))
			anc.trait<-ace(trait,punc.tree)$ace
			trait<-c(trait[tree$tip.label],anc.trait)
		}
		else(trait[1:length(tree$tip.label)]<-trait[tree$tip.label])
		edges<-tree$edge
		rownames(edges)<-1:nrow(edges)
		edge.ages<-tree$edge
		nodes<-dateNodes(tree)
		for(i in 1:nrow(edges))
		{
			edge.ages[i,1]<-nodes[edge.ages[i,1]]
			edge.ages[i,2]<-nodes[edge.ages[i,2]]
		}
		rownames(edge.ages)<-1:nrow(edge.ages)
		trait.end<-array(dim=nrow(edges))
		rownames(trait.end)<-rownames(edges)
		for(i in 1:length(trait.end))
		{
			trait.end[i]<-trait[edges[i,2]]
		}
	}
	disparity<-array(dim=nrow(time),data=0)
	for(i in 1:nrow(time))
	{
		present<-vector(length=0)
		start.time<-time[i,1]
		end.time<-time[i,2]
		for(j in 1:nrow(edge.ages))
		{
			if (edge.ages[j,1]>time[i,2] && edge.ages[j,2]<time[i,1])
			{
				present<-c(present,rownames(edge.ages)[j])
			}
		}
		bin.trait<-array(dim=length(present))
		for(j in 1:length(present))
		{
			if(mode=="BM")
			{
				start.lineage<-edge.ages[present[j],1]
				end.lineage<-edge.ages[present[j],2]
				start.trait<-trait.start[present[j]]
				end.trait<-trait.end[present[j]]
				rate<-(end.trait-start.trait)/(start.lineage-end.lineage)
				rate[is.na(rate)]<-0
				end<-max(c(time[i,2],end.lineage))
				time.to.end<-start.lineage-end
				bin.trait[j]<-start.trait+(time.to.end*rate)
			}
			if(mode=="punctuated")
			{
				bin.trait[j]<-trait.end[present[j]]
			}
		}
		if(metric=="variance")
		{
			disparity[i]<-var(bin.trait)
		}
		else (disparity[i]<-max(bin.trait)-min(bin.trait))		
	}
	names(disparity)<-rownames(time)
	return(disparity)	
}