##### MvMorph multivariate models with PC1-3 only #####
##### extant+extinct dataset #####

##### packages #####
library(tidyverse)
library(plyr)
library(plotrix)
library(geiger)
library(phytools)
library(mvMORPH)

load("Rdata_fore/230614_data_fore_ext.RData")
load("Rdata_fore/simmap_fore/simmap_loc4_mcc_ext_s50.Rdata")
load("Rdata_fore/simmap_fore/simmap_loc3a_mcc_ext_s50.Rdata")
load("Rdata_fore/simmap_fore/simmap_loc3b_mcc_ext_s50.Rdata")
load("Rdata_fore/simmap_fore/simmap_loc3c_mcc_ext_s50.Rdata")

#simmaps

simmap_loc4_mcc_ext_s50 <- make.simmap(tree_prune,loc4, nsim=50) #10 simmap simulations across 1k trees
save(simmap_loc4_mcc_ext_s50, file="simmap_loc4_mcc_ext_s50.Rdata")

simmap_loc3a_mcc_ext_s50 <- make.simmap(tree_prune,loc3a, nsim=50) #10 simmap simulations across 1k trees
save(simmap_loc3a_mcc_ext_s50, file="simmap_loc3a_mcc_ext_s50.Rdata")

simmap_loc3b_mcc_ext_s50 <- make.simmap(tree_prune,loc3b, nsim=50) #10 simmap simulations across 1k trees
save(simmap_loc3b_mcc_ext_s50, file="simmap_loc3b_mcc_ext_s50.Rdata")

simmap_loc3c_mcc_ext_s50 <- make.simmap(tree_prune,loc3c, nsim=50) #10 simmap simulations across 1k trees
save(simmap_loc3c_mcc_ext_s50, file="simmap_loc3c_mcc_ext_s50.Rdata")


######### forelimb Standard PCA ######### 
forelimb_lsr_ext_pca <- princomp(forelimb_lsr_ext, cor = FALSE)
plot(forelimb_lsr_ext_pca)
biplot(forelimb_lsr_ext_pca)
summary(forelimb_lsr_ext_pca)
forelimb_lsr_ext_pca_1_3 <- data.frame(forelimb_lsr_ext_pca$scores[,1:5])
save(forelimb_lsr_ext_pca_1_3, file="forelimb_lsr_ext_pca_1_3.Rdata")
loadings <- data.frame(forelimb_lsr_ext_pca$loadings[,1:5])


#### forelimb mvMorph ####
# BM1 - traits follow brownian motion
forelimb_ext_BM1_pca_ext <- list(length=length(simmap_loc4_mcc_ext_s50))
for(i in 1:length(simmap_loc4_mcc_ext_s50)){
  forelimb_ext_BM1_pca_ext[[i]] <- mvBM(tree = simmap_loc4_mcc_ext_s50[[i]], data = forelimb_lsr_ext_pca_1_3, model= "BM1", param= list(smean = TRUE))
  print(i)
}
save(forelimb_ext_BM1_pca_ext, file="Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_BM1_pca_ext.Rdata")

# OU1
forelimb_ext_OU1_pca_ext <- list(length=length(simmap_loc4_mcc_ext_s50))
for(i in 1:length(simmap_loc4_mcc_ext_s50)){
  forelimb_ext_OU1_pca_ext[[i]] <- mvOU(tree = simmap_loc4_mcc_ext_s50[[i]], data = forelimb_lsr_ext_pca_1_3, model= "OU1", param= list(root = "stationary"))
  print(i)
}
save(forelimb_ext_OU1_pca_ext, file="Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OU1_pca_ext.Rdata")

# OUM_loc4
forelimb_ext_OUM_loc4_pca_ext <- list(length=length(simmap_loc4_mcc_ext_s50))
for(i in 1:length(simmap_loc4_mcc_ext_s50)){
  forelimb_ext_OUM_loc4_pca_ext[[i]] <- mvOU(tree = simmap_loc4_mcc_ext_s50[[i]], data = forelimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(forelimb_ext_OUM_loc4_pca_ext, file="Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc4_pca_ext.Rdata")


# OUM_loc3a
forelimb_ext_OUM_loc3a_pca_ext <- list(length=length(simmap_loc3a_mcc_ext_s50))
for(i in 1:length(simmap_loc3a_mcc_ext_s50)){
  forelimb_ext_OUM_loc3a_pca_ext[[i]] <- mvOU(tree = simmap_loc3a_mcc_ext_s50[[i]], data = forelimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(forelimb_ext_OUM_loc3a_pca_ext, file="Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc3a_pca_ext.Rdata")


# OUM_loc3b
forelimb_ext_OUM_loc3b_pca_ext <- list(length=length(simmap_loc3b_mcc_ext_s50))
for(i in 1:length(simmap_loc3b_mcc_ext_s50)){
  forelimb_ext_OUM_loc3b_pca_ext[[i]] <- mvOU(tree = simmap_loc3b_mcc_ext_s50[[i]], data = forelimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(forelimb_ext_OUM_loc3b_pca_ext, file="Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc3b_pca_ext.Rdata")

# OUM_loc3c
forelimb_ext_OUM_loc3c_pca_ext <- list(length=length(simmap_loc3c_mcc_ext_s50))
for(i in 1:length(simmap_loc3c_mcc_ext_s50)){
  forelimb_ext_OUM_loc3c_pca_ext[[i]] <- mvOU(tree = simmap_loc3c_mcc_ext_s50[[i]], data = forelimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(forelimb_ext_OUM_loc3c_pca_ext, file="Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc3c_pca_ext.Rdata")


##### AIC table for forelimb_exts #####

load("Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_BM1_pca_ext.Rdata")
load("Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OU1_pca_ext.Rdata")
load("Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc4_pca_ext.Rdata")
load("Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc3a_pca_ext.Rdata")
load("Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc3b_pca_ext.Rdata")
load("Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_OUM_loc3c_pca_ext.Rdata")


forelimb_ext_PC_aic.table <- matrix(,length(forelimb_ext_OUM_loc4_pca_ext),6,dimnames=list(c(1: length(forelimb_ext_OUM_loc4_pca_ext)),c("mvBM1", "mvOU1", "mvOUM_loc3a", "mvOUM_loc3b", "mvOUM_loc3c", "mvOUM_loc4")))
for(i in 1:length(forelimb_ext_OUM_loc4_pca_ext)){
  forelimb_ext_PC_aic.table[i,1] <- forelimb_ext_BM1_pca_ext[[i]]$AICc
  forelimb_ext_PC_aic.table[i,2] <- forelimb_ext_OU1_pca_ext[[i]]$AICc
  forelimb_ext_PC_aic.table[i,3] <- forelimb_ext_OUM_loc3a_pca_ext[[i]]$AICc
  forelimb_ext_PC_aic.table[i,4] <- forelimb_ext_OUM_loc3b_pca_ext[[i]]$AICc
  forelimb_ext_PC_aic.table[i,5] <- forelimb_ext_OUM_loc3c_pca_ext[[i]]$AICc
  forelimb_ext_PC_aic.table[i,6] <- forelimb_ext_OUM_loc4_pca_ext[[i]]$AICc
  
}
forelimb_ext_PC_aic.table

forelimb_ext_PC_aicw <- 
  geiger::aicw(c(
    mean(forelimb_ext_PC_aic.table[,1]),
    mean(forelimb_ext_PC_aic.table[,2]),
    mean(forelimb_ext_PC_aic.table[,3]),
    mean(forelimb_ext_PC_aic.table[,4]),
    mean(forelimb_ext_PC_aic.table[,5]),
    mean(forelimb_ext_PC_aic.table[,6])
  ))
rownames(forelimb_ext_PC_aicw) <- c("mvBM1", "mvOU1", "mvOUM_loc3a", "mvOUM_loc3b", "mvOUM_loc3c", "mvOUM_loc4")
forelimb_ext_PC_aicw
write.csv(forelimb_ext_PC_aicw, file="Rdata_fore/mvMorph_models_fore_ext/forelimb_ext_PC_aicw.csv")


# plots

# theta matrix
forelimb_ext_OUM_loc4_pca_ext_theta <- list()
for(i in 1:length(forelimb_ext_OUM_loc4_pca_ext)){
  forelimb_ext_OUM_loc4_pca_ext_theta[[i]] <- matrix(forelimb_ext_OUM_loc4_pca_ext[[i]]$theta)}

forelimb_ext_OUM_loc4_pca_ext_theta_mean <- rowMeans(array( unlist(forelimb_ext_OUM_loc4_pca_ext_theta), c(4,5,100)), dims = 2)
rownames(forelimb_ext_OUM_loc4_pca_ext_theta_mean) <- c("generalist", "arboreal", "glider", "fly")
forelimb_ext_OUM_loc4_pca_ext_theta_mean

# sigma matrix
forelimb_ext_OUM_loc4_pca_ext_sigma <- list()
for(i in 1:length(forelimb_ext_OUM_loc4_pca_ext)){
  forelimb_ext_OUM_loc4_pca_ext_sigma[[i]] <- matrix(forelimb_ext_OUM_loc4_pca_ext[[i]]$sigma)}

forelimb_ext_OUM_loc4_pca_ext_sigma_mean <- rowMeans(array( unlist(forelimb_ext_OUM_loc4_pca_ext_sigma), c(4,4,100)), dims = 2)

# alpha matrix
forelimb_ext_OUM_loc4_pca_ext_alpha <- list()
for(i in 1:length(forelimb_ext_OUM_loc4_pca_ext)){
  forelimb_ext_OUM_loc4_pca_ext_alpha[[i]] <- matrix(forelimb_ext_OUM_loc4_pca_ext[[i]]$alpha)}

forelimb_ext_OUM_loc4_pca_ext_alpha_mean <- rowMeans(array( unlist(forelimb_ext_OUM_loc4_pca_ext_alpha), c(4,4,100)), dims = 2)
log(2)/forelimb_ext_OUM_loc4_pca_ext_alpha_mean


#### plot optima on PC ####
cols_loc4_ext <- setNames(c("#E6C060", #generalist (white)
                                     "#419D78", #arboreal
                                     "#AC97AF", #gliders (tan)
                                     "#274754"), #flying (navy)
                          levels(loc4_ext))

#PCA
plot(forelimb_lsr_ext_pca_1_3[,1], forelimb_lsr_ext_pca_1_3[,2], bg = cols_loc4_ext[as.factor(loc4_ext)], pch = c(21, 23, 24, 25, 22)[as.factor(loc4_ext)], cex = 1.5, las =1, xlim=c(-1.1, 2.5), ylim=c(-1, 1.1))

points(forelimb_ext_OUM_loc4_pca_ext_theta_mean[1,1], forelimb_ext_OUM_loc4_pca_ext_theta_mean[1,2], cex = 4, bg = "#E6C060", pch = 25)
points(forelimb_ext_OUM_loc4_pca_ext_theta_mean[2,1], forelimb_ext_OUM_loc4_pca_ext_theta_mean[2,2], cex = 4, bg = "#419D78", pch = 25)
points(forelimb_ext_OUM_loc4_pca_ext_theta_mean[3,1], forelimb_ext_OUM_loc4_pca_ext_theta_mean[3,2], cex = 4, bg = "#AC97AF", pch = 25)
points(forelimb_ext_OUM_loc4_pca_ext_theta_mean[4,1], forelimb_ext_OUM_loc4_pca_ext_theta_mean[4,2], cex = 4, bg = "#274754", pch = 25)