##############################################################################
##############################################################################
####
#### Function simulates a biogeographic history of a clade over a phylogeny,
#### using the possible events: dispersal (movment of a lineage from one or 
#### more areas to another), vicariance (bifurcation of one lineage covering 
#### multiple areas into two lineages each covering a different subset of the
#### original range) and local extinction (removal of a lineage covering more 
#### than one area from one of those areas). Counts through time of each are
#### output, along with the geographic ranges of the nodes and tips
####
#### Note 1) Regions are treated as discrete entities, and a lineage is either
#### present or absent in each region. Each is named with the letters of the
#### alphabet in order, so the maximum number of areas is 26
####
#### Note 2) As the phylogeny is provided by the user, speciation and 
#### extinction times are fixed (the nodes and terminal tips of the phylogeny)
#### 
#### 
#### Arguments:
#### test.tree - the phylogeny over which the biogeographic history will be 
####             simulated
####
#### no.area - number of areas in the simulated geospace
#### 
#### max.areas - maximum number of areas in which a taxon, lineage, or node
####             may be present. Defaults to the same as the number of areas
####
#### time bins - a two column matrix giving the beginning and end of the time
####             bins in which biogeographic events will be counted
####
#### area.ajs - an optional area adjacency matrix. If left as NULL (default),
####            dispersal may happen between any region. If supplied, lineages
####            may only disperse between  areas adjacent to eachother. Should
####            be a symmetrical matrix of pairwise comparisons between areas,
####            with 1 indicating the areas are adjacent and 0 indicating they
####            are not. Rownames and column names are required, and should
####            be the names of the areas (i.e. letters of the alphabet)
####
#### disp.rate - dispersal rate (in practice the probability that a dispersal
####             event will occur at any point in time along a branch)
####
#### vic.rate - vicariance rate (in practice, the probability that a vicariance
####            event will take place at a node)
####
#### ext.rate - local extinction rate (in practice the probability that a 
####            local extinction event will occur at any point in time along 
####            a branch)
####
####
#### Output is a list of length 4
#### $node.states - a matrix giving the geographic range of each node
#### $disp.counts - the number of dispersal events in each time bin 
#### $vic.counts - the number of vicariance events in each time bin
#### $ext.counts - the number of local extinction events in each time bin
####
##############################################################################
##############################################################################





library(ape)
library(paleotree)
library(spuRs)


diva.sim<-function(test.tree,no.areas,max.areas=no.areas,time.bins,area.adj=NULL,disp.rate=0.03,vic.rate=0.03,ext.rate=0.03)
{
	######create empty arrays to store event counts
	disp.count<-array(dim=nrow(time.bins),data=0)
	vic.count<-array(dim=nrow(time.bins),data=0)
	ext.count<-array(dim=nrow(time.bins),data=0)

	#####names of areas
	areas<-LETTERS[1:no.areas]

	#####total possible combinations of areas (note, area adjacency not yet accounted for)
	states<-areas_list_to_states_list_new(areas = areas,maxareas = max.areas, include_null_range = F,split_ABC = TRUE)



	####create empty matrix to store the geographic ranges of nodes
	nodes<-test.tree$Nnode
	node.states<-matrix(nrow=nodes+length(test.tree$tip.label),ncol=max.areas)

	####create empty marux to store geographic ranges at the bottom of each branch
	bottom.of.branch<-matrix(nrow=nrow(test.tree$edge),ncol=max.areas)



	####generate Q matrix for anagenetic events
	d.e.matrix<-matrix(nrow=length(states),ncol=length(states))

	for(i in 1:nrow(d.e.matrix))
	{
		for(j in 1:ncol(d.e.matrix))
		{
			if(i!=j)
			{
				#####identify if event represents extirpation
				if(length(states[[i]])>length(states[[j]]))
				{
				#####make sure only extirpation from one event at a time
					if(length(states[[i]])-length(states[[j]])==1)
					{
						#####make sure the areas in which it survives were in the ancestral range
						if(length(intersect(states[[i]],states[[j]]))==length(states[[j]]))
						{
							d.e.matrix[i,j]<-ext.rate######ext rate not multiplied by range as q matrix specifies which area it is being removed from
						}
						else(d.e.matrix[i,j]<-0)
					}
					else(d.e.matrix[i,j]<-0)
				}
	

				#####identify if event represents dispersal 
				else if(length(states[[i]])<length(states[[j]]))
				{
				#####make sure only dispersa; from one event at a time
					if(length(states[[j]])-length(states[[i]])==1)
					{
						#####make sure the old range is in the new range
						if(length(intersect(states[[i]],states[[j]]))==length(states[[i]]))
						{
							
							if(is.null(area.adj)==F)#####check only disperse to adjacent areas
							{
								new.area<-which(states[[j]]%in%states[[i]]==F)
								if(1%in%area.adj[states[[i]],states[[j]][new.area]])
								{
									d.e.matrix[i,j]<-disp.rate*length(which(colnames(area.adj)[which(area.adj[states[[j]][new.area],]==1)]%in%states[[i]])) ######dispersal rate multiplied by number of areas adjacent to new area (as these are the areas from which it can disperse)
								}
								else(d.e.matrix[i,j]<-0)
							}
							else(d.e.matrix[i,j]<-disp.rate*length(states[[i]]))######dispersal rate is multiplied by range as it can come from any of the areas
						}
						else(d.e.matrix[i,j]<-0)					
					}
					else(d.e.matrix[i,j]<-0)
				}				
				else(d.e.matrix[i,j]<-0)
			}
		}
	}
	######fill in diagonal	
	for(i in 1:nrow(d.e.matrix))
	{
		d.e.matrix[i,which(is.na(d.e.matrix[i,])==T)]<-0-sum(d.e.matrix[i,],na.rm=T)
	}




	####choose starting range
	no.start.areas<-sample(1:max.areas,1,replace=F)
	node.states[length(test.tree$tip.label)+1,1:no.start.areas]<-sample(areas,no.start.areas,replace=F)



	########start simulation
	for(i in 1:nrow(test.tree$edge))
	{
		test.node<-test.tree$edge[i,1]
		desc.branches<-which(test.tree$edge[,1]==test.node)


		########Vicariance
		
		if(which(test.tree$edge[,1]==test.node)[1]==i)#####Make sure don't do vicariance on same node twice
		{
			if(length(which(is.na(node.states[test.node,])==F))>1)######check if vicariance is possible (i.e. node has range of more than one ares)
			{
				######test if there should be a vicariance event
				vic.test<-runif(1,0,1)	
				if(vic.test<=vic.rate)
				{
					######Do vicariance
					split<-node.states[test.node,sample(which(is.na(node.states[test.node,])==F),1)]###which area being separated
					split.1<-sample(1:2,1)
					split.2<-which(1:2%in%split.1==F)
	
					bottom.of.branch[desc.branches[split.1],1]<-split	
					bottom.of.branch[desc.branches[split.2],1:length(intersect(which(is.na(node.states[test.node,])==F), which(node.states[test.node,]!=split)))]<-node.states[test.node,which(node.states[test.node,]!=split)]
	
					vic.date<-dateNodes(test.tree)[test.node]
					bin<-intersect(which(time.bins[,1]>=vic.date),which(time.bins[,2]<=vic.date))
					vic.count[bin]<-vic.count[bin]+1
				}
				else(bottom.of.branch[desc.branches[1],]<-bottom.of.branch[desc.branches[2],]<-node.states[test.node,])
			}
			else(bottom.of.branch[desc.branches[1],]<-bottom.of.branch[desc.branches[2],]<-node.states[test.node,])
		}

		#####Anagenetic events
		prop.branch<-test.tree$edge.length[i]		
		start<-which(states %in% list(sort(bottom.of.branch[i,]))==T)	
		mcmc.res<-CMCSimulation(d.e.matrix,start-1,prop.branch)
		mcmc.res[,1]<-mcmc.res[,1]+1 #####needed as spuRs treats the first state as 0 instead of 1
		



		if(max(mcmc.res[2:(nrow(mcmc.res)-1),2])<=prop.branch)#####needed as spuRs has to end simulation on a transition, even if simulation carries on after specified time
		{	

			#####store final range at end of branch
			final.state<-states[[mcmc.res[max(which(mcmc.res[,2]<=prop.branch)),1]]]
			node.states[test.tree$edge[i,2],1:length(final.state)]<-final.state

			#####identify which events are dispersal and local extinction
			for(j in 2:max(which(mcmc.res[,2]<=prop.branch)))
			{
				initial.state<-states[[mcmc.res[j-1,1]]]
				final.state<-states[[mcmc.res[j,1]]]
				if(length(initial.state)>length(final.state))
				{
					####store extinction timings
					ext.date<-dateNodes(test.tree)[test.node]-mcmc.res[j,2]
					bin<-intersect(which(time.bins[,1]>=ext.date),which(time.bins[,2]<=ext.date))
					ext.count[bin]<-ext.count[bin]+1	
				}
				else
				{
					####store dispersal timings
					disp.date<-dateNodes(test.tree)[test.node]-mcmc.res[j,2]
					bin<-intersect(which(time.bins[,1]>=disp.date),which(time.bins[,2]<=disp.date))
					disp.count[bin]<-disp.count[bin]+1
				}
			}
		}
		else(node.states[test.tree$edge[i,2],]<-bottom.of.branch[i,])

	}

	results<-list(node.states,disp.count,vic.count,ext.count)
	names(results)<-c("node.states","disp.counts","vic.count","ext.count")
	return(results)
}