#ancestral_state reconstruction using phytools

library(ape)
library(BAMMtools)
library(coda)
library(phytools)
library(hisse)
library(diversitree)

setwd("~/Documents/Lampsiline_ddrad/Ancestral_state")

#Load Tree

musseltree <- read.tree("RAxML_bestTree.Unionids_HQ4_mltree_with_I")

#Load mussel data


musseldata <- read.table("ddrad_ancestral_state_data_waters.csv", sep=",", header=TRUE)
single.species <- read.table("ingroup.csv", sep=",", header=TRUE)

#Removes outgroups from tree
del <- 1:length(musseltree$tip.label)
ingroupnames <- single.species$Sample_name
ingroup <- match(ingroupnames, musseltree$tip.label)
ingroup <- ingroup[which(is.finite(ingroup))]
ingroup.tree=drop.tip(musseltree,del[-ingroup])


#set trait variables for reconstruction


sorted_data = single.species[match(ingroup.tree$tip.label,single.species$Sample_name),]

brood_lure <- sorted_data$Brood.lure
somatic_lure <- sorted_data$somatic.lure
host <- sorted_data$Host

#perform MCMC ancestral character reconstruction

species <- as.character(sorted_data$Species.name)
ingroup.tree$tip.label <- species
named_somatic <- setNames(somatic_lure, ingroup.tree$tip.label)
named_brood <- setNames(brood_lure, ingroup.tree$tip.label)
somaticER <- rerootingMethod(ingroup.tree, named_somatic,model="ER")
broodER <- rerootingMethod(ingroup.tree, named_brood,model="ER")
somaticSYM <- rerootingMethod(ingroup.tree, named_somatic,model="SYM")
broodSYM <- rerootingMethod(ingroup.tree, named_brood,model="SYM")
somaticARD <- rerootingMethod(ingroup.tree, named_somatic,model="ARD")
broodARD <- rerootingMethod(ingroup.tree, named_brood,model="ARD")
print(somaticER)
print(broodER)
print(somaticSYM)
print(broodSYM)


#plot trees with ancestral state estimates

plot.phylo(ingroup.tree,cex=0.5, label.offset=0.0008)
nodelabels(node = as.numeric(rownames(somaticER$marginal.anc)), pie = somaticER$marginal.anc,
    piecol = c("blue", "red"), cex = 0.3)
tiplabels(pie = to.matrix(somatic_lure, sort(unique(somatic_lure))), piecol = c("blue", "red"), 
    cex = 0.2)
    
    
plot.phylo(ingroup.tree,cex=0.5, label.offset=0.0008)
nodelabels(node = as.numeric(rownames(broodER$marginal.anc)), pie = broodER$marginal.anc,
    piecol = c("blue", "red","green","yellow"), cex = 0.3)
tiplabels(pie = to.matrix(brood_lure, sort(unique(brood_lure))), piecol = c("blue", "red","green","yellow"), 
    cex = 0.2)
    

foo<-function(tree,...){
    fsize<-36*par()$pin[2]/par()$pin[1]/Ntip(tree)
    plotTree(tree,fsize=fsize,lwd=1,offset=0.5,...)
}

#creates ultrametric tree
mycalibration <- makeChronosCalib(ingroup.tree,node="root",age.max=1)
ingroup.tree <- chronos(ingroup.tree,lambda=1,model="correlated",calibration=mycalibration,control=chronos.control())
   

#creates figure 5
layout(matrix(1:3,1,3),widths=c(0.42,0.1,0.42))
foo(ingroup.tree,ftype="off") 
nodelabels(node = as.numeric(rownames(somaticSYM$marginal.anc)), pie = somaticSYM$marginal.anc,
    piecol = c("black", "grey"), cex = 0.5)
tiplabels(pie = to.matrix(somatic_lure, sort(unique(somatic_lure))), piecol = c("black", "grey"), 
    cex = 0.3) 
plot.new()
plot.window(xlim=c(-0.1,0.1),ylim=c(1,length(ingroup.tree$tip.label)))
par(cex=1)
text(rep(0,length(ingroup.tree$tip.label)), 1:length(ingroup.tree$tip.label),ingroup.tree$tip.label,cex=0.5,font=3)
foo(ingroup.tree,ftype="off",direction="leftwards")
nodelabels(node = as.numeric(rownames(broodSYM$marginal.anc)), pie = broodSYM$marginal.anc,
    piecol = c("blue", "red","green","yellow"), cex = 0.5)
tiplabels(pie = to.matrix(brood_lure, sort(unique(brood_lure))), piecol = c("blue", "red","green","yellow"), 
    cex = 0.3)


#setting up bamm

chainSwapPercent(bammPath = "~/Documents/Lampsiline_ddrad/BAMM/Run3HQ5_85_25/bamm", controlfile = "~/Documents/Lampsiline_ddrad/BAMM/Run3HQ5_85_25/mussel_diversification.txt",
    nChains = c(2, 4, 6), deltaT = c(0.01, 0.05),
    swapPeriod = c(100, 1000), nGenerations = 20000,
    burnin = 0.2, deleteTempFiles = TRUE)


#interpreting results from BAMM

ratetree <- read.tree("~/Desktop/Bamm/85_25_long/85_25_ingroup.tre")
edata <- getEventData(ratetree, eventdata = "~/Desktop/Bamm/85_25_long/mussel_event_data.txt",burnin=0.1)
mcmcout <- read.csv("~/Desktop/Bamm/85_25_long/mcmc_out.txt",header=TRUE)   


plot(mcmcout$logLik ~ mcmcout$generation)

burnstart <- floor(0.1 * nrow(mcmcout))
postburn <- mcmcout[burnstart:nrow(mcmcout), ]


post_probs <- table(postburn$N_shifts) / nrow(postburn)
names(post_probs)

plot.bammdata(edata, lwd=2)

css <- credibleShiftSet(edata, expectedNumberOfShifts=1, threshold=5, set.limit = 0.95)
css$number.distinct
summary(css)
plot.credibleshiftset(css)

best <- getBestShiftConfiguration(edata, expectedNumberOfShifts=1)
plot.bammdata(best, lwd = 2)
addBAMMshifts(best, cex=2.5)

msc.set <- maximumShiftCredibility(edata, maximize='product')
msc.config <- subsetEventData(edata, index = msc.set$sampleindex)
plot.bammdata(msc.config, lwd=2)


##Creating color coded tips, for all combinations of host infection strategies

host.infection.strat <- rep(0,nrow(sorted_data))
for(i in 1:nrow(sorted_data)){
    if(sorted_data[i,3]=="none" & sorted_data[i,4]=="none"){
        host.infection.strat[i]="none"}else if(sorted_data[i,3]=="none" & sorted_data[i,4]=="mantle"){
        host.infection.strat[i]="mantle"}else if(sorted_data[i,3]=="superconglutinate" & sorted_data[i,4]=="none"){
        host.infection.strat[i]="tethered"}else if(sorted_data[i,3]=="superconglutinate" & sorted_data[i,4]=="mantle"){
        host.infection.strat[i]="tethered_and_mantle"}else if(sorted_data[i,3]=="fragile conglutinate"){
        host.infection.strat[i]="fragile_brood_and_mantle"}else{
        host.infection.strat[i]="complex_brood"
        }
}       

colors = c("red","blue","yellow","green","purple","orange")

#color codes tips based on host infection strategy as seen in figure 4
plot.phylo(ingroup.tree,cex=0.5, label.offset=0.008)
tiplabels(pie = to.matrix(host.infection.strat, sort(unique(host.infection.strat))), piecol = colors, 
    cex = 0.2)

###hisse package for testing state dependent diversification patterns
brood.state <- rep(0,nrow(sorted_data))
for(i in 1:nrow(sorted_data)){
    if(sorted_data[i,3]=="none"){
        brood.state[i]=0}else{
            brood.state[i]=1
        }

}

somatic.state <- rep(0,nrow(sorted_data))
for(i in 1:nrow(sorted_data)){
    if(sorted_data[i,4]=="none"){
        somatic.state[i]=0}else{
            somatic.state[i]=1
    }
}

combo.state <- rep(0,nrow(sorted_data))
for(i in 1:nrow(sorted_data)){
    if(brood.state[i]==1 && somatic.state[i]==1){
        combo.state[i]=1}else{
            combo.state[i]=0
    }
}

broadcast.state <- rep(0,nrow(sorted_data))
for(i in 1:nrow(sorted_data)){
    if(brood.state[i]==0 && somatic.state[i]==0){
        broadcast.state[i]=1}else{
            broadcast.state[i]=0
        
    }
}

named_brood <- setNames(brood.state, ingroup.tree$tip.label)
named_broadcast <- setNames(broadcast.state, ingroup.tree$tip.label)
brood.state = data.frame(ingroup.tree$tip.label,brood.state)
somatic.state = data.frame(ingroup.tree$tip.label,somatic.state)
combo.state = data.frame(ingroup.tree$tip.label,combo.state)
broadcast.state = data.frame(ingroup.tree$tip.label,broadcast.state)

#set sampling frequencies for each group

brood.f = c(0.36,0.51)
somatic.f = c(0.3,0.48)
combo.f = c(0.33,0.62)
broadcast.f = c(0.44,0.31)

#function for comparing 2 and 4 state character dependent and character independent models and graphically displaying best model

hisse.comparison <- function(tree,dat, f=c(1,1)){
    hisse.turnover = c(1,2,3,4)
    hisse.eps = c(1,2,3,4)
    bisse.turnover = c(1,2,0,0)
    bisse.eps = c(1,2,0,0)
    bisse.null.turnover = c(1,1,2,2)
    bisse.null.eps = c(1,1,2,2)

    trans.rates = TransMatMaker(hidden.states=TRUE) 

    trans.rates.nodual = ParDrop(trans.rates, c(3,5,8,10))

    trans.rates.nodual.allequal = ParEqual(trans.rates.nodual, c(1,2,1,3,1,4,1,5,1,6,1,7,1,8))

    trans.rates.bisse = TransMatMaker(hidden.states=FALSE)

    trans.rates.bisse.equal = ParEqual(trans.rates.bisse, c(1,2))

    model.hisse <- hisse(tree, dat, f, hidden.states=TRUE, turnover.anc=hisse.turnover, eps.anc=hisse.eps, trans.rate=trans.rates.nodual.allequal)
    model.bisse <- hisse(tree, dat, f, hidden.states=FALSE, turnover.anc=bisse.turnover, eps.anc=bisse.eps, trans.rate=trans.rates.bisse.equal)
    bisse.null <- hisse(tree, dat, f, hidden.states=TRUE, turnover.anc=bisse.null.turnover, eps.anc=bisse.null.eps, trans.rate=trans.rates.nodual.allequal)
    hisse.null <- hisse.null4(tree, dat, f, trans.type="equal")

    models <- list(model.hisse,model.bisse,bisse.null,hisse.null)
    model.names <- list("hisse model","bisse model","bisse null","hisse null")

    table = data.frame(model.name = rep(NA,4), AIC = rep(NA,4),AICc = rep(NA,4),log.likelihood = rep(NA,4))
    for (i in 1:4){
        table[i,1]=model.names[[i]]
        table[i,2]=models[[i]]$AIC
        table[i,3]=models[[i]]$AICc
        table[i,4]=models[[i]]$loglik
    }
    table = table[order(table$AICc),]
    return(table)

   
}

#perform SSE analyses

brood.models <- hisse.comparison(ingroup.tree,brood.state,f=brood.f)
brood.models

somatic.models <- hisse.comparison(ingroup.tree,somatic.state,f=somatic.f)
somatic.models

combo.models <- hisse.comparison(ingroup.tree,combo.state,f=combo.f)
combo.models

broadcast.models <- hisse.comparison(ingroup.tree,broadcast.state,f=broadcast.f)
broadcast.models

#MCMC BISSE MODEL WITH DIVERSITREE
#brood lure

nbModel <- make.bisse(ingroup.tree,named_brood,sampling.f=brood.f)
p <- starting.point.bisse(ingroup.tree)
nbModelMLFit <- find.mle(nbModel, p)

prior <- make.prior.exponential(1/(2 * 0.4))
mcmcRun <- mcmc(nbModel, nbModelMLFit$par, nsteps = 10000, prior = prior, w = 0.1, print.every = 100)

col <- c("blue", "red")
profiles.plot(mcmcRun[, c("lambda0", "lambda1")], col.line = col, las = 1, legend = "topright")

#broadcast
nbModel <- make.bisse(ingroup.tree,named_broadcast,sampling.f=broadcast.f)
p <- starting.point.bisse(ingroup.tree)
nbModelMLFit <- find.mle(nbModel, p)

prior <- make.prior.exponential(1/(2 * 0.4))
mcmcRun <- mcmc(nbModel, nbModelMLFit$par, nsteps = 10000, prior = prior, w = 0.1, print.every = 100)

col <- c("blue", "red")
profiles.plot(mcmcRun[, c("lambda0", "lambda1")], col.line = col, las = 1, legend = "topright")