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

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

##### load data #####
load("Rdata_hind/230614_data_ext_hind.RData")
load("Rdata_hind/simmap_hind/simmap_loc4_mcc_ext_s50_hind.Rdata")
load("Rdata_hind/simmap_hind/simmap_loc3a_mcc_ext_s50_hind.Rdata")
load("Rdata_hind/simmap_hind/simmap_loc3b_mcc_ext_s50_hind.Rdata")
load("Rdata_hind/simmap_hind/simmap_loc3c_mcc_ext_s50_hind.Rdata")

##### variables #####
loc4 <- data_hind$Loc_mode_4groups
names(loc4) <- data_hind[,1]

loc4_ext <- data_hind$Loc_mode_4groups_ext
names(loc4_ext) <- data_hind[,1]

loc3a <- data_hind$Loc_mode_3agroups
names(loc3a) <- data_hind[,1]

loc3b <- data_hind$Loc_mode_3bgroups
names(loc3b) <- data_hind[,1]

loc3c <- data_hind$Loc_mode_3cgroups
names(loc3c) <- data_hind[,1]

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

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

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

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


#### hindlimb mvMorph ####
# BM1 - traits follow brownian motion
hindlimb_ext_BM1_pca_ext <- list(length=length(simmap_loc4_mcc_ext_s50_hind))
for(i in 1:length(simmap_loc4_mcc_ext_s50_hind)){
  hindlimb_ext_BM1_pca_ext[[i]] <- mvBM(tree = simmap_loc4_mcc_ext_s50_hind[[i]], data = hindlimb_lsr_ext_pca_1_3, model= "BM1", param= list(smean = TRUE))
  print(i)
}
save(hindlimb_ext_BM1_pca_ext, file="Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_BM1_pca_ext.Rdata")

# OU1
hindlimb_ext_OU1_pca_ext <- list(length=length(simmap_loc4_mcc_ext_s50_hind))
for(i in 1:length(simmap_loc4_mcc_ext_s50_hind)){
  hindlimb_ext_OU1_pca_ext[[i]] <- mvOU(tree = simmap_loc4_mcc_ext_s50_hind[[i]], data = hindlimb_lsr_ext_pca_1_3, model= "OU1", param= list(root = "stationary"))
  print(i)
}
save(hindlimb_ext_OU1_pca_ext, file="Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OU1_pca_ext.Rdata")

# OUM_loc4
hindlimb_ext_OUM_loc4_pca_ext <- list(length=length(simmap_loc4_mcc_ext_s50_hind))
for(i in 1:length(simmap_loc4_mcc_ext_s50_hind)){
  hindlimb_ext_OUM_loc4_pca_ext[[i]] <- mvOU(tree = simmap_loc4_mcc_ext_s50_hind[[i]], data = hindlimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(hindlimb_ext_OUM_loc4_pca_ext, file="Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc4_pca_ext.Rdata")


# OUM_loc3a
hindlimb_ext_OUM_loc3a_pca_ext <- list(length=length(simmap_loc3a_mcc_ext_s50_hind))
for(i in 1:length(simmap_loc3a_mcc_ext_s50_hind)){
  hindlimb_ext_OUM_loc3a_pca_ext[[i]] <- mvOU(tree = simmap_loc3a_mcc_ext_s50_hind[[i]], data = hindlimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(hindlimb_ext_OUM_loc3a_pca_ext, file="Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc3a_pca_ext.Rdata")


# OUM_loc3b
hindlimb_ext_OUM_loc3b_pca_ext <- list(length=length(simmap_loc3b_mcc_ext_s50_hind))
for(i in 1:length(simmap_loc3b_mcc_ext_s50_hind)){
  hindlimb_ext_OUM_loc3b_pca_ext[[i]] <- mvOU(tree = simmap_loc3b_mcc_ext_s50_hind[[i]], data = hindlimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(hindlimb_ext_OUM_loc3b_pca_ext, file="Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc3b_pca_ext.Rdata")

# OUM_loc3c
hindlimb_ext_OUM_loc3c_pca <- list(length=length(simmap_loc3c_mcc_ext_s50_hind))
for(i in 1:length(simmap_loc3c_mcc_ext_s50_hind)){
  hindlimb_ext_OUM_loc3c_pca[[i]] <- mvOU(tree = simmap_loc3c_mcc_ext_s50_hind[[i]], data = hindlimb_lsr_ext_pca_1_3, model= "OUM", param= list(root = "stationary"))
  print(i)
}
save(hindlimb_ext_OUM_loc3c_pca, file="Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc3c_pca.Rdata")

##### AIC table for hindlimb_exts #####
load("Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_BM1_pca_ext.Rdata")
load("Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OU1_pca_ext.Rdata")
load("Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc4_pca_ext.Rdata")
load("Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc3a_pca_ext.Rdata")
load("Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc3b_pca_ext.Rdata")
load("Rdata_hind/mvMorph_models_hind_ext/hindlimb_ext_OUM_loc3c_pca.Rdata")


hindlimb_ext_PC_aic.table <- matrix(,length(hindlimb_ext_OUM_loc4_pca_ext),6,dimnames=list(c(1: length(hindlimb_ext_OUM_loc4_pca_ext)),c("mvBM1", "mvOU1", "mvOUM_loc3a", "mvOUM_loc3b", "mvOUM_loc3c", "mvOUM_loc4")))
for(i in 1:length(hindlimb_ext_OUM_loc4_pca_ext)){
  hindlimb_ext_PC_aic.table[i,1] <- hindlimb_ext_BM1_pca_ext[[i]]$AICc
  hindlimb_ext_PC_aic.table[i,2] <- hindlimb_ext_OU1_pca_ext[[i]]$AICc
  hindlimb_ext_PC_aic.table[i,3] <- hindlimb_ext_OUM_loc3a_pca_ext[[i]]$AICc
  hindlimb_ext_PC_aic.table[i,4] <- hindlimb_ext_OUM_loc3b_pca_ext[[i]]$AICc
  hindlimb_ext_PC_aic.table[i,5] <- hindlimb_ext_OUM_loc3c_pca[[i]]$AICc
  hindlimb_ext_PC_aic.table[i,6] <- hindlimb_ext_OUM_loc4_pca_ext[[i]]$AICc
  
}
hindlimb_ext_PC_aic.table

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


# plot optima
# theta matrix
hindlimb_ext_OUM_loc4_pca_ext_theta <- list()
for(i in 1:length(hindlimb_ext_OUM_loc4_pca_ext)){
  hindlimb_ext_OUM_loc4_pca_ext_theta[[i]] <- matrix(hindlimb_ext_OUM_loc4_pca_ext[[i]]$theta)}

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

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

hindlimb_ext_OUM_loc4_pca_ext_sigma_mean <- rowMeans(array( unlist(hindlimb_ext_OUM_loc4_pca_ext_sigma), c(4,4,100)), dims = 2)

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

hindlimb_ext_OUM_loc4_pca_ext_alpha_mean <- rowMeans(array( unlist(hindlimb_ext_OUM_loc4_pca_ext_alpha), c(4,4,100)), dims = 2)
log(2)/hindlimb_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(hindlimb_lsr_ext_pca_1_3[,1], hindlimb_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(hindlimb_ext_OUM_loc4_pca_ext_theta_mean[1,1], hindlimb_ext_OUM_loc4_pca_ext_theta_mean[1,2], cex = 4, bg = "#E6C060", pch = 25)
points(hindlimb_ext_OUM_loc4_pca_ext_theta_mean[2,1], hindlimb_ext_OUM_loc4_pca_ext_theta_mean[2,2], cex = 4, bg = "#419D78", pch = 25)
points(hindlimb_ext_OUM_loc4_pca_ext_theta_mean[3,1], hindlimb_ext_OUM_loc4_pca_ext_theta_mean[3,2], cex = 4, bg = "#AC97AF", pch = 25)
points(hindlimb_ext_OUM_loc4_pca_ext_theta_mean[4,1], hindlimb_ext_OUM_loc4_pca_ext_theta_mean[4,2], cex = 4, bg = "#274754", pch = 25)