# the functions for running the analyses are in this file
source("C:\\Users\\dcintron\\Desktop\\code-star-v2\\code-star\\primary-analyses-code\\functions-gest-overall-noint-mlobs.R")

# data
star_subset <- readRDS(file = "C:\\Users\\dcintron\\Desktop\\code-star-v2\\code-star\\data\\star_subset_miss.rds")

# bootstrap Point Estimates and CIs for CDE at 0 (interactions)
# setup cluster for parallel processing
library(parallel)
numCores <- detectCores(logical = F)
cl <- makeCluster(numCores - 2) 
clusterEvalQ(cl, {
  library(lavaan)
  library(formatR)
  library(knitr)
  library(kableExtra)
  library(data.table)
  library(xtable)
  library(geepack)
  library(DTRreg)
  library(mice)
})
clusterExport(cl, list("star_subset", "o_model_v","Fuller_MoM_p_ZWint",
                       "Gest_T3_Bart_MoM_latentMY_XMint_M0obs","lavPredict_raw",
                       "OneEst","RegCoefTable"))
nboots <- 2e3
set.seed(1231)
ptm <- proc.time()[3]

res_boots <- parLapply(cl, 1:nboots, function(boot_id) {
  
  # data imputation and mbar_calculation
  star_subset <- star_subset[-c(1:18,23:45,49)]
  
  # removing irrelevant attribute of variables (a result of scaling process)
  attr(star_subset$W1_DISC_DAILY_AVG_Z, "dim") <- NULL
  attr(star_subset$W1_DISC_LIFE_AVG_Z, "dim") <- NULL
  attr(star_subset$W1_PROM_DEPR_AVG_Z, "dim") <- NULL
  
  # Impute data
  tryCatch({
  imp1 <- mice::mice(star_subset, method = "pmm", maxit = 1, m = 1)
  star_subset_imp <- mice::complete(imp1)
  
  # rename variables for Loh et al approach
  star_subset_imp$m1 <- star_subset_imp$W1_DISC_LIFE_AVG_Z
  star_subset_imp$m2 <- star_subset_imp$W1_DISC_DAILY_AVG_Z 
  star_subset_imp$m3 <- star_subset_imp$W1_PROM_DEPR_AVG_Z 
  
  # IDs from sampling with replacement
  perm_ids <- sample(x=nrow(star_subset_imp),replace=TRUE)
  
  # extract resampled dataset for each resample 
  onedat_boot_id <- star_subset_imp[perm_ids,]
  
  # run gest function with with m_bar set to mean (~0) and resampled dataset
  list(OneEst(data=onedat_boot_id,verbose=TRUE,cde_m=0),star_subset_imp)
  },
  error=function(e) {
    NULL
  })
})
proc.time()[3]-ptm
stopCluster(cl)

# check for any conditions that failed to impute data
# 1 of 2000 had an issue
# remove from analyses
table(sapply(res_boots,is.null))
which(sapply(res_boots,is.null))
res_boots_nonull <- res_boots[!sapply(res_boots,is.null)]

# cleaning the results
res_imps <-lapply(res_boots_nonull, `[[`, 2) # imputed data
res_ests <- lapply(res_boots_nonull, `[[`, 1) # estimates stored in res_boots
res_XMint <- lapply(lapply(res_ests, `[[`, 2), `[[`, "MoM_XMint_M0obs") # assuming exposure by mediator interactions
res_noXMint <- lapply(lapply(res_ests, `[[`, 2), `[[`, "MoM_noXMint_M0obs") # assuming no exposure by mediator interactions

# median of estimates from bootstraps for point estimates
# point estimates for interactions
res_PE_XMint <- res_XMint[[1]]
for (i in 1:nrow(res_XMint[[1]])) {
  for (j in 2:ncol(res_XMint[[1]])) {
    res_PE_ij <- median(unlist(lapply(res_XMint, function(x) x[i,j])),na.rm=F)
    res_PE_XMint[i,j] <- paste(round(res_PE_ij,2),collapse=",")
  }
}

# confidence intervals for interactions
res_CI_XMint <- res_XMint[[1]]
for (i in 1:nrow(res_XMint[[1]])) {
  for (j in 2:ncol(res_XMint[[1]])) {
    res_CI_ij <- quantile(unlist(lapply(res_XMint, function(x) x[i,j])),
                          probs=c(.025,.975),na.rm=TRUE)
    res_CI_XMint[i,j] <- paste(round(res_CI_ij,2),collapse=",")
  }
}

# point estimates for no interactions
res_PE_noXMint <- res_noXMint[[1]]
for (i in 1:nrow(res_noXMint[[1]])) {
  for (j in 2:ncol(res_noXMint[[1]])) {
    res_PE_ij <- median(unlist(lapply(res_noXMint, function(x) x[i,j])))
    res_PE_noXMint[i,j] <- paste(round(res_PE_ij,2),collapse=",")
  }
}

# confidence intervals for no interactions
res_CI_noXMint <- res_noXMint[[1]]
for (i in 1:nrow(res_noXMint[[1]])) {
  for (j in 2:ncol(res_noXMint[[1]])) {
    res_CI_ij <- quantile(unlist(lapply(res_noXMint, function(x) x[i,j])),
                          probs=c(.025,.975),na.rm=TRUE)
    res_CI_noXMint[i,j] <- paste(round(res_CI_ij,2),collapse=",")
  }
}

# printing results
res_PE_XMint
res_CI_XMint
res_PE_noXMint
res_CI_noXMint

# saving results
save.image(file = "C:\\Users\\dcintron\\Desktop\\code-star-v2\\code-star\\results\\cde-xmint-latent-overall-imp-noint-120721.RData")
#load(file = "C:\\Users\\dcintron\\Desktop\\code-star-v2\\code-star\\results\\cde-xmint-latent-overall-imp-noint-120721.RData")

#################
# Total Effects #
#################

set.seed(1231)
# creating a formula list. Element 1) Y ~ X. Element 2) Y ~ X + baseline covariates
te_formula <- list()
te_formula[[1]] <- W1_D_SENAS_AVG_COGNITION_Z ~ factor(HS_BH) + factor(HS_BH_GR) + 
  z_W1_INTERVIEW_AGE + factor(z_W1_D_GENDER) + factor(z_GRF_C1) + factor(z_GRF_C3) + 
  factor(z_MED_C1) + factor(z_MED_C2) + factor(z_MED_C3) +factor(z_PED_C1) + 
  factor(z_PED_C2) + factor(z_PED_C3) + factor(z_W1_MATERNAL_EMPLOYMENT) + 
  factor(z_W1_PATERNAL_EMPLOYMENT) 

# fit formulas over all imputed datasets
lm_te <- list()
for(i in 1:length(res_imps)){
  lm_te[[i]] <- lapply(te_formula,lm,data=res_imps[[i]])
}
lm_te_coefs <- unlist(lm_te, recursive = F)
coefs_bach <- unlist(lapply(lm_te_coefs, function(x) coef(x)[2])) # HS_BH
coefs_grad <- unlist(lapply(lm_te_coefs, function(x) coef(x)[3])) # HS_BH_GR
coefs_age <- unlist(lapply(lm_te_coefs, function(x) coef(x)[4])) 
coefs_gender <- unlist(lapply(lm_te_coefs, function(x) coef(x)[5])) 
coefs_gufp1 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[6])) 
coefs_gufp3 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[7])) 
coefs_med1 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[8])) 
coefs_med2 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[9])) 
coefs_med3 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[10])) 
coefs_ped1 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[11])) 
coefs_ped2 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[12])) 
coefs_ped3 <- unlist(lapply(lm_te_coefs, function(x) coef(x)[13])) 
coefs_me <- unlist(lapply(lm_te_coefs, function(x) coef(x)[14])) 
coefs_pe <- unlist(lapply(lm_te_coefs, function(x) coef(x)[15])) 

# point estimates and CI for bootstrapped imputed data
(coefs_bach_PE <- median(coefs_bach))
(coefs_bach_CI <- quantile(coefs_bach,probs=c(.025,.975),na.rm=TRUE))
(coefs_grad_PE <- median(coefs_grad))
(coefs_grad_CI <- quantile(coefs_grad,probs=c(.025,.975),na.rm=TRUE))

median(coefs_age)
quantile(coefs_age,probs=c(.025,.975),na.rm=TRUE)

median(coefs_gender)
quantile(coefs_gender,probs=c(.025,.975),na.rm=TRUE)

median(coefs_gufp1)
quantile(coefs_gufp1,probs=c(.025,.975),na.rm=TRUE)

median(coefs_gufp3)
quantile(coefs_gufp3,probs=c(.025,.975),na.rm=TRUE)

median(coefs_med1)
quantile(coefs_med1,probs=c(.025,.975),na.rm=TRUE)

median(coefs_med2)
quantile(coefs_med2,probs=c(.025,.975),na.rm=TRUE)

median(coefs_med3)
quantile(coefs_med3,probs=c(.025,.975),na.rm=TRUE)

median(coefs_ped1)
quantile(coefs_ped1,probs=c(.025,.975),na.rm=TRUE)

median(coefs_ped2)
quantile(coefs_ped2,probs=c(.025,.975),na.rm=TRUE)

median(coefs_ped3)
quantile(coefs_ped3,probs=c(.025,.975),na.rm=TRUE)

median(coefs_me)
quantile(coefs_me,probs=c(.025,.975),na.rm=TRUE)

median(coefs_pe)
quantile(coefs_pe,probs=c(.025,.975),na.rm=TRUE)

#################################################################
# Education on Mediators controlling for early-life confounders #
#################################################################

##############
# Depression #
##############

# creating a formula list. Element 1) Y ~ X. Element 2) Y ~ X + baseline covariates
med_formula <- list()
med_formula[[1]] <- 'dep  =~ o_m3_1 + o_m3_2 + o_m3_3 + o_m3_4 + o_m3_5 + o_m3_6 + o_m3_7 + o_m3_8
                     eddisc =~ o_m2_1 + o_m2_2 + o_m2_3 + o_m2_4 + o_m2_5 + o_m2_6 + o_m2_7 + o_m2_8 + o_m2_9
                     dep ~ HS_BH + HS_BH_GR + 
                     eddisc + W1_DISC_LIFE_AVG_Z + 
                     z_W1_INTERVIEW_AGE + z_W1_D_GENDER + z_GRF_C1 + z_GRF_C3 + 
                     z_MED_C1 + z_MED_C2 + z_MED_C3 + z_PED_C1 + z_PED_C2 + 
                     z_PED_C3 + z_W1_MATERNAL_EMPLOYMENT + z_W1_PATERNAL_EMPLOYMENT'


# setup cluster for parallel processing
# fit lavaan model for mediator
library(parallel)
numCores <- detectCores()
cl <- makeCluster(numCores - 2) 
clusterEvalQ(cl, {
  library(lavaan)
})
clusterExport(cl, list("res_imps","med_formula"))

set.seed(1231)
ptm <- proc.time()[3]
med_effect_dep <- parLapply(cl, 1:length(res_imps), function(i) {
  tryCatch(lavaan::cfa(data=res_imps[[i]],
                       ordered = TRUE,
                       estimator = "DWLS", 
                       std.lv = TRUE,
                       control=list(iter.max=100),
                       model = med_formula[[1]]), error=function(e) NULL)
})
proc.time()[3]-ptm
stopCluster(cl)

# function to pull parameters from lavaan output
get_est <- function(fit, coef_pos) {
  param <- parameterEstimates(fit, standardized = TRUE)[coef_pos,"est"]
  param
}
library(lavaan)
est_bach <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(18)))
est_grad <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(19)))
est_eddisc <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(20)))
est_mldisc <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(21)))
est_age <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(22)))
est_gender <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(23)))
est_gufp1 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(24)))
est_gufp3 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(25)))
est_med1 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(26)))
est_med2 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(27)))
est_med3 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(28)))
est_ped1 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(29)))
est_ped2 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(30)))
est_ped3 <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(31)))
est_me <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(32)))
est_pe <- unlist(lapply(med_effect_dep, get_est, coef_pos = c(33)))



(coefs_bach_PE <- median(est_bach))
(coefs_bach_CI <- quantile(est_bach,probs=c(.025,.975),na.rm=TRUE))
(coefs_grad_PE <- median(est_grad))
(coefs_grad_CI <- quantile(est_grad,probs=c(.025,.975),na.rm=TRUE))
(coefs_eddisc_PE <- median(est_eddisc))
(coefs_eddisc_CI <- quantile(est_eddisc,probs=c(.025,.975),na.rm=TRUE))
(coefs_mldisc_PE <- median(est_mldisc))
(coefs_mldisc_CI <- quantile(est_mldisc,probs=c(.025,.975),na.rm=TRUE))

median(est_age)
quantile(est_age,probs=c(.025,.975),na.rm=TRUE)

median(est_gender)
quantile(est_gender,probs=c(.025,.975),na.rm=TRUE)

median(est_gufp1)
quantile(est_gufp1,probs=c(.025,.975),na.rm=TRUE)

median(est_gufp3)
quantile(est_gufp3,probs=c(.025,.975),na.rm=TRUE)

median(est_med1)
quantile(est_med1,probs=c(.025,.975),na.rm=TRUE)

median(est_med2)
quantile(est_med2,probs=c(.025,.975),na.rm=TRUE)

median(est_med3)
quantile(est_med3,probs=c(.025,.975),na.rm=TRUE)

median(est_ped1)
quantile(est_ped1,probs=c(.025,.975),na.rm=TRUE)

median(est_ped2)
quantile(est_ped2,probs=c(.025,.975),na.rm=TRUE)

median(est_ped3)
quantile(est_ped3,probs=c(.025,.975),na.rm=TRUE)

median(est_me)
quantile(est_me,probs=c(.025,.975),na.rm=TRUE)

median(est_pe)
quantile(est_pe,probs=c(.025,.975),na.rm=TRUE)

###########################
# Everyday Discrimination #
###########################

# creating a formula list. Element 1) Y ~ X. Element 2) Y ~ X + baseline covariates
med_formula <- list()
med_formula[[1]] <- 'eddisc =~ o_m2_1 + o_m2_2 + o_m2_3 + o_m2_4 + o_m2_5 + o_m2_6 + o_m2_7 + o_m2_8 + o_m2_9
                     eddisc ~ W1_DISC_LIFE_AVG_Z +  + HS_BH + HS_BH_GR + 
                     z_W1_INTERVIEW_AGE + z_W1_D_GENDER + z_GRF_C1 + z_GRF_C3 + 
                     z_MED_C1 + z_MED_C2 + z_MED_C3 + z_PED_C1 + z_PED_C2 + 
                     z_PED_C3 + z_W1_MATERNAL_EMPLOYMENT + z_W1_PATERNAL_EMPLOYMENT'


# setup cluster for parallel processing
# fit lavaan model for mediator
library(parallel)
numCores <- detectCores()
cl <- makeCluster(numCores - 2) 
clusterEvalQ(cl, {
  library(lavaan)
})
clusterExport(cl, list("res_imps","med_formula"))

set.seed(1231)
ptm <- proc.time()[3]
med_effect_ed <- parLapply(cl, 1:length(res_imps), function(i) {
  tryCatch(lavaan::cfa(data=res_imps[[i]],
                                     ordered = TRUE,
                                     estimator = "DWLS", 
                                     std.lv = TRUE,
                                     control=list(iter.max=100),
                                     model = med_formula[[1]]), error=function(e) NULL)
})
proc.time()[3]-ptm
stopCluster(cl)
table(unlist(lapply(res_boots,is.null)))

# function to pull parameters from lavaan output
library(lavaan)
get_est <- function(fit, coef_pos) {
  param <- parameterEstimates(fit, standardized = TRUE)[coef_pos,"est"]
  param
}

est_bach <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(11)))
est_grad <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(12)))
est_mldisc <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(10)))

est_age <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(13)))
est_gender <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(14)))
est_gufp1 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(15)))
est_gufp3 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(16)))
est_med1 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(17)))
est_med2 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(18)))
est_med3 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(19)))
est_ped1 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(20)))
est_ped2 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(21)))
est_ped3 <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(22)))
est_me <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(23)))
est_pe <- unlist(lapply(med_effect_ed, get_est, coef_pos = c(24)))

(coefs_bach_PE <- median(est_bach))
(coefs_bach_CI <- quantile(est_bach,probs=c(.025,.975),na.rm=TRUE))
(coefs_grad_PE <- median(est_grad))
(coefs_grad_CI <- quantile(est_grad,probs=c(.025,.975),na.rm=TRUE))
(coefs_mldisc_PE <- median(est_mldisc))
(coefs_mldisc_CI <- quantile(est_mldisc,probs=c(.025,.975),na.rm=TRUE))

median(est_age)
quantile(est_age,probs=c(.025,.975),na.rm=TRUE)

median(est_gender)
quantile(est_gender,probs=c(.025,.975),na.rm=TRUE)

median(est_gufp1)
quantile(est_gufp1,probs=c(.025,.975),na.rm=TRUE)

median(est_gufp3)
quantile(est_gufp3,probs=c(.025,.975),na.rm=TRUE)

median(est_med1)
quantile(est_med1,probs=c(.025,.975),na.rm=TRUE)

median(est_med2)
quantile(est_med2,probs=c(.025,.975),na.rm=TRUE)

median(est_med3)
quantile(est_med3,probs=c(.025,.975),na.rm=TRUE)

median(est_ped1)
quantile(est_ped1,probs=c(.025,.975),na.rm=TRUE)

median(est_ped2)
quantile(est_ped2,probs=c(.025,.975),na.rm=TRUE)

median(est_ped3)
quantile(est_ped3,probs=c(.025,.975),na.rm=TRUE)

median(est_me)
quantile(est_me,probs=c(.025,.975),na.rm=TRUE)

median(est_pe)
quantile(est_pe,probs=c(.025,.975),na.rm=TRUE)

###########################
# Major Discrimination #
###########################

# creating a formula list. Element 1) Y ~ X. Element 2) Y ~ X + baseline covariates
med_formula <- list()
med_formula[[1]] <- 'W1_DISC_LIFE_AVG_Z ~ HS_BH + HS_BH_GR  + 
                     z_W1_INTERVIEW_AGE + z_W1_D_GENDER + z_GRF_C1 + z_GRF_C3 + 
                     z_MED_C1 + z_MED_C2 + z_MED_C3 + z_PED_C1 + z_PED_C2 + 
                     z_PED_C3 + z_W1_MATERNAL_EMPLOYMENT + z_W1_PATERNAL_EMPLOYMENT'


# setup cluster for parallel processing
# fit lavaan model for mediator
library(parallel)
numCores <- detectCores()
cl <- makeCluster(numCores - 2) 
clusterEvalQ(cl, {
  library(lavaan)
})
clusterExport(cl, list("res_imps","med_formula"))

set.seed(1231)
ptm <- proc.time()[3]
med_effect_mld <- parLapply(cl, 1:length(res_imps), function(i) {
  tryCatch(lavaan::cfa(data=res_imps[[i]],
                       ordered = TRUE,
                       estimator = "DWLS", 
                       std.lv = TRUE,
                       control=list(iter.max=100),
                       model = med_formula[[1]]), error=function(e) NULL)
})
proc.time()[3]-ptm
stopCluster(cl)
table(unlist(lapply(res_boots,is.null)))

# function to pull parameters from lavaan output
library(lavaan)
get_est <- function(fit, coef_pos) {
  param <- parameterEstimates(fit, standardized = TRUE)[coef_pos,"est"]
  param
}

est_bach <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(1)))
est_grad <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(2)))

est_age <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(3)))
est_gender <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(4)))
est_gufp1 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(5)))
est_gufp3 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(6)))
est_med1 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(7)))
est_med2 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(8)))
est_med3 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(9)))
est_ped1 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(10)))
est_ped2 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(11)))
est_ped3 <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(12)))
est_me <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(13)))
est_pe <- unlist(lapply(med_effect_mld, get_est, coef_pos = c(14)))

(coefs_bach_PE <- median(est_bach))
(coefs_bach_CI <- quantile(est_bach,probs=c(.025,.975),na.rm=TRUE))
(coefs_grad_PE <- median(est_grad))
(coefs_grad_CI <- quantile(est_grad,probs=c(.025,.975),na.rm=TRUE))

median(est_age)
quantile(est_age,probs=c(.025,.975),na.rm=TRUE)

median(est_gender)
quantile(est_gender,probs=c(.025,.975),na.rm=TRUE)

median(est_gufp1)
quantile(est_gufp1,probs=c(.025,.975),na.rm=TRUE)

median(est_gufp3)
quantile(est_gufp3,probs=c(.025,.975),na.rm=TRUE)

median(est_med1)
quantile(est_med1,probs=c(.025,.975),na.rm=TRUE)

median(est_med2)
quantile(est_med2,probs=c(.025,.975),na.rm=TRUE)

median(est_med3)
quantile(est_med3,probs=c(.025,.975),na.rm=TRUE)

median(est_ped1)
quantile(est_ped1,probs=c(.025,.975),na.rm=TRUE)

median(est_ped2)
quantile(est_ped2,probs=c(.025,.975),na.rm=TRUE)

median(est_ped3)
quantile(est_ped3,probs=c(.025,.975),na.rm=TRUE)

median(est_me)
quantile(est_me,probs=c(.025,.975),na.rm=TRUE)

median(est_pe)
quantile(est_pe,probs=c(.025,.975),na.rm=TRUE)

###################################
# Correlation Matrix 
###################################

for(i in 1:length(res_imps)) {
  res_imps[[i]]$ED_LEVEL <- 1
  res_imps[[i]][res_imps[[i]]$HS_BH==0 & res_imps[[i]]$HS_BH_GR==0, "ED_LEVEL"] <- 0
  res_imps[[i]][res_imps[[i]]$HS_BH==1, "ED_LEVEL"] <- 1
  res_imps[[i]][res_imps[[i]]$HS_BH_GR==1, "ED_LEVEL"] <- 2
}

cors <- list()
cors <- lapply(seq_along(res_imps), function(i) cor(res_imps[[i]][c(1:4,7,5:6,61)]))
array_cors <- array(do.call(cbind, cors), dim=c(dim(cors[[1]]), length(cors)))
apply(array_cors, c(1, 2), mean, na.rm = TRUE)

# creating a formula list. Element 1) Y ~ X. Element 2) Y ~ X + baseline covariates
med_formula <- list()
med_formula[[1]] <- 'dep  =~ o_m3_1 + o_m3_2 + o_m3_3 + o_m3_4 + o_m3_5 + o_m3_6 + o_m3_7 + o_m3_8
                     eddisc =~ o_m2_1 + o_m2_2 + o_m2_3 + o_m2_4 + o_m2_5 + o_m2_6 + o_m2_7 + o_m2_8 + o_m2_9
                     HS_BH ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     HS_BH_GR ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     HS_BH_FEMALE ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     HS_BH_GR_FEMALE ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_W1_INTERVIEW_AGE ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_W1_D_GENDER ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_GRF_C1 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_GRF_C3 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_MED_C1 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_MED_C2 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_MED_C3 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_PED_C1 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_PED_C2 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_PED_C3 ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_W1_MATERNAL_EMPLOYMENT ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                     z_W1_PATERNAL_EMPLOYMENT ~ dep + eddisc + W1_DISC_LIFE_AVG_Z
                    '