library(metafor)
library(dplyr)
library(ggplot2)
library(lavaan)
library(psych)
# source("Utilities-Main.R")

### Perform 3-level meta-analysis.
### This can be used to ensure level-3 variance is not zero.

calculation_HLM <- function(input, valueX, valueY, cluster_var, ind_var, mods = NULL,
                            var_names = list(rxx = "rxx",
                                             ryy = "ryy",
                                             rxy = "rxy",
                                             N = "N")) {
  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]))
  
  
  ## 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
  
  # TODO Back to this later to see how tau may be calculated (Based on H&S page 151-152)
  
  # var_rc <- sum(ni*(ri.m-rho)^2)/sum(ni)
  
  # End of TODO
  names(cluster_var) <- "cluster"
  names(ind_var) <- "ind"
  if (!is.null(mods)) {
    df <- data.frame(ri.m = ri.m, var.i2 = var.i2, cluster = c(cluster_var), ind = c(ind_var), mods = mods, N = ni)   
  }
  else {
    df <- data.frame(ri.m = ri.m, var.i2 = var.i2, cluster = c(cluster_var), ind = c(ind_var), N = ni)
  }
  
  
  df <- df %>% filter(!is.na(cluster))

  if (is.null(mods)) {
    corrx <- rma.mv(yi = ri.m,V = var.i2, slab = cluster, method = "ML",
                    random = ~ 1 | cluster/ind, test = "t", data = df)
    
    corrx_novar2 <- rma.mv(yi = ri.m,V = var.i2, slab = cluster, method = "ML",
                           random = list(~ 1 | ind, ~ 1 | cluster), sigma2 = c(0,NA), test = "t", data = df)
    corrx_novar3 <- rma.mv(yi = ri.m,V = var.i2, slab = cluster, method = "ML",
                           random = list(~ 1 | ind, ~ 1 | cluster), sigma2 = c(NA,0), test = "t", data = df)    
  }
  ### If moderator is specified in HLM...but still, single moderator in this case.
  ### If multiple moderators are needed you can use getdata() and copy the syntax below and run analyses.
  else {
    ## We used test = "t" because it mimics the Knapp & Hartung (2003) method
    ## K-N method was chosen because Aguinis & Gonzales-Mule (2018) recommended it for ME models
    ## when testing boundary condition hypotheses.
    
    ## Use an iterative estimator of residual heterogeneity, like ML/REML.
    # corrx <- rma.mv(yi = ri.m,V = var.i2, slab = cluster, method = "REML", W = 1 / var.i2,
    #                 random = ~ 1 | cluster/ind, mods = ~ mods, test = "t", struct = "UN", data = df)
    corrx <- rma.mv(yi = ri.m,V = var.i2, slab = cluster, method = "REML", W = 1 / var.i2,
                    random = ~ 1 | cluster/ind, mods = mods, test = "t", data = df)   
    corrx_novar2 <- rma.mv(yi = ri.m,V = var.i2, slab = cluster, method = "REML", W = 1 / var.i2,
                           random = list(~ 1 | ind, ~ 1 | cluster), mods = mods, sigma2 = c(0,NA), test = "t", data = df)
    corrx_novar3 <- rma.mv(yi = ri.m,V = var.i2, slab = cluster, method = "REML", W = 1 / var.i2,
                           random = list(~ 1 | ind, ~ 1 | cluster), mods = mods, sigma2 = c(NA,0), test = "t", data = df)       
  }
  robust_se <- robust(corrx, cluster = df$cluster)
  # https://www.metafor-project.org/doku.php/tips:i2_multilevel_multivariate 
  W <- diag(1/df$var.i2)
  X <- model.matrix(corrx)
  P <- W - W %*% X %*% solve(t(X) %*% W %*% X) %*% t(X) %*% W
  var_dist <- round(100 * corrx$sigma2 / (sum(corrx$sigma2) + (corrx$k - corrx$p)/sum(diag(P))), 3)
  ## ICC calculation is referenced from:
  ## https://www.metafor-project.org/doku.php/analyses:konstantopoulos2011
  return(list(result = corrx, ICC = round(corrx$sigma2[1] / sum(corrx$sigma2), 3),
              var_dist = list(Between = var_dist[1], Within = var_dist[2]),
              anova_level2 = anova(corrx, corrx_novar2),
              anova_level3 = anova(corrx, corrx_novar3),
              robust_se = robust_se))
}

### Curvilinear test of a continuous moderator.
### Plot also included.
curvicalc_mod <- function(input, mods, valueX, valueY,
                          var_names = list(rxx = "rxx",
                                           ryy = "ryy",
                                           rxy = "rxy",
                                           N = "N")) {
  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)
  
  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

  res_mods = rma(ri.m, var.i2, weights = 1/var.i2, mods = ~mods + I(mods^2), method = "HS", test = "knha")
  k = length(na.omit(mods))
  c = length(unique(na.omit(mods)))
  
  df <- data.frame(ri.m, mods, ni)
  linear_model <- lm(ri.m ~ mods + I(mods^2), weights = ni, 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)
  
  range <- seq(min(mods, na.rm = TRUE) - 1, max(mods, na.rm = TRUE) + 1, by = 0.001)
  
  return(list(result = res_mods, num_studies = k, num_countries = c, 
              plot = plot(mods, ri.m) + 
                lines(range, predict(linear_model, newdata = data.frame(mods = range)), col = "black"),
              details = data.frame(ci.lower = res_mods$ci.lb, ci.upper = res_mods$ci.ub),
              predict_high = pred_high, predict_low = pred_low))
}

### Secondary meta-analysis. Using formulas provided by Schmidt & Oh (2013).
### Outputs are arranged in alignment with the tables reported in their results section.
calculation_secondary <- function(secondary_results) {
  secondary_df <- data.frame()
  for (i in 1:length(secondary_results)) {
    if (!is.null(secondary_results[i][[1]])) {
      df <- as.data.frame(secondary_results[i], row.names = names(secondary_results[i]))
      colnames(df) <- names(secondary_results[i][[1]])
      secondary_df <- rbind(secondary_df, df)
    }
  }
  secondary_df <- secondary_df %>% filter(k > 1)
  secondary_df <- secondary_df %>% mutate(across(.cols = c(k:sd_rho, se_rho), \(x) as.numeric(x)))
  
  secondary_df$wi <- secondary_df$k/(secondary_df$var_rho)      # Equation 6c.
  
  secondary_df$grand_rho <- sum(secondary_df$wi * secondary_df$rho)/sum(secondary_df$wi)   # Equation 6b.
  
  secondary_df$samp_error_ind <- (secondary_df$rho)^2 / secondary_df$k
  
  secondary_df$samp_error <- nrow(secondary_df) / sum(secondary_df$wi)   # Equation 6e.
  
  # Equation 6a.
  secondary_df$obs_var <- sum(secondary_df$wi * (secondary_df$rho - secondary_df$grand_rho)^2)/sum(secondary_df$wi)
  
  # Equation 6.
  secondary_df$true_var <- ifelse(secondary_df$obs_var - secondary_df$samp_error >= 0, secondary_df$obs_var - secondary_df$samp_error, 0)
  
  secondary_df$r_rhorho <- secondary_df$true_var / secondary_df$obs_var
  
  secondary_df$ProVar <- 1 - secondary_df$r_rhorho
  
  secondary_df$rho_corrected <- secondary_df$r_rhorho * (secondary_df$rho - secondary_df$grand_rho) + secondary_df$grand_rho
  
  secondary_df$mag_change <- secondary_df$rho_corrected - secondary_df$rho
  secondary_df$prop_change <- abs(secondary_df$rho_corrected/secondary_df$rho - 1)
  return(secondary_df)
}
