setwd("D:/Dropbox/Research_Projects/Awaji Hadrosaur/Manuscripts/Supplements")
library(Claddis)
library(paleotree)
library(geiger)
library(phytools)

#========================================================
data <- read_nexus_matrix("Data S1.nex")
data <- prune_cladistic_matrix(data, 
                               taxa2prune = "Lophorhothon_ontogeny")
#Inputs for the analyses=========================================================
#setting the lower limit of the character completeness
limit <- 0.5
#defining number of trees to be replicated
ntree <- 100
#deining time-scaling methods
TSM <- "equal"
#Input pectoral girdle + forelimb character numbers
characters <- c(248:286)

is_yamatosaurus <- FALSE #includes Yamatohadros in the analysis if TRUE
strict_or_majority <- "Majority" #consensus rule
#===============================================================================
Matrix <- matrix(0, nrow(data$matrix_1$matrix), length(characters))
rownames(Matrix)<-rownames(data$matrix_1$matrix)
colnames(Matrix)<-characters

for (j in 1:nrow(data$matrix_1$matrix)){
  for (i in colnames(Matrix)){
    ifelse(is.na(data$matrix_1$matrix[j,as.numeric(i)]), Matrix, Matrix[j, paste(i)]<-Matrix[j, paste(i)] + 1) 
    #counting number of NA (unknown) characters
  }
}

temp <- rownames(Matrix[apply(Matrix, 1, mean)<=limit,]) #listing taxon names that is scored less than the limit
if (is_yamatosaurus == TRUE){
  temp <- temp[temp != "Yamatosaurus_izanagii"]
}
allMatrix <- prune_cladistic_matrix(data, 
                                       taxa2prune = c(temp), 
                                       characters2prune = c(1:ncol(data$matrix_1$matrix))[-c(as.numeric(colnames(Matrix)))])

alltrees <- read.nexus("Data S1.nex")
age <- read.csv("Data S2.csv", row.names = 1)
obj <- name.check(alltrees[[1]], age)
obj

con_tree <- consensus(alltrees, p = if (strict_or_majority == "Strict"){1} else {0.5})
con_tree <- multi2di(con_tree, random = F)

obj <- name.check(con_tree, allMatrix$matrix_1$matrix)
con_tree <- drop.tip(con_tree, obj$tree_not_data)
obj <- name.check(con_tree, age)
age <- age[!rownames(age)%in%c(obj$data_not_tree),]
obj <- name.check(con_tree, age)
obj

con_trees <- timePaleoPhy(con_tree, age, type = "equal", 
                          vartime = 1, ntrees = ntree, randres = F, 
                          dateTreatment = "minMax")
con_trees <- lapply(con_trees, ladderize)

results <- list()
#Branch evolutionary rate analysis 
for (i in 1:length(con_trees)) {
  results[[i]] <- test_rates(time_tree = con_trees[[i]], 
                                     cladistic_matrix = allMatrix, 
                                     time_bins = seq(from = con_trees[[i]]$root.time, to = 0, length.out = 5), 
                                     branch_partitions  = lapply(X = as.list(x = 1:nrow(con_trees[[i]]$edge)), as.list), 
                                     change_times = "random", 
                                     alpha = 0.01, 
                                     polymorphism_state =  "random", 
                                     uncertainty_state =  "missing",
                                     inapplicable_state  = "missing", 
                                     time_binning_approach = "lloyd")
}

yamato <- ifelse(is_yamatosaurus, "w_yamato_", "wo_yamato_")
save.image(paste0(TSM, yamato, ".RData" ))
#load(paste0(TSM, ".RData"))

summ_results <- as.data.frame(matrix(NA, 
                                       length(results[[1]]$branch_test_results)*
                                         nrow(summary(results)), 
                                       6, 
                                       dimnames = list(NULL, 
                                                       c("treeID","ER_present", 
                                                         "ER_other", "AIC", 
                                                         "AICc", "Branch_ID"))))

for (j in 1:nrow(summary(results))){
  for (i in 1:length(results[[j]]$branch_test_results)){
    summ_results[i + length(results[[j]]$branch_test_results)*(j-1), 1] <- j
    summ_results[i + length(results[[j]]$branch_test_results)*(j-1), 2] <- results[[j]]$branch_test_results[[i]]$rates[1]
    summ_results[i + length(results[[j]]$branch_test_results)*(j-1), 3] <- results[[j]]$branch_test_results[[i]]$rates[2]
    summ_results[i + length(results[[j]]$branch_test_results)*(j-1), 4] <- results[[j]]$branch_test_results[[i]]$aic
    summ_results[i + length(results[[j]]$branch_test_results)*(j-1), 5] <- results[[j]]$branch_test_results[[i]]$aicc
    summ_results[i + length(results[[j]]$branch_test_results)*(j-1), 6] <- i
  }
}

median_values <- matrix(NA, max(summ_results$Branch_ID), ncol(summ_results))
colnames(median_values) <- colnames(summ_results)
for (i in 1:max(summ_results$Branch_ID)){
  median_values[i,]<-apply(summ_results[summ_results$Branch_ID==i,], 2, median)
}

gradTree <- con_trees[[1]]
gradTree$edge.length<-rep(1, length(con_trees[[1]]$edge.length))

#exporting
pdf(paste0(TSM, yamato,  "_corrected_tree.pdf"))
par(oma = c(0, 0, 0, 0), mfrow = c(1, 1))
plotBranchbyTrait(gradTree, median_values[,2], mode = "edges", palette = "rainbow", cex = 0.8)
edgelabels()
dev.off()