# set this to wherever your files are saved!
setwd("C:/Users/mrnvd/Desktop/R/Mariana")

# source useful R packages 
library(paleotree)
library(strap)


#source additional R scripts for functions not in any R package (make sure it's in your working directory!)
source("functions.R")
 

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

# Plot results from analyses

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

# Select a taxon to analyze (choose one at a time)

genus <- paste("Anazygidae")

# get min taxon age for selected genus

#anchor <- getMinTaxonAge(genus)
anchor = 443.07

#paste0(genus,".t")
setwd("C:/Users/mrnvd/Desktop/R/Mariana")
# read in t and p files grab a sample of time-scaled trees
t.file <- read.nexus("anazygids_ne.nexus", tree.names = TRUE)
t.file
p.file <- read.table("anazygidae_pstat.txt", header = TRUE)
p.file
trees <- t.file[5000:20001]
for (i in 1:length(trees)){
  oldRootAge <- max(dateNodes(trees[[i]]))
  #print(oldRootAge)
  trees[[i]]$root.time <- oldRootAge + anchor #
  print(trees[[i]]$root.time)
}
# adjust branches from units of (# of changes * time) to units of (time) by dividing 
# each tree by the mean clock rate
#clock.rate = p.file[5,2]
#clock.rate
#t.file[["edge.length"]] = t.file[["edge.length"]] / clock.rate
#for (i in 1:length(t.file)){
#  clock.rate <- p.file$ORCucldMean[i]
#  t.file[[i]]$edge.length <- t.file[[i]]$edge.length / clock.rate
#}

# grab post-burnin trees
#trees <- t.file[5000:20001]

# rescale the root age so node ages are in 'deep time'

#for (i in 1:length(trees)){
#  oldRootAge <- max(dateNodes(trees[[i]]))
#  print(oldRootAge)
  #trees[[i]]$root.time <- oldRootAge + anchor #
# }


#oldRootage = max(dateNodes(t.file))
#t.file$root.time = oldRootage + anchor

#
# Read in geologic time scale and store info for help with plotting later
#
geol_time<-read.table("Geologic_Timescale.txt",header=TRUE)
print(geol_time)
Stages<-cbind(geol_time[,3],geol_time[,4])
print(Stages)
row.names(Stages)<-geol_time[,2]
print(Stages)
midpoints<-(Stages[1:(length(Stages)-1)]+Stages[2:length(Stages)])/2
print(midpoints)
midpoints <- midpoints[4:7]# Sandbian, Katian, and Hirnatian
int.start <- c(467.30, 458.18,452.75,445.21, 443.07, 433.40)

####

# Diversity curves

####

# get diversity curves for multiple trees and plot the resulting median diversity with 95% quantiles

k <- multiDiv(trees, int.times = Stages[5:7,], drop.ZLB = TRUE, plot = TRUE) # Since it's only one tree

# store diversity data for each tree to plot later
div <- k$median.curve[,1]
lower.conf <- k$median.curve[,2]
upper.conf <- k$median.curve[,3]

# begin the plot
plot(midpoints, div, xlim = rev(c(440, 460)),ylim = c(-0.05, max(upper.conf)), type = "n", axes = F, xlab = "", ylab = "")
# make boxes for each stage
rect(440, 0, 460, 20, col = rep(c("grey95", "grey97")), border = NA)

# make polygons of upper and lower 95% quantiles
polygon(c(midpoints,rev(midpoints)),c(div,rev(upper.conf)), col = "seashell2", lty = 2)
polygon(c(midpoints,rev(midpoints)),c(div,rev(lower.conf)), col = "seashell2", lty = 2)


# plot lines & points
lines(midpoints, div, xlim = c(440,460), lwd = 2, col = "black")
points(midpoints, div, xlim = c(440,460), pch = 22, cex = 1.5, col = "black", bg = "black")

# plot axes
axis(1, col = "black", line = 0, at=c(0, 440, 450,460), lwd=2, lwd.ticks = 2, cex.axis=1.5, col.ticks = "black")
axis(2, col = "black", line = 0.5, at = seq(0, max(upper.conf), round(1.7)), lwd=2, lwd.ticks = 2, cex.axis=1.5, col.ticks = "black")

mtext("Age (Ma)", side = 1, line = 2.5, cex=1.5)
mtext("Diversity (# of lineages)", side = 2, line = 2.8, cex=1.5)


pdf(file = "LineageDiversityGSAlog.pdf", height =5.5, width = 6, useDingbats = FALSE)
dev.off()



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

# Net diversification and sampling plots

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


# read in the p.stats file
print(p.file)

# sort 'param' into quantities to plote
netDiv <- p.file[1:4,2]
upper.div <- p.file[1:4,5]
lower.div <- p.file[1:4,4]
extinction <- p.file[5:8,2]
upper.e <- p.file[5:8,5]
lower.e <- p.file[5:8,4]
relSamp <- p.file[9:12,2]
upper.s <- p.file[9:12,5]
lower.s <- p.file[9:12,4]

# convert FBD parameters to speciation, extinction, and sampling rates
meanMu <- calcMu(netDiv, extinction)
meanLambda <- calcLambda(netDiv, meanMu)
meanPsi <- calcPsi(meanMu, relSamp)

# if you want medians instead
#meanMu <- calcMu(netDiv[,6], turnover[,6])
#meanLambda <- calcLambda(netDiv[,6], meanMu)
#meanPsi <- calcPsi(meanMu, relSamp[,6])

# set up plotting
midpoints = c(451,450,449.5,443.07)
# Net diversification
# plot
plot(midpoints, netDiv, xlim = c(460,440),ylim = c(-0.05, max(upper.div)), type = "n", axes = F, xlab = "", ylab = "")

rect(int.start[-1], rep(0, 9), int.start[-10], rep(100, 9), col = rep(c("grey95", "grey97")), border = NA)

segments(midpoints,netDiv,midpoints, lower.div, col = "darkblue", lty = 2, lwd = 2)
segments(midpoints,netDiv,midpoints, upper.div, col = "darkblue", lty = 2, lwd = 2)

lines(midpoints, netDiv, xlim = c(485.4,427.4), lwd = 2, col = "darkblue")
points(midpoints, netDiv, xlim = c(485.4,427.4), pch = 21, cex = 1.3, col = "black", bg = "darkblue")


axis(1, at=seq(440, 460, 5), col = "black", line = -1.5)
axis(2, col = "black", line = 0.5, at = seq(0, max(upper.div), 0.05))

mtext("Age (Ma)", side = 1, line = 1.7, cex=1.5)
mtext("Net diversification (Lmy)", side = 2, line = 2.8, cex=1.5)


#pdf(file = "NetDiversificationGSA.pdf", height = 6, width = 5.5, useDingbats = FALSE)
#dev.off()



# Relative fossilization

# plot Sampling
plot(midpoints, relSamp, xlim = c(460,440),ylim = c(-0.00005, max(upper.s) ), type = "n", axes = F, xlab = "", ylab = "")


rect(int.start[-1], rep(0, 9), int.start[-10], rep(100, 9), col = rep(c("grey95", "grey97")), border = NA)

segments(midpoints,relSamp,midpoints, upper.s, col = "sienna1", lty = 2, lwd = 2)
segments(midpoints,relSamp,midpoints, lower.s, col = "sienna1", lty = 2, lwd = 2)


lines(midpoints, relSamp, xlim = c(440,460), lwd = 2, col = "sienna1")
points(midpoints, relSamp, xlim = c(440,460), pch = 21, cex = 1.3, col = "black", bg = "sienna1")

axis(1, col = "black", line = -1)
axis(2, col = "black", line = 0.5, at = seq(0, 2, 0.2))

mtext("Age (Ma)", side = 1, line = 1.7, cex=1.5)
mtext("Relative fossilization & sampling (r / (q + r)", side = 2, line = 2.8, cex=1.5)


#pdf(file = "RelativeSampling.pdf", height = 6, width = 5.5, useDingbats = FALSE)
#dev.off()


# Sampling probability per interval

int <- Stages[5:7,1] - Stages[5:7,2]

R <- 1 - exp(-meanPsi * int)

# plot Sampling
plot(midpoints, R, xlim = c(460,440),ylim = c(-0.0005, 1), type = "n", axes = F, xlab = "", ylab = "")


rect(int.start[-1], rep(0, 9), int.start[-10], rep(100, 9), col = rep(c("grey95", "grey97")), border = NA)


lines(midpoints, R, xlim = c(485.4,427.4), lwd = 2, col = "sienna1")
points(midpoints, R, xlim = c(485.4,427.4), pch = 21, cex = 1.3, col = "black", bg = "sienna1")

axis(1, col = "grey75", line = 0)
axis(2, col = "grey75", line = 0.5, at = seq(0, 1, 0.2))

mtext("Age (Ma)", side = 1, line = 2)
mtext("Per-interval sampling probability (R)", side = 2, line = 2.7)

cor.test(R, div, method = "spearman") # 






###

# PLOT SPECIATION & EXTINCTION WITH BOXPLOTS

###

# 
#params <- read.table(paste0(genus,".p"), header = TRUE)

# Sort the file
param <- read.table("anazygidae_BDSKY_final.log", header = T)
param <- param[1250:5001,]

# get values
netDiv <- param[,6:9]
turnover <- param[,10:13]
relSamp <- param[,14:17]


# convert to lambda, mu, and psi
meanMu <- calcMu(netDiv, turnover)
meanLambda <- calcLambda(netDiv, meanMu)
meanPsi <- calcPsi(meanMu, relSamp)

# set up plotting


#STAGES <- c("Tremadoc", "Floian", "Dap/Darr", "Sandbian", "Katian/H", "Llan")
STAGES <- c("K1", "K2", "K3", "K4/Hirnantian")

# Speciation & extinction rates
# plot
plot(midpoints, midpoints, xlim = c(452,443),ylim = c(-10,10), type = "n", axes = F, xlab = "", ylab = "")

rect(int.start[-1], rep(-10, 50), int.start[-6], rep(100, 9), col = rep(c("grey95", "grey97")), border = NA)
rect(452, rep(-10, 50), 451, rep(100, 9), col = "burlywood1", border = NA)
rect(451, rep(-10, 50), 450, rep(100, 9), col = "peachpuff", border = NA)
rect(450, rep(-10, 50), 449.5, rep(100, 9), col = "burlywood1", border = NA)
rect(449.5, rep(-10, 50), 443.07, rep(100, 9), col = "peachpuff", border = NA)

# Redefining the midpoints list to fall within the skyline interval
# Reasoning: the boxes for time bin 2 and 3 were overlapping (too close to each other on the X-axis)
midpoints.lambda = c(451.7,450.8,449.8,447)
midpoints.mu = c(451.27,450.3,449.3,446.5)


boxplot(meanLambda, col = "lightblue", add = TRUE, at = midpoints.lambda, names = STAGES, border = TRUE, boxwex = .4, outline = FALSE, frame = FALSE, axes = FALSE)
boxplot(meanMu, col = "salmon", add = TRUE, at = midpoints.mu, names = STAGES, border = TRUE, boxwex = .4, outline = FALSE, frame = FALSE, axes = FALSE)

axis(1, col = "black", line = -1, at=455:443)
axis(2, col = "black", line = 0.5, at = seq(-10, 10, 2))

mtext("Geologic time (Ma)", side = 1, line = 1.3, cex=1.5)
mtext("Speciation & extinction rates (Lmy)", side = 2, line = 2.7, cex=1.5)


#pdf(file = "DiversificationRatesGSA.pdf", height = 6, width = 6.5, useDingbats = FALSE)
dev.off()



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

# plot a timescaled MCCT

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

# use the MCC from TreeAnnotator to find the MCCT in the posterior and get branch lengths in units of time
# new function--need to move to functions.R file to source. Requires dateNodes from paleotree.
get_MCCT <- function(genus){
  
  MCCT <- read.nexus("anazygids_ne.nexus")
  # use findTree to figure out which posterior tree == MCCT
  which_is_MCCT <- suppressWarnings(findTree(MCCT, t.file))
  MCCT <- t.file[[suppressWarnings(findTree(MCCT, t.file))]]
  # rescale the tree in units of time
  #MCCT$edge.length / p.file$clockrate[which_is_MCCT]
  # re-adjust the root to send the node ages back to 'deep time'
  MCCT$root.time <- max(dateNodes(MCCT)) + getMinTaxonAge(genus)
  
  return(MCCT)
  
}

MCCT <- get_MCCT(genus)

trees <- t.file[5000:20001]
MCCT <- read.tree("catazyga_cons_MCC_after_ape.tre")
MCCT$root.time = 453.48
#for (i in 1:length(trees)){
#  oldRootAge <- max(dateNodes(trees[[i]]))
  #print(oldRootAge)
#  trees[[i]]$root.time <- oldRootAge + anchor #
  #print(M[[i]]$MCCT)
#}


#basic plot
plot(ladderize(MCCT)); axisPhylo()
slices = read.table("global_ordovician_slices_plot.txt", header=T)
slices$Name = as.factor(slices$Name)
slices = as.data.frame(slices)
summary(slices)
FA = read.table("anazygidae_FA_new.txt", header=T)
# w/ geologic timescale, can play around w/ options to make a prettier plot
geoscalePhylo(ladderize(MCCT), units=c("Epoch","Age"),x.lim = c(458, 441), 
              cex.age=2, cex.tip=1.5, boxes="Age",
              label.offset = 0.02, arotate=0, erotate=0, cex.ts=2, ages=FA, width=1.3,
              tick.scale="")



# #####
# plot the distribution of divergence times for genus origination
########


# call getCladeAge function from source
ages <- getCladeAge(trees)

z <- hist(ages,ylim = c(0,1), xlim = c(max(ages)+5,min(ages)-5), xlab = "Geologic time (Ma)", ylab = "Probability density", freq = FALSE, main = "", col = "salmon")

#get mean age
mean(ages)

# get ~estimate of time interval with largest probability density
z$mids[which(z$density == max(z$density))]

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

# Sampled ancestors

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

# get probability of a taxon being a sampled ancestor
samp = samp.Anc(trees)

#plot a histogram 
k <- hist(samp.Anc(trees), col = "darkgrey", xlab = "Pr(Ancestral taxon)", main = "")
#pdf(file = "SAs.pdf", height =5.5, width = 6, useDingbats = FALSE)
#dev.off()


####

# add time to zero-length branches (i.e., sampled ancestors). 'blens' parameter sets length for ZLB (recommended to be very small but non-zero!)


# Example for a single tree
MCCT$edge.length # the zero-length branches?
# add time
new.MCCT <- modify.edges(MCCT, blens = 0.01)
new.MCCT$edge.length # now they're gone & you can use PCMs without divide by zero / matrix inversion issues!


# also works for a distribution of trees
new.trees <- modify.edges(trees, blens = 0.01)