library(ape)
library(sde)
## Loading required package: MASS
## Loading required package: stats4
## Loading required package: fda
## Loading required package: splines
## Loading required package: Matrix
## 
## Attaching package: 'fda'
## The following object is masked from 'package:graphics':
## 
##     matplot
## Loading required package: zoo
## 
## Attaching package: 'zoo'
## The following objects are masked from 'package:base':
## 
##     as.Date, as.Date.numeric
## sde 2.0.15
## Companion package to the book
## 'Simulation and Inference for Stochastic Differential Equations With R Examples'
## Iacus, Springer NY, (2008)
## To check the errata corrige of the book, type vignette("sde.errata")
library(phytools)
## Loading required package: maps
library(MASS)
library(stats4)
library(fda)
library(splines)
library(Matrix)
library(fds)
## Loading required package: rainbow
## Loading required package: pcaPP
## Loading required package: RCurl
## Loading required package: bitops
library(rainbow)
library(pcaPP)
library(RCurl)
options(warn = -1)

#plot tree using path space
PlotTreeHistory<-function(phy=phy,pathdata=pathdata,main=main,colors=colors){
  #pathdata=oubpathdata
  ntips<-length(phy$tip.label)
  edge.number<-dim(phy$edge)[1]
  start<-rep(0,length(edge.number))
  end<-rep(0,length(edge.number))
  
  mtx<-(cbind(phy$edge,ceiling(phy$edge.length)+1,start,end))
  mtx<-mtx[dim(mtx)[1]:1,]
  colnames(mtx)<-c("anc","des","step","start","end")
  mtx<-as.data.frame(mtx)
  mtx$start[1:2]<-1
  mtx$end[1:2]<-mtx$step[1:2]
  for(edgeIndex in 3:dim(mtx)[1]){
    mtx$start[edgeIndex] <-  mtx$end[mtx$anc[edgeIndex]==mtx$des]
    mtx$end[edgeIndex] <-    mtx$end[mtx$anc[edgeIndex]==mtx$des]+mtx$step[edgeIndex]-1
  }
  ylim<-c(1, max(mtx$end) )
  unlistpath<-unlist(pathdata)
  xlim<-c(min(unlistpath),1.2*max(unlistpath))
  plot(NA,xlim=xlim,ylim=ylim,xlab="Trait value", ylab="Times Steps",main=main,frame.plot=FALSE,yaxt='n',xaxt='n',ann=FALSE)
  points(pathdata[[1]][1] ,1, lwd=2,pch=5,col= 1,cex=1.5)
  for(edgeIndex in 1:edge.number){
    path<-unlist(pathdata[edgeIndex]) 
    points(path,mtx$start[edgeIndex]:mtx$end[edgeIndex], type="l",lwd=2,col=colors[edgeIndex])
    points(path[length(path)],mtx$end[edgeIndex]:mtx$end[edgeIndex], cex=1.5,lwd=2,
           pch=c(16,16,16,16,16)[edgeIndex],
           col=colors[edgeIndex] #c("black","black","blue", "orange","green","black", "black","brown")[edgeIndex]
           )#,col=colors[edgeIndex])
  }    
}#end of plot history


#Ornstein Uhlenbeck bridge
oub<-function(Theta, a=0,b=0,t0=0,T=1,N=100){
  alpha<-Theta[1]
  theta<-Theta[2]
  sigma<-Theta[3]
  if(T<=t0){stop("wrong times")}
  dt<-(T-t0)/N
  t<-seq(t0,T,length=N+1)
  d<-expression(-1*x)#alp
  s<-expression(2)#sigma=2
  X<-sde.sim(X0=a,drift=d,sigma=s,N=N)
  Lambda<-exp(-alpha*(T-t))-exp(-alpha*(T+t))
  Lambda<-Lambda/(1-exp(-2*alpha*T))
  oubpath<-X-Lambda*(X[N+1]-b)
  oubpath<-as.numeric(oubpath)
  return(oubpath)
}

condOU<-function(Theta,phy=phy,tips=tips){
  alpha<-Theta[1]
  theta<-Theta[2]
  sigma<-Theta[3]
  n<-length(phy$tip.label)
  m<-dim(phy$edge)[1]+1
  covmtx<-sigma^2/2/alpha*exp(-alpha*dist.nodes(phy))
  
  cov22 <- covmtx[1:n,1:n]
  cov11<-covmtx[(n+1):m,(n+1):m]
  cov12<-covmtx[1:n,(n+1):m]
  
  inv_cov22 <-solve(cov22)
  
  condmean <- theta + t(cov12)%*%inv_cov22%*%(tips-rep(theta,n))
  condvar <- cov11 - t(cov12)%*%inv_cov22%*%cov12
  return(list(condmean=condmean,condvar=condvar))
}



oubtree<-function(Theta,phy=phy,tips=tips){
  #anclist<-condOU(Theta=Theta,phy=phy,tips=tips)
  #fullnodedata<-c(tips,anclist$condmean)
  anclist<-ace(x=tips,phy=phy)
  fullnodedata<-c(tips,anclist$ace)
  anc<-phy$edge[,1]
  des<-phy$edge[,2]
  ntips<-length(phy$tip.label)
  edge.number<-dim(phy$edge)[1]
  edge.length<-phy$edge.length
  pathdata<-list()
  for(edgeIndex in edge.number:1){
    brlen<-edge.length[edgeIndex]
    N<-ceiling(brlen)
    b<-fullnodedata[des[edgeIndex]]
    a<-fullnodedata[anc[edgeIndex]]
    assign(paste("path",edgeIndex,sep=""),oub(Theta,a=a,b=b,t0=0,T=brlen,N=N))
    pathdata<-c(pathdata, list(get(paste("path",edgeIndex,sep=""))) )
  }
  return(pathdata)
}
######################################
################ Main ################
######################################

#phy<-read.newick(text="(((D:0.3,E:0.3):0.2,C:0.5));")
phy<-read.newick(text="((D:0.3,E:0.3):0.2,C:0.5);")
plot(phy)

par(mfrow=c(1,3))
par(mar=c(2,2,1,1))
colors<-c("white","red","blue","orange")
min.length<-100
while(min(phy$edge.length)<min.length){
  phy$edge.length<-1.005*phy$edge.length
}
phy<-reorder(phy,"postorder")
set.seed(20210906)
## OUB tree
tips<-c(9.8,8.2,-1.1)
anclist<-ace(x=tips,phy=phy)
anclist$ace
##        4        5 
## 4.841171 7.217642
fullnodedata<-c(tips,anclist$ace)
round(fullnodedata,2)
##                       4     5 
##  9.80  8.20 -1.10  4.84  7.22
Theta=c(0.01,mean(tips),1) #(alp,theta,sigma)
oubpathdata<-oubtree(Theta,phy=phy,tips=tips)
## sigma.x not provided, attempting symbolic derivation.
## sigma.x not provided, attempting symbolic derivation.
## 
## sigma.x not provided, attempting symbolic derivation.
## 
## sigma.x not provided, attempting symbolic derivation.
PlotTreeHistory(phy=phy,pathdata=oubpathdata,main="Trait Change",colors=colors)

phy$tip.label<-paste(" Spe.",c(1:3),sep="")
phy$edge.length
## [1] 150.7302 150.7302 100.4868 251.2170
plot(phy,edge.width=5,cex=1.5,edge.color=rev(colors),direction="upward")



tips<-c(2.8,7.2,2.1)
anclist<-ace(x=tips,phy=phy)
anclist$ace
##        4        5 
## 3.805878 4.488231
fullnodedata<-c(tips,anclist$ace)
round(fullnodedata,2)
##                   4    5 
## 2.80 7.20 2.10 3.81 4.49
Theta=c(0.01,mean(tips),1) #(alp,theta,sigma)
oubpathdata<-oubtree(Theta,phy=phy,tips=tips)
## sigma.x not provided, attempting symbolic derivation.
## 
## sigma.x not provided, attempting symbolic derivation.
## 
## sigma.x not provided, attempting symbolic derivation.
## 
## sigma.x not provided, attempting symbolic derivation.
PlotTreeHistory(phy=phy,pathdata=oubpathdata,main="Trait Change",colors=colors)