library(metafor)
library(dplyr)
library(ggplot2)
library(tidyr)
library(lavaan)

data_FFM <- readxl::read_excel("../../DataSet/FFM_DATA.xlsx", sheet = "Coding_Sheet")
data_HEXACO <- readxl::read_excel("../../DataSet/HEXACO_DATA.xlsx", sheet = "HEXACO Coding")
####### Main calculation functions #######

### Main estimation function for the meta-analytic effect.
### You can specify the name of: rxx, ryy, rxy, N.

### Hunter-Schmidt estimator was used, with inverse sampling variance as weighting factor.
### Fixed effect estimates are also provided - but not used in our situation.

### If rxx/ryy is missing, we use mean imputation of the input data frame.
### If you provide a data frame where no study has rxx or ryy - impute 1 as conservative means.
calculation <- function(input, valueX, valueY, method = "HS",
                        var_names = list(rxx = "rxx",
                                         ryy = "ryy",
                                         rxy = "rxy",
                                         N = "N")) {
  digits <- 2
  if (is.null(input)) return(NULL)
  if (nrow(input) < 1) return(NULL)
  ri <- as.numeric(unlist(input[var_names$rxy]))
  ni <- as.numeric(unlist(input[var_names$N]))
  rxx <- as.numeric(unlist(input[var_names$rxx]))
  ryy <- as.numeric(unlist(input[var_names$ryy]))
  
  ## Mean Imputation
  rxx = rxx %>% tidyr::replace_na(mean(na.omit(rxx)))
  ryy = ryy %>% tidyr::replace_na(mean(na.omit(ryy)))
  
  ## If mean(na.omit()) fails, result will become all NaN. In this case we replace NaN with 1
  if (is.na(rxx[1])) {
    # print("Artificial Alpha X was imputed")
    if (missing(valueX)) valueX <- 1
    # print(valueX)
    rxx <- valueX
  }
  if (is.na(ryy[1])) {
    # print("Artificial Alpha Y was imputed")
    if (missing(valueY)) valueY <- 1
    # print(valueY)
    ryy <- valueY
  }
  
  data <- cbind(ri, ni, rxx, ryy)
  k = length(ri)
  N = sum(ni)
  
  ri <- ifelse(ri==0,0.0000001, ri)
  
  M1 <- sum(ni*ri)/sum(ni)  # weighted mean, uncorrected (bare bones)
  # print(sprintf("Bare bones correlation:%.3f", M1))
  
  var_M1 <- sum(ni*(ri-M1)^2)/sum(ni)
  
  var.i <-((1-M1^2)^2)/(ni-1) # individual study sampling variance (bare bones)
  ai <- sqrt(rxx*ryy) # individual study attenuation for unreliability
  var.i2 <- var.i/ai^2 # corrected individual study sampling variance for unreliability
  ri.m <- ri/ai # corrected individual study ES for unreliability
  ri.m_weighted <- sum(ri.m*ni)/sum(ni)
  
  
  corrx <- rma(yi = ri.m,vi = var.i2,method = method, weights = 1/(var.i2))
  # corrx <- rma(yi = ri.m,vi = var.i2, method = "HS", weights = ni - 3)
  rho <- corrx$b
  
  sd_r <- sqrt(var_M1)
  se_rho <- corrx$se
  sd_rho <- sqrt(corrx$tau2)
  
  avg_var.i2 <- sum(ni*var.i2)/sum(ni)
  var_corrected_r <- sum(ni * (ri.m-ri.m_weighted)^2)/sum(ni)
  var_rho <- var_corrected_r - avg_var.i2
  sd_rho2 <- ifelse(var_rho < 0, 0, sqrt(var_rho))
  
  
  var_rho1 <- sum(ni*(ri.m - c(rho))^2)/sum(ni)
  
  CI_u <- rho + 1.96 * se_rho
  CI_l <- rho - 1.96 * se_rho
  CI_95 <- paste("[",roundif(CI_l, digits),",", 
                 roundif(CI_u, digits),"]",sep = "", collapse = "")
  
  CV_u <- rho + 1.28 * sd_rho
  CV_l <- rho - 1.28 * sd_rho
  CV_80 <- paste("[", roundif(CV_l, digits),",", 
                 roundif(CV_u, digits),"]",sep = "", collapse = "")
  
  CV_u2 <- rho + 1.28 * sd_rho2
  CV_l2 <- rho - 1.28 * sd_rho2
  CV_80_2 <- paste("[", roundif(CV_l2, digits),",", 
                 roundif(CV_u2, digits),"]",sep = "", collapse = "")
  
  
  ### Results for fixed effects
  corrx_fe <- rma(yi = ri.m,vi = var.i2,method = "FE", weights = 1/(var.i2))
  rho_fe <- corrx_fe$b
  
  sd_r_fe <- sqrt(var_M1)
  se_rho_fe <- corrx_fe$se
  sd_rho_fe = sqrt(corrx_fe$tau2)
  
  var_rho1_fe <- sum(ni*(ri.m - c(rho_fe))^2)/sum(ni)
  
  CI_u_fe <- rho_fe + 1.96 * se_rho_fe
  CI_l_fe <- rho_fe - 1.96 * se_rho_fe
  CI_95_fe <- paste("[",roundif(CI_l_fe, digits),",", 
                    roundif(CI_u_fe, digits),"]",sep = "", collapse = "")
  
  CV_u_fe <- rho_fe + 1.28 * sd_rho_fe
  CV_l_fe <- rho_fe - 1.28 * sd_rho_fe
  CV_80_fe <- paste("[", roundif(CV_l_fe, digits),",", 
                    roundif(CV_u_fe, digits),"]",sep = "", collapse = "")  
  
  
  return(list(result = corrx, 
              result_fe = corrx_fe, 
              details = data.frame(k = k, N = N, r = round(M1, 3), 
                                   sd_r = roundif(sqrt(var_M1), digits),
                                   rho = round(rho, 3), 
                                   sd_rho = roundif(sqrt(corrx$tau2), digits), CI_95 = CI_95, CV_80 = CV_80, 
                                   se_rho = roundif(corrx$se, digits),
                                   rxx = mean(rxx), ryy = mean(ryy),
                                   var_rho = var_rho1
                                   )))
  
}

### Single, continuous moderator analysis.
### Simple effects are estimated using a fixed effect linear model.
### If you need multiple moderators, you can use get_data(), add other moderators and imitate the syntax in calculation_mod().
calculation_mod <- function(input, mods, valueX, valueY, xlab = "", ylab = "", main = "", method = "HS",
                            var_names = list(rxx = "rxx",
                                             ryy = "ryy",
                                             rxy = "rxy",
                                             N = "N"), use_Z_mod = FALSE) {
  digits <- 2
  if (is.null(input)) return(NULL)
  if (nrow(input) < 2) return(NULL)
  ri <- as.numeric(unlist(input[var_names$rxy]))
  ni <- as.numeric(unlist(input[var_names$N]))
  rxx <- as.numeric(unlist(input[var_names$rxx]))
  ryy <- as.numeric(unlist(input[var_names$ryy]))
  data <- cbind(ri, ni, rxx, ryy)
  
  ## Mean Imputation
  rxx = rxx %>% tidyr::replace_na(mean(na.omit(rxx)))
  ryy = ryy %>% tidyr::replace_na(mean(na.omit(ryy)))
  
  ## If mean(na.omit()) fails, result will become all NaN. In this case we replace NaN with 1 (Alpha = 1)
  if (is.na(rxx[1])) {
    # print("Artificial Alpha X was imputed")
    if (missing(valueX)) valueX <- 1
    # print(valueX)
    rxx <- valueX
  }
  if (is.na(ryy[1])) {
    # print("Artificial Alpha Y was imputed")
    if (missing(valueY)) valueY <- 1
    # print(valueY)
    ryy <- valueY
  }
  
  
  k = length(ri)
  N = sum(ni)
  
  ri <- ifelse(ri==0,0.0000001, ri)
  
  ### Schmidt & Hunter (2014)
  
  # Mean uncorrected correlation (p. 144)
  # "This is the sample size weight average of the observed correlations."
  M1 <- sum(ni*ri)/sum(ni)  # weighted mean, uncorrected (bare bones)  
  print(sprintf("Bare bones correlation:%.3f", M1))
  
  # Formula 3.2: Corresponding sampling variance across studies
  var_M1 <- sum(ni*(ri-M1)^2)/sum(ni)
  
  
  # This calculates the sampling error variance in the UNCORRECTED correlation, where
  # M1 is the mean uncorrected correlation across studies, and
  # ni is the sample size for each study. (p. 144)
  var.i <-((1-M1^2)^2)/(ni-1) # individual study sampling variance (bare bones)
  
  # Formula 3.25: Compound attenuation factor.
  ai <- sqrt(rxx*ryy) # individual study attenuation for unreliability
  
  # Formula 3.27: Sampling error variance in the CORRECTED correlation
  var.i2 <- var.i/(ai^2) # corrected individual study sampling variance for unreliability
  
  # Formula 3.26: From observed correlation to corrected correlation
  ri.m <- ri/ai # corrected individual study ES for unreliability
  
  ### p. 145-146: What weights should be used?
  # (p. 147) When studies differ greatly on one or more of the artifacts corrected,
  # then more complicated weighting makes better use of the information in the studies.
  
  # 
  weights = 1/var.i2
  # res_mods = rma(ri.m, var.i2, weights = ni-3, mods = mods, method = method, test = "knha")
  # 
  if (use_Z_mod == FALSE) {
    res_mods = rma(ri.m, var.i2, weights = 1/var.i2, mods = mods, method = "HS", test = "knha")
    k = length(na.omit(mods))
    c = length(unique(na.omit(mods)))
    df <- data.frame(ri.m, mods, ni, weights)
    df <- df %>% filter(!is.na(mods))
    # linear_model <- lm(ri.m ~ mods, weights = ni, data = df)
    linear_model <- lm(ri.m ~ mods, weights = weights, data = df)
  }
  else {
    z.m <- escalc(ri = ri, measure = "ZCOR", ni = ni)
    res_mods = rma(z.m, var.i2, weights = ni-3, mods = mods, method = "HS", test = "knha")
    k = length(na.omit(mods))
    c = length(unique(na.omit(mods)))
    df <- data.frame(ri.m, mods, ni, weights)
    df <- df %>% filter(!is.na(mods))
    # linear_model <- lm(ri.m ~ mods, weights = ni, data = df)
    linear_model <- lm(ri.m ~ mods, weights = weights, data = df)
  }



  
  Low <- mean(mods, na.rm = TRUE) - sd(mods, na.rm = TRUE)
  High <- mean(mods, na.rm = TRUE) + sd(mods, na.rm = TRUE)
  
  pred_low <- predict(linear_model, data.frame(mods = Low),
                      interval = "confidence", level = .95, se.fit = TRUE)
  pred_high <- predict(linear_model, data.frame(mods = High),
                       interval = "confidence", level = .95, se.fit = TRUE)
  detail_CI <- c(low_estimate = roundif(pred_low$fit[1],2), 
                 high_estimate = roundif(pred_high$fit[1],2),
                 low_CV = paste("[",roundif(pred_low$fit[2],2),",",roundif(pred_low$fit[3],2),"]",sep = ""),
                 high_CV = paste("[",roundif(pred_high$fit[2],2),",",roundif(pred_high$fit[3],2),"]",sep = ""))
  
  ## plot = plot(mods, ri.m, pch = 20, xlab = xlab, ylab = ylab, main = main, family = "serif") + abline(res_mods) + grid() temporarily removed
  return(list(result = res_mods, result_fixed = summary(linear_model), num_studies = k, num_countries = c, 
              details = data.frame(ci.lower = res_mods$ci.lb, ci.upper = res_mods$ci.ub),
              linear_estimate = linear_model, R2 = res_mods$R2 / 100, beta = sqrt(res_mods$R2 / 100),
              predict_high = pred_high, predict_low = pred_low, detail_CI = detail_CI))
}


### Single, categorical moderator analysis. Identical to performing subgroup analyses.
### Naming dependence removed so that we don't need to build a new function for every new moderator.
cat_mod <- function(data, mod_name) {
  if (is.null(mod_name)) {
    return(NULL)
  }
  nm <- c()
  results <- list()
  for (value in unlist(unique(data[mod_name]))) {
    if (is.na(value)) {next}
    else {
      df <- data[which(data[mod_name] == value),]
      result <- calculation(df)$details
      nm <- c(nm, value)
      results <- c(results, list(result))
    }
  }
  results <- as.list(results)
  names(results) <- nm
  # info <- extract_info(results)
  return(list(results = results))
}






####### Auxiliary functions #######

### Format the output.
### Normally we want two digit output. You can set digit = 2.
### But if the estimate happens to be <.01 and will be rounded to .00, we force it into three digits.
roundif <- function(number, digit) {
  result <- format(round(number, digit), nsmall = digit)
  if (as.numeric(result) == 0) {
    pos <- number >= 0
    if (pos) {
      return(sprintf("%s",format(round(number, digit+1), nsmall = digit+1)))
    }
    else {
      flag <- format(round(number, digit + 1), nsmall = digit + 1)
      if (as.numeric(flag) >= 0) {
        return(sprintf("-%s",format(round(number, digit+1), nsmall = digit+1)))
      }
      else {
        return(sprintf("%s",format(round(number, digit+1), nsmall = digit+1)))
      }
    }
  }
  else return(result)
}

### Identical to calculation() before calculation of estimates.
### If you want the data frame and append more continuous moderators afterwards, you can use this function.
get_data <- function(input, mods, valueX, valueY,var_names = list(rxx = "rxx",
                                                                  ryy = "ryy",
                                                                  rxy = "rxy",
                                                                  N = "N")) {
  if (is.null(input)) return(NULL)
  if (nrow(input) < 2) return(NULL)
  ri <- as.numeric(unlist(input[var_names$rxy]))
  ni <- as.numeric(unlist(input[var_names$N]))
  rxx <- as.numeric(unlist(input[var_names$rxx]))
  ryy <- as.numeric(unlist(input[var_names$ryy]))
  data <- cbind(ri, ni, rxx, ryy)
  
  ## Mean Imputation
  rxx = rxx %>% tidyr::replace_na(mean(na.omit(rxx)))
  ryy = ryy %>% tidyr::replace_na(mean(na.omit(ryy)))
  
  ## If mean(na.omit()) fails, result will become all NaN. In this case we replace NaN with 1 (Alpha = 1)
  if (is.na(rxx[1])) {
    # print("Artificial Alpha X was imputed")
    if (missing(valueX)) valueX <- 1
    # print(valueX)
    rxx <- valueX
  }
  if (is.na(ryy[1])) {
    # print("Artificial Alpha Y was imputed")
    if (missing(valueY)) valueY <- 1
    # print(valueY)
    ryy <- valueY
  }
  
  
  k = length(ri)
  N = sum(ni)
  
  ri <- ifelse(ri==0,0.0000001, ri)
  
  M1 <- sum(ni*ri)/sum(ni)  # weighted mean, uncorrected (bare bones)
  print(sprintf("Bare bones correlation:%.3f", M1))
  
  var_M1 <- sum(ni*(ri-M1)^2)/sum(ni)
  
  var.i <-((1-M1^2)^2)/(ni-1) # individual study sampling variance (bare bones)
  ai <- sqrt(rxx*ryy) # individual study attenuation for unreliability
  var.i2 <- var.i/ai^2 # corrected individual study sampling variance for unreliability
  ri.m <- ri/ai # corrected individual study ES for unreliability
  
  return(data.frame(ri.m = ri.m, var.i2 = var.i2, weights = 1/var.i2))
}

### Helper function for cat_mod to extract the effect sizes and variances for each subgroup.
### Variances are calculated using the Knapp-Hartung (2003) method.
extract_info <- function(subgroups, valueX, valueY) {
  effects <- c()
  variances <- c()
  for (g in subgroups) {
    print(g)
    cal <- calculation(g, valueX, valueY)$details
    effects <- c(effects, cal$rho)
    variances <- c(variances,
                   (cal$rxx - (cal$r)^2)*(cal$ryy - (cal$r)^2)/((cal$N-1)*((cal$rxx)^2*(cal$ryy)^2)))
  }
  return(list(effects = effects, variances = variances))
  
}

### z-test for categorical moderators.
ztest_mod <- function(rs, vars) {
  print(sprintf("Correlations: %.3f, %.3f",rs[1],rs[2]))
  z = (rs[1]-rs[2])/sqrt(vars[1]+vars[2])
  print(sprintf("Z value: %.4f",z))
  return((1 - pnorm(abs(z)))*2)
}



####### Wrapper functions - analyze all FFM/HEXACO traits at once #######
### Batch analyze main effects for FFM traits. You can also specify a categorical moderator.
### O = 1, C = 2, E = 3, A = 4, N = 5
BF_analysis <- function(input, valueX, valueY, mod = NULL, IV_name, var_names = list(rxx = "rxx",
                                                                                     ryy = "ryy",
                                                                                     rxy = "rxy",
                                                                                     N = "N")) {
  if (missing(valueX)) {
    valueX <- mean(as.numeric(unlist(input[var_names$rxx]), na.rm = TRUE))
  }
  if (missing(valueY)) {
    valueY <- mean(as.numeric(unlist(input[var_names$ryy]), na.rm = TRUE))
  }
  C_data <- input %>% filter(get(IV_name) == 2)
  A_data <- input %>% filter(get(IV_name) == 4)
  N_data <- input %>% filter(get(IV_name) == 5)
  O_data <- input %>% filter(get(IV_name) == 1)
  E_data <- input %>% filter(get(IV_name) == 3) 
  
  return(list(overall = list(C = calculation(C_data, var_names = var_names)$details,
                             A = calculation(A_data, var_names = var_names)$details,
                             N = calculation(N_data, var_names = var_names)$details,
                             O = calculation(O_data, var_names = var_names)$details,
                             E = calculation(E_data, var_names = var_names)$details),
              mod_result = list(C = cat_mod(C_data, mod),
                                A = cat_mod(A_data, mod),
                                N = cat_mod(N_data, mod),
                                O = cat_mod(O_data, mod),
                                E = cat_mod(E_data, mod))))
}

HEX_analysis <- function(input, valueX, valueY, mod = NULL, IV_name, var_names = list(rxx = "rxx",
                                                                                      ryy = "ryy",
                                                                                      rxy = "rxy",
                                                                                      N = "N")) {
  if (missing(valueX)) {
    valueX <- mean(as.numeric(unlist(input[var_names$rxx]), na.rm = TRUE))
  }
  if (missing(valueY)) {
    valueY <- mean(as.numeric(unlist(input[var_names$ryy]), na.rm = TRUE))
  }
  C_data <- input %>% filter(get(IV_name) == 2)
  A_data <- input %>% filter(get(IV_name) == 4)
  E_data <- input %>% filter(get(IV_name) == 5)
  O_data <- input %>% filter(get(IV_name) == 1)
  X_data <- input %>% filter(get(IV_name) == 3) 
  H_data <- input %>% filter(get(IV_name) == 6)
  
  return(list(overall = list(H = calculation(H_data, var_names = var_names)$details,
                             E = calculation(E_data, var_names = var_names)$details,
                             X = calculation(X_data, var_names = var_names)$details,
                             A = calculation(A_data, var_names = var_names)$details,
                             C = calculation(C_data, var_names = var_names)$details,
                             O = calculation(O_data, var_names = var_names)$details),
              mod_result = list(H = cat_mod(H_data, mod),
                                E = cat_mod(E_data, mod),
                                X = cat_mod(X_data, mod),
                                A = cat_mod(A_data, mod),
                                C = cat_mod(C_data, mod),
                                O = cat_mod(O_data, mod))))
}

### Batch HLM (3-level meta-analysis) for FFM traits. You need to specify index and cluster variables.
BF_analysis_HLM <- function(input, valueX, valueY, cluster_var, ind_var, IV_name, mods = NULL, var_names = list(rxx = "rxx",
                                                                                                                ryy = "ryy",
                                                                                                                rxy = "rxy",
                                                                                                                N = "N")) {
  if (missing(valueX)) {
    valueX <- mean(as.numeric(unlist(input[var_names$rxx]), na.rm = TRUE))
  }
  if (missing(valueY)) {
    valueY <- mean(as.numeric(unlist(input[var_names$ryy]), na.rm = TRUE))
  }
  C_data <- input %>% filter(get(IV_name) == 2)
  A_data <- input %>% filter(get(IV_name) == 4)
  N_data <- input %>% filter(get(IV_name) == 5)
  O_data <- input %>% filter(get(IV_name) == 1)
  E_data <- input %>% filter(get(IV_name) == 3) 

  
  if (!is.null(mods)) {
    return(list(C = calculation_HLM(C_data, cluster_var = C_data[cluster_var], ind_var = C_data[ind_var], mods = C_data[mods], var_names = var_names),
                A = calculation_HLM(A_data, cluster_var = A_data[cluster_var], ind_var = A_data[ind_var], mods = A_data[mods], var_names = var_names),
                N = calculation_HLM(N_data, cluster_var = N_data[cluster_var], ind_var = N_data[ind_var], mods = N_data[mods], var_names = var_names),
                O = calculation_HLM(O_data, cluster_var = O_data[cluster_var], ind_var = O_data[ind_var], mods = O_data[mods], var_names = var_names),
                E = calculation_HLM(E_data, cluster_var = E_data[cluster_var], ind_var = E_data[ind_var], mods = E_data[mods], var_names = var_names))) 
  }
  else {
    return(list(C = calculation_HLM(C_data, cluster_var = C_data[cluster_var], ind_var = C_data[ind_var], var_names = var_names),
                A = calculation_HLM(A_data, cluster_var = A_data[cluster_var], ind_var = A_data[ind_var], var_names = var_names),
                N = calculation_HLM(N_data, cluster_var = N_data[cluster_var], ind_var = N_data[ind_var], var_names = var_names),
                O = calculation_HLM(O_data, cluster_var = O_data[cluster_var], ind_var = O_data[ind_var], var_names = var_names),
                E = calculation_HLM(E_data, cluster_var = E_data[cluster_var], ind_var = E_data[ind_var], var_names = var_names)))     
  } 
}

### Batch HLM (3-level meta-analysis) for FFM traits. You need to specify index and cluster variables.
HEX_analysis_HLM <- function(input, valueX, valueY, cluster_var, ind_var, IV_name, mods = NULL, var_names = list(rxx = "rxx",
                                                                                                                 ryy = "ryy",
                                                                                                                 rxy = "rxy",
                                                                                                                 N = "N")) {
  if (missing(valueX)) {
    valueX <- mean(as.numeric(unlist(input[var_names$rxx]), na.rm = TRUE))
  }
  if (missing(valueY)) {
    valueY <- mean(as.numeric(unlist(input[var_names$ryy]), na.rm = TRUE))
  }
  C_data <- input %>% filter(get(IV_name) == 2)
  A_data <- input %>% filter(get(IV_name) == 4)
  X_data <- input %>% filter(get(IV_name) == 3)
  O_data <- input %>% filter(get(IV_name) == 1)
  E_data <- input %>% filter(get(IV_name) == 5) 
  H_data <- input %>% filter(get(IV_name) == 6) 
  
  if (!is.null(mods)) {
    return(list(C = calculation_HLM(C_data, cluster_var = C_data[cluster_var], ind_var = C_data[ind_var], mods = C_data[mods], var_names = var_names),
                A = calculation_HLM(A_data, cluster_var = A_data[cluster_var], ind_var = A_data[ind_var], mods = A_data[mods], var_names = var_names),
                X = calculation_HLM(X_data, cluster_var = X_data[cluster_var], ind_var = X_data[ind_var], mods = X_data[mods], var_names = var_names),
                O = calculation_HLM(O_data, cluster_var = O_data[cluster_var], ind_var = O_data[ind_var], mods = O_data[mods], var_names = var_names),
                E = calculation_HLM(E_data, cluster_var = E_data[cluster_var], ind_var = E_data[ind_var], mods = E_data[mods], var_names = var_names),
                H = calculation_HLM(H_data, cluster_var = H_data[cluster_var], ind_var = H_data[ind_var], mods = H_data[mods], var_names = var_names))) 
  }
  else {
    return(list(C = calculation_HLM(C_data, cluster_var = C_data[cluster_var], ind_var = C_data[ind_var], var_names = var_names),
                A = calculation_HLM(A_data, cluster_var = A_data[cluster_var], ind_var = A_data[ind_var], var_names = var_names),
                X = calculation_HLM(X_data, cluster_var = X_data[cluster_var], ind_var = X_data[ind_var], var_names = var_names),
                O = calculation_HLM(O_data, cluster_var = O_data[cluster_var], ind_var = O_data[ind_var], var_names = var_names),
                E = calculation_HLM(E_data, cluster_var = E_data[cluster_var], ind_var = E_data[ind_var], var_names = var_names),
                H = calculation_HLM(H_data, cluster_var = H_data[cluster_var], ind_var = H_data[ind_var], var_names = var_names)))     
  } 
}


####### File writing functions - Helpful in building Tables #######

### Output main effect results to a file.
format_file <- function(result, write = FALSE, filename) {
  full_df <- data.frame()
  null_ix <- c()
  for (i in 1:length(result)) {
    if (is.null(result[[i]])) {null_ix <- c(null_ix, i)}
    df <- as.data.frame(result[i])
    names(df) <- substring(names(df), 3)
    full_df <- rbind(full_df, df)
  }
  if (length(null_ix) == 0) {
    rownames(full_df) <- names(result)
  }
  else {
    rownames(full_df) <- names(result)[-null_ix]
  }
  if (write) {
    write.csv(full_df, filename)
    return(full_df)
  }
  else {
    return(full_df)
  }
}

### Output categorical moderator results to a file.
### This one presents results for each personality factor.
format_file_mods <- function(mod_result, write = FALSE, arrange = TRUE, order, filename) {
  full_df <- data.frame()
  null_ix <- c()
  for (i in 1:length(mod_result)) {      # i represents each FFM/HEXACO dimension
    if (is.null(mod_result[i])) {null_ix <- c(null_ix, i)}
    BF_df <- data.frame()
    na_ix <- c()
    for (j in 1:length(mod_result[[i]]$results)) {       # j represents each value in the categorical moderator
      df <- as.data.frame(mod_result[[i]]$results[[j]])
      # colnames(df) <- c("k", "N", "r", "sd_r", "rho",
      #              "sd_rho","CI_95","CV_80","se_rho","rxx","ryy","var_rho","agreement")
      # rownames(df) <- c(names(mod_result)[i])
      if (nrow(df) == 0) {
        na_ix <- c(na_ix, j)
      }
      BF_df <- rbind(BF_df, df)
    }
    if (length(na_ix) > 0) {
      rownames(BF_df) <- paste(names(mod_result)[i], "_", names(mod_result[[i]]$results)[-na_ix], sep = "")
    }
    else {
      rownames(BF_df) <- paste(names(mod_result)[i], "_", names(mod_result[[i]]$results), sep = "")
    }
    
    full_df <- rbind(full_df, BF_df)
  }
  if (arrange) {
    full_df$names <- rownames(full_df)
    full_df <- full_df %>% separate(names, c("Factor","Category")) %>% 
      arrange(Category, match(Factor, order)) %>% select(-c(Factor, Category))
    
  }
  # rownames(full_df) <- rep(names(mod_result), each = length(mod_result[[1]]$results))
  if (write) {
    write.csv(full_df, filename)
    return(full_df)
  }
  else {
    return(full_df)
  }
}


