### MCMC

rm(list = ls()) # clear environment 
Sys.setenv(LANG = "en") # set language 

# load packages
library(phytools)
library(geiger) 
library(diversitree)
setwd("~/Documents/Uni/PHD/Statistics/MCMC")


#--------------------------------------------------
# Read in data
#--------------------------------------------------

# load tree
originaltrees1 <- read.nexus("Reginato_9-gene_Dated_fixed_topology.tre")

taxa1 <- originaltrees1$ALL_FOSSILS # select tree 

# load traits
temp1 <- read.csv("MelasTraitsMean.csv", header=T, sep=";", row.names = 1)

# remove species with NAs
temp1=temp1[c(1:348,350:400,402:413),]

# get syndromes and tip names
states <- temp1$Pred.Syn.Disp
names(states)<-temp1$Phylo.Name

## remove species from tree and data not included in tree or data
to.drop <- taxa1$tip.label[!(taxa1$tip.label %in% names(states))]
trees1 <- lapply(originaltrees1, drop.tip, to.drop)
tree1 <- trees1$ALL_FOSSILS

### rename syndromes to 1 - 4!
states2<-treedata(taxa1,states,sort=T,warnings=T)$data

states1<-states2
states1[which(states1=="Buzz")]<-1
states1[which(states1=="Passerine")]<-2
states1[which(states1=="Nectar foraging vertebrates")]<-3
states1[which(states1=="Generalist")]<-4

states1 <- as.vector(states1)
states1 <- as.integer(states1)
names(states1)<-rownames(states2)

## change sampling frection! sampling fraction - probability of an extent specie to be included - the more species of a syndroe sampled the higher the sampling franction - dont know how many buzz species included in my set in % of total buzz species
sampling <- table(states1) # 285 Buzzc(1), 15 Gen (4), 30 Nec (3), 6 Pass (2): Of all melas 95.5% buzz, 2.5% Nectar, 0.9% Pass, 1.1% Gen
sampling[1]<-285/5594 # 95,5%   total 5594 buzz  / 285 in data  
sampling[2]<- 6/53 # 0.9%  total 53  passerine / 6 in data
sampling[3]<- 28/146 # 2.5%  total 146  Nectar / 28 in data
sampling[4]<- 15/64 # 1.1% total 64 generalist / 15 in data

# make model
function(treenum)
  tree1 <- trees1[[treenum]]
age.orig1 <- max(branching.times(tree1))
tree1$edge.length <- tree1$edge.length / age.orig1 
age <- 1

lik1 <- make.musse(tree1, states1, 4, sampling.f=sampling) # base model
lik.sym1 <- constrain(lik1, lambda1 ~ lambda2, lambda1 ~ lambda3, lambda1 ~ lambda4, lambda2 ~ lambda3, lambda2 ~ lambda4, lambda3 ~ lambda4,mu1 ~ mu2, mu1 ~ mu3, mu1 ~ mu4, mu2 ~ mu3, mu2 ~ mu4, mu3 ~ mu4, q12 ~ q21, q13 ~ q31, q14 ~ q41, q23 ~ q32, q24 ~ q42, q34 ~ q43) # SYM
likfreetrans <- constrain(lik1,lambda1 ~ lambda2, lambda1 ~ lambda3, lambda1 ~ lambda4, lambda2 ~ lambda3, lambda2 ~ lambda4, lambda3 ~ lambda4,mu1 ~ mu2, mu1 ~ mu3, mu1 ~ mu4, mu2 ~ mu3, mu2 ~ mu4, mu3 ~ mu4) # ARD  
constraint <- constrain(lik1, lambda1 ~ lambda2, lambda1 ~ lambda3, lambda1 ~ lambda4, lambda2 ~ lambda3, lambda2 ~ lambda4, lambda3 ~ lambda4,mu1 ~ mu2, mu1 ~ mu3, mu1 ~ mu4, mu2 ~ mu3, mu2 ~ mu4, mu3 ~ mu4, q13 ~ 1, q14 ~ 1, q12 ~ 0.75, q21 ~ 0.25, q31 ~ 0.25, q41 ~ 0.25, q34 ~ 0.5, q43 ~ 0.5, q23 ~ 0.01, q32  ~ 0.01, q42 ~ 0.01, q24 ~ 0.01) # 

# fit models
p1 <- starting.point.musse(tree1, 4)
fit.likfreetrans <- find.mle(likfreetrans, p1[argnames(likfreetrans)]) 
fit.lik.sym1 <- find.mle(lik.sym1, p1[argnames(lik.sym1)])    
fit.lik.const <- find.mle(constraint, p1[argnames(constraint)]) 

# compare AIC of models

aciFreetrans <- AIC(fit.likfreetrans) # ARD -171.6969
acilik.sym1<-AIC(fit.lik.sym1) # SYM -201.2776
aciConst <- AIC(fit.lik.const) # Const. 113.3236

aicc<-setNames(c(aciFreetrans,acilik.sym1,aciConst),c("ARD","SYM","Constraint"))
aic.w(aicc) # SYM 

anova(fit.likfreetrans,fit.lik.sym1,fit.lik.const) # 

#MCMC with consensus

argnames(lik.sym1) # parameters of model
# use best model to get prior rates
# make mcmc to get priors and use SYM - w = how many function evaluations are required between sample updates (w = 5 -> 5 function evaluations between sample updates) 8 because of 8 parameters
temp <- mcmc(lik.sym1, rep(1, 8), nsteps=250, w=rep(5, 8), prior=make.prior.exponential(age), upper=rep(age*100, 8), print.every=0)
temp<- temp[-seq(50),] # burn the first 50 objects

# Parameters for priors 
prior.rate <- sapply(temp[, argnames(lik.sym1)], median)
prior.new <- make.prior.exponential(1/prior.rate)


lik.names <- argnames(lik.sym1)
# use best model to get mcmc controls 
# prelim chain to determine mcmc controls and remove burnin - get values (x.init = par.start, w, upper) for real chain
temp <- mcmc(lik.sym1 , rep(age/2, length(lik.names)), nsteps=250, w=age, prior=prior.new, upper=age*100, print.every=0)
temp <- temp[-seq(50),] # remove burning
w.new <- round(apply(temp[, lik.names], 2, max) - apply(temp[, lik.names], 2, min), 2) # number of function evaluations  required between sample updates
upper.new <- round(apply(temp[, lik.names], 2, max) * 5, 2) # Upper bounds on parameter space 
par.start <- as.numeric(temp[which.max(temp$p), lik.names])

# run mcmc on consensus tree
mcmc.lik.SYM <- mcmc(lik.sym1 , par.start, nsteps=5000, w=w.new, upper=upper.new, prior=prior.new, print.every=100)

write.csv(mcmc.lik.SYM, "~/Documents/Uni/PHD/Statistics/MCMC/MCMC5000SYM.csv")

##### Analyse MCMC output ####
### profil plot, plot parameters

colors<-setNames(c("blue","yellow","green","red"),1)
# lambda
pdf("~/Documents/Uni/PHD/Statistics/MCMC/Profilplot_lambda_SYM_FinFin.pdf")
profiles.plot(mcmc.lik.SYM[2], col.line=colors, las=1, legend.pos="topright", main = "Profilplot Lambda Consenus SYM")
dev.size(units = c("cm"))
dev.off()

# mu
pdf("~/Documents/Uni/PHD/Statistics/MCMC/Profilplot_mu_SYM_FinFin.pdf")
profiles.plot(mcmc.lik.SYM[3], col.line=colors, las=1, legend.pos="topright", main = "Profilplot mu Consenus SYM")
dev.size(units = c("cm"))
dev.off()

#q
pdf("~/Documents/Uni/PHD/Statistics/MCMC/Profilplot_q_SYM_FinFin.pdf")
colorsq <-setNames(c("yellow","green","blue","red","gray0","violet","deeppink","orange","cadetblue1","chocolate4","aquamarine4","azure4"),1)
profiles.plot(mcmc.lik.SYM[4:9], col.line=colorsq,las=1, legend.pos="topright",main = "Profilplot q Consenus SYM")
dev.size(units = c("cm"))
dev.off()

#p
pdf("~/Documents/Uni/PHD/Statistics/MCMC/MCMC_Parameters_p_SYM_FinFin.pdf")
plot(mcmc.lik.SYM$i, mcmc.lik.SYM$p, type="l", xlab="generation", ylab="log(L)", main = "MCMC_Parameters_p SYM")
dev.size(units = c("cm"))
dev.off()

#################################################

### run MCMC over 100 randomly sampled trees

# the mcmc output files will be placed here
setwd()

outdir <- "mcmc_results_final/"
if (!file.exists(outdir))
  dir.create(outdir)

library(parallel)  # for mclapply() below

# Fill with leading zeros to get an n-digit character string
# http://tolstoy.newcastle.edu.au/R/help/00a/1076.html
zfill <- function(x, n)
{
  nc <- nchar(x)
  zeros <- paste(rep(0, n), collapse = "")
  paste(substring(zeros, nchar(x) + 1, n), substring(x, 1, nchar(x)), sep = "")
}
originaltrees <- read.nexus("all_fossils_Reginato_dated_sameTopology.trees")

originaltrees<-sample(originaltrees,size=100) # sample 100 random trees

taxa <- originaltrees[[1]] # select a single tree for tip lables

temp <- read.csv("MelasTraitsMean.csv", header=T, sep=";", row.names = 1)

# remove species with NAs
temp=temp[c(1:348,350:400,402:413),]

# get syndromes and tip names
states <- temp$Pred.Syn.Disp
names(states)<-temp$Phylo.Name

## remove species from tree and data not included in tree or data
to.drop <- taxa$tip.label[!(taxa$tip.label %in% names(states))]
trees <- lapply(originaltrees, drop.tip, to.drop)

### rename syndromes to 1 - 4!
states2<-treedata(taxa,states,sort=T,warnings=T)$data

states1<-states2
states1[which(states1=="Buzz")]<-1
states1[which(states1=="Passerine")]<-2
states1[which(states1=="Nectar foraging vertebrates")]<-3
states1[which(states1=="Generalist")]<-4

states1 <- as.vector(states1)
states1 <- as.integer(states1)
names(states1)<-rownames(states2)

## change sampling frection! sampling fraction - probability of an extent specie to be included - the more species of a syndroe sampled the higher the sampling franction - dont know how many buzz species included in my set in % of total buzz species
sampling <- table(states1) # 285 Buzzc(1), 15 Gen (4), 30 Nec (3), 6 Pass (2): Of all melas 95% buzz, 2.5% Nectar, 0.9% Pass, 1.1% Gen
sampling[1]<-285/5594 # 95,5%   total 5594 buzz species / 285 in data   # Nr. Species Melas Ulloa-Ulloa 2022 5858 sp.
sampling[2]<- 6/53 # 0.9%  total  species passerine / 6 in data
sampling[3]<- 28/146 # 2.5%  total 137 species Nectar / 28 in data
sampling[4]<- 15/64 # 1.1% total generalist / 15 in data

# get priors
tree <- trees[[10]] # select tree Nr 10
tree$edge.length <- tree$edge.length / max(branching.times(tree))
age <- 1

lik <- make.musse(tree, states1, 4, sampling.f=sampling) 
lik.sym <- constrain(lik, lambda1 ~ lambda2, lambda1 ~ lambda3, lambda1 ~ lambda4, lambda2 ~ lambda3, lambda2 ~ lambda4, lambda3 ~ lambda4,mu1 ~ mu2, mu1 ~ mu3, mu1 ~ mu4, mu2 ~ mu3, mu2 ~ mu4, mu3 ~ mu4, q12 ~ q21, q13 ~ q31, q14 ~ q41, q23 ~ q32, q24 ~ q42, q34 ~ q43) # SYM

# make mcmc to get priors using best model
temp <- mcmc(lik.sym, rep(1, 8), nsteps=250, w=rep(5, 8), prior=make.prior.exponential(age), upper=rep(age*100, 8), print.every=0)
temp<- temp[-seq(50),] # burn the first 50 objects

# Parameters for SYM Model : "lambda4""mu4""q21" "q31" "q32" "q41" "q42" "q43"  
# Use same values for clado and ana change; different for lambda and mu
prior.rate <- sapply(temp[, argnames(lik.sym)], median)
prior.new <- make.prior.exponential(1/prior.rate)

#
# The function to run for each tree
#--------------------------------------------------

run.mcmc <- function(treenum)
{
  tree <- trees[[treenum]]
  age.orig <- max(branching.times(tree))
  tree$edge.length <- tree$edge.length / age.orig 
  age <- 1
  
  lik.classefull <- make.musse(tree, states1, 4, sampling.f=sampling) 
  
  # tinclude best model (SYM)
  lik <- constrain(lik.classefull, lambda1 ~ lambda2, lambda1 ~ lambda3, lambda1 ~ lambda4, lambda2 ~ lambda3, lambda2 ~ lambda4, lambda3 ~ lambda4,mu1 ~ mu2, mu1 ~ mu3, mu1 ~ mu4, mu2 ~ mu3, mu2 ~ mu4, mu3 ~ mu4, q12 ~ q21, q13 ~ q31, q14 ~ q41, q23 ~ q32, q24 ~ q42, q34 ~ q43)
  lik.names <- argnames(lik)
  
  # prelim chain to determine mcmc controls and remove burnin - get values (x.init = par.start, w, upper) for real chain
  temp <- mcmc(lik, rep(age/2, length(lik.names)), nsteps=250, w=age, prior=prior.new, upper=age*100, print.every=0)
  temp <- temp[-seq(50),] # remove burning
  w.new <- round(apply(temp[, lik.names], 2, max) - apply(temp[, lik.names], 2, min), 2) # number of function evaluations  required between sample updates
  upper.new <- round(apply(temp[, lik.names], 2, max) * 5, 2) # Upper bounds on parameter space 
  par.start <- as.numeric(temp[which.max(temp$p), lik.names]) # start state?
  
  message(paste("Done with prelim for tree", treenum))
  
  # real chain 
  ans <- mcmc(lik, par.start, nsteps=5000, w=w.new, upper=upper.new, prior=prior.new, print.every=0)
  # rescale by tree age
  ans[, lik.names] <- ans[, lik.names] / age.orig
  
  # write the results for this tree
  ans <- data.frame(treeid=treenum, ans)
  outfile <- paste(outdir, "MUSSE_tree", zfill(treenum, 3), ".csv", sep="")
  write.csv(ans, file=outfile, row.names=F)
  
  message(paste("Done fitting tree", treenum))
}

#--------------------------------------------------
# Actually running the analysis
#--------------------------------------------------

treesubset <- 1:100 # in trees 1026 trees - now run with the first 100?

# easy parallelization: give each core a tree to run
junk <- mclapply(treesubset, run.mcmc)

############################################
##### analyse output
setwd("~/Documents/Uni/PHD/Statistics/MCMC")
#### calculate mean of parameters

data_files <- list.files("~/Documents/Uni/PHD/Statistics/MCMC/mcmc_results_finalFinal")  # Identify file names
data_files   

test <- read.csv("~/Documents/Uni/PHD/Statistics/MCMC/mcmc_results_final/MUSSE_tree092.csv")
colMeans(test)[3:10]

# df for output
paraMean <- data.frame(matrix(nrow=100,ncol=9));colnames(paraMean)<- c("lambda4","mu4", "q21","q31","q32","q41","q42","q43","p")

#loop
for(i in 1:length(data_files)) {                             
  ParaData <- assign(paste0("data", i),                                   
                     read.csv(paste0("~/Documents/Uni/PHD/Statistics/MCMC/mcmc_results_final/",
                                     data_files[i])))
  paraMean[i,] <- colMeans(ParaData)[3:11]
  
}

# wirte
write.csv(paraMean, "~/Documents/Uni/PHD/Statistics/MCMC/MeanParametersOver100Trees_with5000Steps_Fin.csv") #save out the results


# profiles.plot
### profil plot, plot parameters
colors<-setNames(c("blue","yellow","green","red"),1)
# lambda
pdf("~/Documents/Uni/PHD/Statistics/MCMC/Profilplot_lambda_SYM_Fin100Trees5000.pdf")
profiles.plot(paraMean[1], col.line=colors, las=1, legend.pos="topright", main = "Profilplot Lambda Mean 100 Trees 5000 Steps SYM")
dev.size(units = c("cm"))
dev.off()

# mu
pdf("~/Documents/Uni/PHD/Statistics/MCMC/Profilplot_mu_SYM_Fin100Trees5000.pdf")
profiles.plot(paraMean[2], col.line=colors, las=1, legend.pos="topright", main = "Profilplot mu Mean 100 Trees 5000 Steps SYM")
dev.size(units = c("cm"))
dev.off()

#q
pdf("~/Documents/Uni/PHD/Statistics/MCMC/Profilplot_q_SYM_Fin100Trees5000.pdf")
colorsq <-setNames(c("yellow","green","blue","red","gray0","violet","deeppink","orange","cadetblue1","chocolate4","aquamarine4","azure4"),1)
profiles.plot(paraMean[3:8], col.line=colorsq,las=1, legend.pos="topright",main = "Profilplot q Mean 100 Trees 5000 Steps SYM")
dev.size(units = c("cm"))
dev.off()

#p
pdf("~/Documents/Uni/PHD/Statistics/MCMC/MCMC_Parameters_p_SYM_Fin100Trees5000.pdf")
plot(row.names(paraMean), paraMean$p, type="l", xlab="generation", ylab="log(L)", main = "MCMC_Parameters_p Mean 100 Trees 5000 Steps SYM")
dev.size(units = c("cm"))
dev.off()


### Plot ACE using best MuSSE model and consensus tree

# Export the marginal ancestral reconstruction at the nodes of the tree
st <- t(asr.marginal(lik.sym1,coef(fit.lik.sym1)))

# color for tips
synd <- temp1$Pred.Syn.Disp 
names(synd) <- temp1$Phylo.Name 

datcol1<-treedata(tree1,synd,sort=T,warnings=T)$data
datcol <- as.data.frame(datcol1)

# Color Tips
synd1 <- as.data.frame(datcol$V1)
tip_col <- synd1$`datcol$V1`

tip_col[which(tip_col=="Buzz")]<-"cornflowerblue" # cornflowerblue
tip_col[which(tip_col=="Passerine")]<-"darkgoldenrod" # darkgoldenrod
tip_col[which(tip_col=="Nectar foraging vertebrates")]<-"deeppink" # deeppink
tip_col[which(tip_col=="Generalist")]<-"chartreuse" # chartreuse
        
# Color Nodeds
piecol=c("cornflowerblue","darkgoldenrod","deeppink","chartreuse")
      
quartz()
pdf("~/Documents/Uni/PHD/Statistics/MCMC/MuSSE_CharacterMapping_SYM_FinFinTEST_FIN.pdf")
plot(tree1, type="phylogram", cex = 0.3, edge.width=0.5, show.tip.label = T,use.edge.length = T,node.depth=2, label.offset = 3)
nodelabels(pie=st,piecol=piecol, cex=0.3)
tiplabels(pch = 22, bg = tip_col, cex = 0.7, adj = 1, offset = 0.5)
axisPhylo()
axisPhylo(side=1,root.time=NULL,backward=TRUE,las=1)
#legend("topleft", c("Buzz", "Passerine", "Nectar foraging vertebrates", "Generalist"),pch=16, cex = 0.5, col=piecol, bty="n")
      
dev.size(units = c("cm"))
dev.off()
      

