# Load necessary libraries
library(lavaan)
library(car)
library(lmtest)
library(sandwich)
library(openxlsx)
library(readxl)

# Read DF
df2 <- read_excel("/Users/carolineseton/Desktop/df2.xlsx")

# Mediation Analyses Unadjusted Models
# Define variables
predictors <- c("CFRS", "PCA_RC2")  
mediators <- c("BERQ_ActivelyApproach", "BERQ_Withdrawal", 
               "CERQ_SelfBlame", "CERQ_Planning", "CERQ_Reappraisal", "CERQ_Catastrophising")
outcomes <- c("GAD", "GDS", "WEMWBS")  

# Function to check assumptions
check_assumptions <- function(model) {
  results <- list()
  results$VIF <- vif(model)
  results$Linearity <- resettest(model)$p.value
  results$Homoscedasticity <- bptest(model)$p.value
  residuals <- residuals(model)
  results$Normality <- shapiro.test(residuals)$p.value
  return(results)
}

# Function to compute Aroian test
aroian_test <- function(a, sa, b, sb) {
  z_value <- (a * b) / sqrt((b^2 * sa^2) + (a^2 * sb^2) + (sa^2 * sb^2))
  p_value <- 2 * (1 - pnorm(abs(z_value)))
  return(c(z_value, p_value))
}

# Create storage lists
mediation_results <- list()
assumption_tests <- list()
aroian_results <- list()

# Loop through predictors, mediators, and outcomes
for (predictor in predictors) {
  for (mediator in mediators) {
    for (outcome in outcomes) {
      cat("\nRunning mediation for:", predictor, "->", mediator, "->", outcome, "\n")
      
      # Define mediation model
      mediation_model <- paste0(
        mediator, " ~ a*", predictor, "\n",
        outcome, " ~ b*", mediator, " + c*", predictor, "\n",
        "indirect := a * b", "\n",
        "total_indirect := indirect"
      )
      
      # Fit mediation model with bootstrapping
      fit <- sem(mediation_model, data = df2, estimator = "ML", missing = "fiml", se = "bootstrap", bootstrap = 1000)
      
      # Extract estimates
      fit_summary <- parameterEstimates(fit, standardized = TRUE, ci = TRUE, level = 0.95)
      
      # Extract relevant paths
      a_path <- fit_summary[fit_summary$lhs == mediator & fit_summary$rhs == predictor & fit_summary$op == "~",]
      b_path <- fit_summary[fit_summary$lhs == outcome & fit_summary$rhs == mediator & fit_summary$op == "~",]
      c_prime_path <- fit_summary[fit_summary$lhs == outcome & fit_summary$rhs == predictor & fit_summary$op == "~",]
      indirect_path <- fit_summary[grepl("indirect", fit_summary$label),]
      
      # Compute Aroian Test
      aroian_values <- aroian_test(a_path$est, a_path$se, b_path$est, b_path$se)
      
      # Store mediation results
      mediation_results[[paste(predictor, mediator, outcome, sep = "_")]] <- data.frame(
        Predictor = predictor,
        Mediator = mediator,
        Outcome = outcome,
        
        Path_A_Beta = a_path$est,
        Path_A_SE = a_path$se,
        Path_A_p = a_path$pvalue,
        Path_A_CI_Lower = a_path$ci.lower,
        Path_A_CI_Upper = a_path$ci.upper,
        
        Path_B_Beta = b_path$est,
        Path_B_SE = b_path$se,
        Path_B_p = b_path$pvalue,
        Path_B_CI_Lower = b_path$ci.lower,
        Path_B_CI_Upper = b_path$ci.upper,
        
        Direct_Beta = c_prime_path$est,
        Direct_SE = c_prime_path$se,
        Direct_p = c_prime_path$pvalue,
        Direct_CI_Lower = c_prime_path$ci.lower,
        Direct_CI_Upper = c_prime_path$ci.upper,
        
        Indirect_Beta = indirect_path$est,
        Indirect_SE = indirect_path$se,
        Indirect_p = indirect_path$pvalue,
        Indirect_CI_Lower = indirect_path$ci.lower,
        Indirect_CI_Upper = indirect_path$ci.upper,
        
        Aroian_Z = aroian_values[1],
        Aroian_p = aroian_values[2],
        
        Significance = ifelse(indirect_path$ci.lower > 0 | indirect_path$ci.upper < 0, "Significant", "Not Significant")
      )
      
      # Store Aroian results
      aroian_results[[paste(predictor, mediator, outcome, sep = "_")]] <- data.frame(
        Predictor = predictor,
        Mediator = mediator,
        Outcome = outcome,
        Aroian_Z = aroian_values[1],
        Aroian_p = aroian_values[2]
      )
      
      # Check assumptions
      model_formula <- paste(outcome, "~", predictor, "+", mediator)
      model_lm <- lm(as.formula(model_formula), data = df2)
      assumption_tests[[paste(predictor, mediator, outcome, sep = "_")]] <- check_assumptions(model_lm)
    }
  }
}

# Convert lists to data frames
mediation_results_df <- do.call(rbind, mediation_results)
aroian_results_df <- do.call(rbind, aroian_results)

# Convert assumption tests into a data frame format
assumption_results_df <- data.frame(
  Mediation = character(),
  VIF = character(),
  Linearity = numeric(),
  Homoscedasticity = numeric(),
  Normality = numeric(),
  stringsAsFactors = FALSE
)

for (test_name in names(assumption_tests)) {
  test_results <- assumption_tests[[test_name]]
  
  assumption_results_df <- rbind(assumption_results_df, data.frame(
    Mediation = test_name,
    VIF = paste(names(test_results$VIF), round(test_results$VIF, 3), collapse = ", "),
    Linearity = round(test_results$Linearity, 3),
    Homoscedasticity = round(test_results$Homoscedasticity, 3),
    Normality = round(test_results$Normality, 3)
  ))
}

# Save results to Excel files
write.xlsx(mediation_results_df, "/Users/carolineseton/Desktop/Study 2/mediation_results.xlsx")
write.xlsx(aroian_results_df, "/Users/carolineseton/Desktop/Study 2/aroian_test_results.xlsx")
write.xlsx(assumption_results_df, "/Users/carolineseton/Desktop/Study 2/assumption_test_results.xlsx")