############################################################################
#data & R prep

#load Bayesian stats packages
library(pacman)
p_load(rstan,brms)
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

#set workspace
#load data
data <- read.csv("bonobo fWHR dataset_Martin et al.csv")

#standardize fWHR
data$fWHR<-(data$fWHR - mean(data$fWHR))/sd(data$fWHR)

#center normDS (original David's score measure) within groups
library(dplyr)
mean_DS<-summarise((group_by(data,Group)), DS = mean(normDS))
mean_DS<-mean_DS[[2]]
data$meanDS<-mean_DS[data$Group]
data$wgDS<-data$normDS-data$meanDS

#aggregate to mean fWHR
data.agg <- aggregate(fWHR ~ Name + Group + Sex + Age + AssR + wgDS,
                      FUN = mean, data = data)

#get measurement error (z-scores)
data.agg$sd <- aggregate(fWHR ~ Name, FUN = sd, data = data)[,2]

#mean imputation of expected error for single photo subjects
data.agg$sd[is.na(data.agg$sd)]<-mean(!is.na(data.agg$sd))

#standardize to 2 SDs
data.agg$AssR<-scale(data.agg$AssR, scale=sd(data.agg$AssR)*2)
data.agg$wgDS<-scale(data.agg$wgDS, scale=sd(data.agg$wgDS)*2)
data.agg$Age<-scale(data.agg$Age,scale=sd(data.agg$Age)*2)

#create data frame with NA body weight
library(data.table)
setDT(data.agg); setDT(data)
data.agg[data, on = .(Name = Name), weight := weight]
data<-data.frame(data)
data.agg<-data.frame(data.agg)

#standardize to 2 SDs on non-missing data mean and sd
data.agg$weight<-(data.agg$weight-mean(data.agg$weight[!is.na(data.agg$weight)]))/
  (2*sd(data.agg$weight[!is.na(data.agg$weight)]))

############################################################################
#main effect model w/o body weight

m0<- brm(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS,
         family=gaussian(), data=data.agg,
         prior=c(prior("normal(0,2)",class="Intercept"),
                  prior("normal(0,2)",class="b"),
                 prior("cauchy(0,2)",class="sigma")),
         warmup=2000,iter=5000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#save
saveRDS(m0, "m0.RDS")

#summary
fixef(m0, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m0, "SexMale > 0")


#Cohen's f2
#R2 of full model
R2tot <- bayes_R2(m0,summary=FALSE)

#R2 without sex
m0.2<-update(m0, formula. = ~ . - Sex)
R2red <- bayes_R2(m0.2,summary=FALSE)

#f2
f2<-(R2tot-R2red)/(1-R2tot)
median(f2)

############################################################################
#main effect model w/ body weight imputation####

#imputation model
weight.m<-bf(weight|mi() ~ fWHR + Sex + Age + AssR + wgDS) + gaussian()

#fWHR model
fwhr.m<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS + mi(weight))+ gaussian()

#multi-response prior
prior=c(prior("normal(0,2)",class="Intercept", resp="fWHR"),
        prior("normal(0,2)",class="b",resp="fWHR"),
        prior("cauchy(0,2)",class="sigma",resp="fWHR"),
        prior("normal(0,2)",class="Intercept", resp="weight"),
        prior("cauchy(0,2)",class="sigma",resp="weight"))

#main model
m1<- brm(weight.m + fwhr.m + set_rescor(FALSE),
         data=data.agg,prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#save model
saveRDS(m1,"m1.RDS")

#summary
fixef(m1, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m1, "fWHR_SexMale > 0")
hypothesis(m1, "fWHR_miweight > 0",class="bsp")
hypothesis(m1, "fWHR_Age < 0")
hypothesis(m1, "fWHR_AssR > 0")
hypothesis(m1, "fWHR_wgDS > 0")


#Cohen's f2

#R2 of full model
R2tot<-bayes_R2(m1, resp="fWHR", summary=FALSE)
median(R2tot)

#refit without predictors
#can't use update() for multivariate models

#Sex
fwhr.m<-bf(fWHR|se(sd, sigma=TRUE) ~ Age + AssR + wgDS + mi(weight))+ gaussian()
m1.2<- brm(weight.m + fwhr.m + set_rescor(FALSE),
         data=data.agg,prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))
R2red<- bayes_R2(m1.2, resp="fWHR", summary=FALSE)
f2<-(R2tot-R2red)/(1-R2tot)
median(f2)
saveRDS(m1.2,"m1_2.RDS")

#weight
fwhr.m<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS)+ gaussian()
m1.3<- brm(weight.m + fwhr.m + set_rescor(FALSE),
           data=data.agg,prior=prior,
           warmup=2000,iter=4000, chains=4, seed=9,
           control=list(adapt_delta=0.99, max_treedepth=20))
R2red<- bayes_R2(m1.3, resp="fWHR", summary=FALSE)
f2<-(R2tot-R2red)/(1-R2tot)
median(f2)
saveRDS(m1.3,"m1_3.RDS")

#AssR
fwhr.m<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + wgDS + mi(weight))+ gaussian()
m1.4<- brm(weight.m + fwhr.m + set_rescor(FALSE),
           data=data.agg,prior=prior,
           warmup=2000,iter=4000, chains=4, seed=9,
           control=list(adapt_delta=0.99, max_treedepth=20))
R2red<- bayes_R2(m1.4, resp="fWHR", summary=FALSE)
f2<-(R2tot-R2red)/(1-R2tot)
median(f2)
saveRDS(m1.4,"m1_4.RDS")

#wgDS
fwhr.m<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + mi(weight))+ gaussian()
m1.5<- brm(weight.m + fwhr.m + set_rescor(FALSE),
           data=data.agg,prior=prior,
           warmup=2000,iter=4000, chains=4, seed=9,
           control=list(adapt_delta=0.99, max_treedepth=20))
R2red<- bayes_R2(m1.5, resp="fWHR", summary=FALSE)
f2<-(R2tot-R2red)/(1-R2tot)
median(f2)
saveRDS(m1.5,"m1_5.RDS")

#Age
fwhr.m<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + AssR + wgDS + mi(weight))+ gaussian()
m1.6<- brm(weight.m + fwhr.m + set_rescor(FALSE),
           data=data.agg,prior=prior,
           warmup=2000,iter=4000, chains=4, seed=9,
           control=list(adapt_delta=0.99, max_treedepth=20))
R2red<- bayes_R2(m1.6, resp="fWHR", summary=FALSE)
f2<-(R2tot-R2red)/(1-R2tot)
median(f2)
saveRDS(m1.6,"m1_6.RDS")

############################################################################
#model + assR*Sex

#fWHR model
fwhr.m2<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS +
              mi(weight) + AssR*Sex)+ gaussian()
#model
m2<- brm(weight.m + fwhr.m2 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))
#summary
fixef(m2, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m2, "fWHR_SexMale:AssR < 0")

#Cohens f2
R2red<- bayes_R2(m2, resp="fWHR", summary=FALSE)
f2<-(R2red-R2tot)/(1-R2red) #switch order due to added vs removed effect
median(f2)

saveRDS(m2, "m2.RDS")

############################################################################
#model + wgDS*Sex

#fWHR model
fwhr.m3<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS +
              mi(weight) + wgDS*Sex)+ gaussian()

#model
m3<- brm(weight.m + fwhr.m3 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#summary
fixef(m3, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m3, "fWHR_SexMale:wgDS < 0")

#Cohens f2
R2red<- bayes_R2(m3, resp="fWHR", summary=FALSE)
f2<-(R2red-R2tot)/(1-R2red)
median(f2)

#save
saveRDS(m3, "m3.RDS")

############################################################################
#model + AssR*wgDS

#fWHR model
fwhr.m4<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS + AssR*wgDS +
              mi(weight))+ gaussian()

#model
m4<- brm(weight.m + fwhr.m4 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#summary
fixef(m4, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m4, "fWHR_AssR:wgDS > 0")

#Cohens f2
R2red<- bayes_R2(m4, resp="fWHR", summary=FALSE)
f2<-(R2red-R2tot)/(1-R2red)
median(f2)

#save
saveRDS(m4, "m4.RDS")

############################################################################
#model + age*Assr

#fWHR model
fwhr.m5<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS +
              mi(weight) + Age*AssR)+ gaussian()

#model
m5<- brm(weight.m + fwhr.m5 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#summary
fixef(m5, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m5, "fWHR_Age:AssR > 0")

#Cohens f2
R2red<- bayes_R2(m5, resp="fWHR", summary=FALSE)
f2<-(R2red-R2tot)/(1-R2red)
median(f2)

#save
saveRDS(m5, "m5.RDS")

############################################################################
#model + age*wgDS

#fWHR model
fwhr.m6<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS +
              mi(weight) + Age*wgDS)+ gaussian()

#model
m6<- brm(weight.m + fwhr.m6 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#summary
fixef(m6, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m6, "fWHR_Age:wgDS > 0")

#Cohens f2
R2red<- bayes_R2(m6, resp="fWHR", summary=FALSE)
f2<-(R2red-R2tot)/(1-R2red)
median(f2)

#save
saveRDS(m6, "m6.RDS")

############################################################################
#model + age*sex*AssR

#fWHR model
fwhr.m7<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS +
              mi(weight) + Sex*Age*AssR)+ gaussian()

#model
m7<- brm(weight.m + fwhr.m7 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#summary
fixef(m7, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m7, "fWHR_SexMale:Age:AssR > 0")

#Cohens f2
R2red<- bayes_R2(m7, resp="fWHR", summary=FALSE)
f2<-(R2red-R2tot)/(1-R2red)
median(f2)

#save
saveRDS(m7, "m7.RDS")

############################################################################
#model + Sex*Age*wgDS

#fWHR model
fwhr.m8<-bf(fWHR|se(sd, sigma=TRUE) ~ Sex + Age + AssR + wgDS +
              mi(weight) + Sex*Age*wgDS)+ gaussian()

#model
m8<- brm(weight.m + fwhr.m8 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#summary
fixef(m8, summary=TRUE, robust=TRUE, probs=c(0.05,0.95))
hypothesis(m8, "fWHR_SexMale:Age:wgDS < 0")

#Cohens f2
R2red<- bayes_R2(m8, resp="fWHR", summary=FALSE)
f2<-(R2red-R2tot)/(1-R2red)
median(f2)

#save
saveRDS(m8, "m8.RDS")

############################################################################
#model + t2(age*AssR, by = Sex)

#recalculate main model w/o measurement error term to facilitate WAIC calculation
fwhr.m1nome<-bf(fWHR ~ Sex + Age + AssR + wgDS + mi(weight))+ gaussian()
m1.nome<- brm(weight.m + fwhr.m1nome + set_rescor(FALSE),
         data=data.agg,prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#tensor model
fwhr.m9<-bf(fWHR ~ t2(Age,AssR, by=Sex) + wgDS + mi(weight))+ gaussian()
m9<- brm(weight.m + fwhr.m9 + set_rescor(FALSE),
         data=data.agg,
         prior=prior,
         warmup=2000,iter=4000, chains=4, seed=9,
         control=list(adapt_delta=0.99, max_treedepth=20))

#compare IC
compare_ic(WAIC(m1.nome,resp="fWHR"), WAIC(m9,resp="fWHR"))

#save
saveRDS(m1.nome, "m1nome.RDS")
saveRDS(m9, "m9.RDS")


############################################################################
#model + t2(age*wgDS, by = Sex)

#tensor model
fwhr.m10<-bf(fWHR ~ t2(Age,wgDS, by=Sex) + AssR + mi(weight))+ gaussian()
m10<- brm(weight.m + fwhr.m10 + set_rescor(FALSE),
           data=data.agg,
           prior=prior,
           warmup=2000,iter=4000, chains=4, seed=9,
           control=list(adapt_delta=0.99, max_treedepth=20))

#compare IC
compare_ic(WAIC(m1.nome,resp="fWHR"), WAIC(m10,resp="fWHR"))

#save
saveRDS(m10, "m10.RDS")

############################################################################
#main effect model w/ random intercept & slopes for zoo
#remove measurement error for WAIC comparison

#imputation model
weight.m2<-bf(weight|mi() ~ fWHR + Sex + Age + AssR + wgDS + (AssR + wgDS||Group)) + gaussian()

#fWHR model
fwhr.m11<-bf(fWHR ~ Sex + Age + AssR + wgDS + mi(weight) + (AssR + wgDS||Group))+ gaussian()

#main model
m11<- brm(weight.m2 + fwhr.m11 + set_rescor(FALSE),
         data=data.agg,
         prior=c(prior("normal(0,2)",class="Intercept", resp="fWHR"),
                 prior("normal(0,2)",class="b",resp="fWHR"),
                 prior("cauchy(0,2)",class="sd",resp="fWHR"),
                 prior("cauchy(0,2)",class="sigma",resp="fWHR"),
                 prior("normal(0,2)",class="Intercept", resp="weight"),
                 prior("normal(0,2)",class="b",resp="weight"),
                 prior("cauchy(0,2)",class="sd",resp="weight"),
                 prior("cauchy(0,2)",class="sigma",resp="weight")),
         warmup=6000,iter=8000, chains=4, seed=1,
         control=list(adapt_delta=0.99, stepsize = 0.0001, max_treedepth=30))

#variance of zoo intercepts
post<-posterior_samples(m11)
zoovar<-post$sd_Group__fWHR_Intercept^2 #sd^2=var
median(zoovar);mad(zoovar)

#model comparison
compare_ic(WAIC(m1.nome,resp="fWHR"), WAIC(m11,resp="fWHR"))

#save
saveRDS(m11, "m11.RDS")


############################################################################
#simple correlations

library(psych)

#missing data

#weight ~ fWHR
cor(data.agg$weight, data.agg$fWHR, use="pairwise.complete.obs")

#weight ~ Sex
biserial(data.agg$weight, data.agg$Sex)

#fWHR ~ Sex
biserial(data.agg$fWHR, data.agg$Sex)


#if you have any questions or concerns
#please contact Jordan S. Martin
#jsm.primatology@gmail.com