# # # # # # # # # # # # # # # # # # # # # # # # # #
#                                                 #
#     UTILITY FUNCTIONS                           #
#                                                 #
# # # # # # # # # # # # # # # # # # # # # # # # # #


# getIPD ----------------------------------------------------------------------

# This function requires the study shorthand (study), i.e. "buntrock2015", and
# a Github personal access token (ask Mathias)

getIPD = function(study, token){
  
  # Requirements
  require(gh)
  require(base64enc)
  require(jsonlite)
  require(readr)
  require(magrittr)
  
  # Define path
  ghBase = "GET https://api.github.com/repos/prevdep-ipd/data/contents/tidy/"
  suffix = "/data.json"
  path = paste0(ghBase, study, suffix)
  
  # Send API request
  req = try({gh::gh(path, 
                    .token = token,
                    .accept = "application/vnd.github.raw")}, 
            silent = TRUE)
  
  if (identical(class(req)[1], "try-error"))
    stop(study, " could not be retrieved from Github. Either the entry does not exist, ",
         "or the personal access token is not valid.")
  
  # Make pretty
  req %>% 
    rawToChar() %>% jsonlite::fromJSON() %>% 
    apply(2, function(x){
      suppressWarnings({
        readr::parse_guess(x, locale = readr::locale(decimal_mark = "."))
      }) -> col
      ifelse(is.null(attr(col, "problems")), return(col), return(x))
    }, simplify = FALSE) %>%
    as.data.frame() -> dataClean
  
  return(dataClean)
  
}


# getIPDMetadata ------------------------------------------------------------------

getIPDMetadata = function(study, token){
  
  # Requirements
  require(gh)
  require(base64enc)
  require(jsonlite)
  require(readr)
  require(magrittr)
  
  # Define path
  ghBase = "GET https://api.github.com/repos/prevdep-ipd/data/contents/tidy/"
  suffix = "/metadata.json"
  x = paste0(ghBase, study, suffix)
  
  res <- try({gh::gh(x, .token = token)}, silent = TRUE) 
  base64enc::base64decode(res$content) %>% 
    rawToChar() %>% jsonlite::fromJSON() %>% return()
  
}


# getIPDInfo ------------------------------------------------------------------

# This function requires the metadata key, i.e. "rob", and
# a Github personal access token (ask Mathias)

getIPDInfo = function(key, token){
  
  # Requirements
  require(gh)
  require(base64enc)
  require(jsonlite)
  require(readr)
  require(magrittr)
  
  # Define path
  ghBase = "GET https://api.github.com/repos/prevdep-ipd/data/contents/tidy/"
  suffix = "/metadata.json"
  
  # Send API request to get available folders
  req = try({gh::gh(ghBase, .token = token)}, silent = TRUE)
  
  # Get shorthands
  shorthands = vapply(req, "[[", "", "name")
  rm(req)
  
  # Get specific metadata
  paths = paste0(ghBase, shorthands, suffix)
  lapply(as.list(paths)[1], function(x){
    try({gh::gh(x, .token = token)}, silent = TRUE) %>% 
      {base64enc::base64decode(.$content)} %>% 
      rawToChar() %>% jsonlite::fromJSON() %>% 
      {list(author = .[["author"]], key = .[[key]])}
  }) %>% 
    {x = lapply(., function(y) y$key);
    names(x) = lapply(., function(y) y$author); x} %>%  
    return()
  
}

# getIPDStudies ---------------------------------------------------------------

getIPDStudies = function(token){
  
  # Requirements
  require(gh)
  require(base64enc)
  require(jsonlite)
  require(readr)
  require(magrittr)
  
  # Define path
  ghBase = "GET https://api.github.com/repos/prevdep-ipd/data/contents/tidy/"
  
  # Send API request to get available folders
  req = try({gh::gh(ghBase, .token = token)}, silent = TRUE)
  if (identical(class(req)[1], "try-error")) {stop(req[1])}
  shorthands = vapply(req, "[[", "", "name")
  return(shorthands)
  
}


# getIPDKeys ------------------------------------------------------------------

# This function requires a Github personal access token (ask Mathias)

getIPDKeys = function(token){
  
  # Requirements
  require(gh)
  require(base64enc)
  require(jsonlite)
  require(readr)
  require(magrittr)
  require(tibble)
  
  # Define path
  ghBase = "GET https://api.github.com/repos/prevdep-ipd/data/contents/tidy/"
  suffix = "/metadata.json"
  
  # Send API request to get available folders
  req = try({gh::gh(ghBase, .token = token)}, silent = TRUE)
  
  # Get shorthands
  shorthands = vapply(req, "[[", "", "name")
  rm(req)
  
  # Get specific metadata
  paths = paste0(ghBase, shorthands, suffix)
  lapply(as.list(paths), function(x){
    try({gh::gh(x, .token = token)}, silent = TRUE) %>% 
      {base64enc::base64decode(.$content)} %>% 
      rawToChar() %>% jsonlite::fromJSON() %>% 
      {list(author = .[["author"]], keys = names(.))}
  }) %>% 
    {x = lapply(., function(y) y$key);
    names(x) = lapply(., function(y) y$author); x} %>%  
    return()
  
}


w.geomean = function(x, w, ...) exp(weighted.mean(log(x), w, ...))


NNT = function (d, CER = NULL, event.e = NULL, n.e = NULL, event.c = NULL,
                n.c = NULL, names = NULL, method = NULL)
{
  if (!missing(CER) & is.numeric(CER)){
    if (sum(CER < 0 | CER > 1) > 0){
      stop("'CER' range must be between 0 and 1.")
    }
  }
  if (missing(event.e)) {
    if (missing(CER)) {
      NNT = 1/((2 * pnorm(d/sqrt(2)) - 1))
      class(NNT) = c("NNT", "kk", "numeric")
    }
    else {
      if (missing(method)) {
        NNT = 1/(pnorm(d + qnorm(CER)) - CER)
        class(NNT) = c("NNT", "fl", "numeric")
      }
      else {
        if (method == "KraemerKupfer") {
          NNT = 1/((2 * pnorm(d/sqrt(2)) - 1))
          class(NNT) = c("NNT", "kk", "numeric")
        }
        else {
        }
        if (method == "Furukawa") {
          if (missing(CER) | class(CER) != "numeric") {
            stop("To use Furukawa's method, provide a numeric value for CER. \n")
          }
          else {
            NNT = 1/(pnorm(d + qnorm(CER)) - CER)
            class(NNT) = c("NNT", "fl", "numeric")
          }
        }
        else {
        }
      }
    }
  }
  else {
    if (class(event.e) == "numeric" & missing(d)) {
      EER = event.e/n.e
      CER = event.c/n.c
      NNT = abs(1/(EER - CER))
      class(NNT) = c("NNT", "raw", "numeric")
      if (missing(method)) {
      }
      else {
        if (method %in% c("KraemerKupfer", "Furukawa")) {
          class(NNT) = c("NNT", "unspecified", "numeric")
        }
        else {
        }
      }
    }
  }
  if (missing(names)) {
    return(NNT)
  }
  else {
    data = data.frame(Name = names, NNT = NNT)
    class(data) = c(class(NNT)[1:2], "data.frame")
    class(data$NNT) = "numeric"
    NNT = data
    return(NNT)
  }
}

clusterScale = function (x, cluster, adj = FALSE, group = NULL) 
{
  isname <- c(length(x) == 1, length(cluster) == 1, length(group) == 
                1) & c(is.character(x), is.character(cluster), is.character(group))
  if (any(isname)) {
    parent <- parent.frame()
    if (isname[1]) 
      x <- eval(parse(text = x), parent)
    if (isname[2]) 
      cluster <- eval(parse(text = cluster), parent)
    if (isname[3]) 
      group <- eval(parse(text = group), parent)
  }
  if (!is.numeric(cluster)) 
    cluster <- as.integer(cluster)
  cluster <- match(cluster, unique(cluster))
  n.obs <- rowsum(as.integer(!is.na(x)), cluster)
  gm <- rowsum(x, cluster, na.rm = T)/n.obs
  gsd <- unlist(lapply(split(x, cluster), sd))
  gm[is.nan(gm)] <- NA
  gsd[is.nan(gsd)] <- NA
  gm <- gm[cluster]
  gsd <- gsd[cluster]
  return(as.numeric((x - gm)/gsd))
}


extract.facs = function(x){
  
  reslist = list()
  for (i in 1:length(x)){
    
    name = x[[i]]$descript
    value = as.character(x[[i]]$values$value)
    freq = x[[i]]$values$frequency
    prop = freq/sum(freq)
    
    reslist[[i]] = data.frame(name, value, freq, prop)
    
  }
  
  res = do.call(rbind, reslist)
  return(res)
  
}

pooled.sd = function(y, group) {
  sd = as.numeric(lapply(split(y, group), function(x){sd(x, na.rm = TRUE)}))
  n = as.numeric(lapply(split(y, group), function(x){length(x)}))
  sqrt(sum((n-1)*sd^2)/sum(n-1))
}


lm.extractor.tci = function(m, study = NA) {
  require(tidyverse)
  require(stringr)
  pooled.sd(m$model[,1], m$model[,2]) -> sd.p
  data.frame(study, 
             as_tibble(summary(m)$coefficients, rownames="coef") %>% 
               filter(str_starts(coef,"group:")) %>% pull("Estimate"),
             as_tibble(summary(m)$coefficients, rownames="coef") %>% 
               filter(str_starts(coef,"group:")) %>% pull("Std. Error")) %>% 
    {colnames(.) = c("study", "es", "se");.} %>% 
    mutate(es = es/sd.p, se = se/sd.p)
}

lm.extractor.fp = function(tmp, which, name = NA) {
  tmp %>% 
    {table(fp=.[[which]], group=.[["group"]])} %>% 
    {data.frame(study = name, event.e = .[2,2], event.c = .[2,1],
                n.e = sum(.[,2]), n.c = sum(.[,1]))} 
}

lm.extractor.ri = function(tmp, which, name = NA) {
  tmp %>% 
    {table(ri=.[[which]], group=.[["group"]])} %>% 
    {data.frame(study = name, event.e = .[2,2], event.c = .[2,1],
                n.e = sum(.[,2]), n.c = sum(.[,1]))} 
}

lm.extractor = function(m, study = NA) {
  require(tidyverse)
  pooled.sd(m$model[,1], m$model[,2]) -> sd.p
  data.frame(study, 
             summary(m)$coefficients["group",1],
             summary(m)$coefficients["group",2]) %>% 
    {colnames(.) = c("study", "es", "se");.} %>% 
    mutate(es = es/sd.p, se = se/sd.p)
}

lm.extractor.rd = function(tmp, which, name = NA) {
  tmp %>% 
    {table(rd=.[[which]], group=.[["group"]])} %>% 
    {data.frame(study = name, event.e = .[2,2], event.c = .[2,1],
                n.e = sum(.[,2]), n.c = sum(.[,1]))} 
}


eff.extractor = function(m, bgroup.name = "group1", model.name = "model"){
  
  m = mitml::testEstimates(m, var.comp = T)
  est = m$estimates
  
  return(
    data.frame(
      name = model.name,
      b = est[bgroup.name, "Estimate"],
      b.lower = confint(m)[bgroup.name, 1],
      b.upper = confint(m)[bgroup.name, 2],
      t = est[bgroup.name, "t.value"],
      p = est[bgroup.name, 5],
      tau.intercept = m$var.comp[1],
      tau.slope = m$var.comp[3]))
  
}


mod.extractor = function(m, mod.name = "age", model.name = "model"){
  
  m = mitml::testEstimates(m, var.comp = T)
  est = m$estimates
  rows = rownames(est)[grep(mod.name, rownames(est))]
  
  return(
    data.frame(
      name = model.name,
      b = est[rows, "Estimate"],
      b.lower = confint(m)[rows, 1],
      b.upper = confint(m)[rows, 2],
      t = est[rows, "t.value"],
      p = est[rows, 5],
      tau.intercept = m$var.comp[1],
      tau.slope = m$var.comp[3]))
  
}



desc.extractor = function(variable, group, implist = NULL) {
  
  require(purrr)
  require(dplyr)
  
  catch = function(argname, matchcall, data, encl) {
    eval(matchcall[[match(argname, names(matchcall))]], data, enclos = encl)
  }
  
  parentFrame = sys.frame(sys.parent())
  mC = match.call() 
  if (is.null(implist)) { 
    data = parentFrame
    var = try({catch("variable", mC, data, parentFrame)}, silent = TRUE)
    grp = try({catch("group", mC, data, parentFrame)}, silent = TRUE)
    if (!is.numeric(var))
      stop("'variable' must be numeric. Use 0/1 for binary variables.")
    data.frame(variable = as.character(mC$variable)[length(mC$variable)], 
               grp = "overall",
               m = mean(var, na.rm=TRUE), sd = sd(var, na.rm=TRUE), 
               n = length(var[!is.na(var)])) -> last.row
    split(var, grp) -> split.dat
    names(split.dat) -> col
    lapply(split.dat, function(x) {
      cbind(m = mean(x, na.rm=TRUE), 
            sd = sd(x, na.rm=TRUE), 
            n = length(x[!is.na(x)]))
    }) %>% 
      do.call(rbind, .) %>% data.frame() %>% 
      cbind(variable = as.character(mC$variable)[length(mC$variable)], 
            grp = col, .) %>% 
      rbind(., last.row) %>% return()
  } else {
    m = length(implist)
    var = mC$variable %>% as.character()
    grp = mC$group %>% as.character()
    purrr::map(implist, function(x){
      split(x[[var]], x[[grp]]) %>% 
        lapply(function(y) {
          cbind(m = mean(y, na.rm=TRUE), 
                sd = sd(y, na.rm=TRUE), 
                n = length(y[!is.na(y)]))
        }) %>% do.call(rbind, .) %>% data.frame
    }) %>% Reduce(`+`, .) %>% {./m} -> ret
    names(split(implist[[1]][[var]], 
                implist[[1]][[grp]])) -> col
    cbind(variable = as.character(mC$variable), 
          grp = col, data.frame(ret)) %>% return()
  }
}


# Reliable Change Index
# Define RCI function
rci = function(pre, post, data, rtt, sdNorm){
  if ("rci" %in% names(data)) {warning("variable rci already exists in data frame and will be overwritten")}
  difference <- data[,post]-data[,pre]
  Serr <- sdNorm*sqrt(1-rtt)
  Sdiff <- sqrt((2*Serr^2))
  rci <- difference/Sdiff
  return(rci)
}

calculateRci = function(pre, post, rtt, sdNorm){
  if ("rci" %in% names(data)) {warning("variable rci already exists in data frame and will be overwritten")}
  difference <- post-pre
  sdNorm = sd(pre, na.rm = TRUE)
  Serr <- sdNorm*sqrt(1-rtt)
  Sdiff <- sqrt((2*Serr^2))
  rci <- difference/Sdiff
  return(rci)
}


# Backtransformation of factor dummy vars
backtrans.dummy = function(varname, dat){
  
  varname = varname
  dat = dat
  
  # Select dummy columns
  cols = dat %>% dplyr::select(grep(varname, colnames(.))) %>% 
    map_dfr(function(x) as.numeric(x)-1)
  
  # Calculate zero class variable and order
  cols[[paste0(varname, ".0")]] = 1 - rowSums(cols)
  cols = cols[,order(colnames(cols))]
  
  # Get levels from variable names
  levs = unlist(strsplit(colnames(cols), '[.]'))[2 * (1:ncol(cols))]
  
  # Merge
  res = factor(as.matrix(as.data.frame(cols)) %*% 1:ncol(cols))
  
  levels(res) = as.numeric(levels(res))-1
  
  return(res)
  
}


# Coord Radar
coord_radar <- function (theta = "x", start = 0, direction = 1) {
  theta <- match.arg(theta, c("x", "y"))
  r <- if (theta == "x") "y" else "x"
  ggproto("CordRadar", CoordPolar, theta = theta, r = r, start = start, 
          direction = sign(direction),
          is_linear = function(coord) TRUE)
}



# Accuracy Cutoff CV
acc_cutoff = function(y, ypred){
  
  n = length(y)
  y = as.numeric(as.character(y))
  ypred = as.numeric(as.character(ypred))
  
  res = list()
  for (i in 1:100/100){
    
    x = i*100
    acc = sum(ifelse(ypred > i, 1, 0) == y)/n
    
    res[[x]] = data.frame(thresh = i,
                          accuracy = acc)
    
  }
  
  res.df = do.call(rbind, res)
  res.df = res.df[order(-res.df$accuracy),]
  
  return(res.df)
  
}


CV.acc_cutoff = function(model, X, y, k = 10){
  
  require(dplyr)
  
  if(class(sl_l) != "SuperLearner"){
    stop("class of model must be 'SuperLearner'")
  }
  
  # Define split folds
  ksize = ceiling(nrow(X)*(1/k))
  fold = sample(rep(1:k, times = ksize))[1:nrow(X)]
  
  cv.list = list()
  # Define CV loop
  for (i in 1:k){
    
    data = X[fold != i, ]
    val_dat = X[fold == i, ]
    y_i = y[fold != i]
    y_e = y[fold == i]
    
    # Get internal estimate
    ypred_i = SuperLearner::predict.SuperLearner(model, newdata = data)$pred %>% as.numeric()
    thresh = acc_cutoff(y_i, ypred_i)[1,1]
    acc_internal = acc_cutoff(y_i, ypred_i)[1,2]
    
    # Get external estimate
    ypred_e = SuperLearner::predict.SuperLearner(model, newdata = val_dat)$pred %>% as.numeric()
    acc_external = sum(ifelse(ypred_e > thresh, 1, 0) == y_e)/length(ypred_e)
    
    cv.list[[i]] = data.frame(fold = i,
                              thresh, acc_internal, acc_external)
    
  }
  
  cv.runs = do.call(rbind, cv.list)
  
  return(list(cv.runs = cv.runs,
              results = data.frame(tresh = mean(cv.runs$thresh),
                                   tresh.sd = sd(cv.runs$thresh),
                                   acc_internal = mean(cv.runs$acc_internal),
                                   acc_internal.sd = sd(cv.runs$acc_internal),
                                   acc_external = mean(cv.runs$acc_external),
                                   acc_external.sd = sd(cv.runs$acc_external))))
  
}


# Function for iml
f_pred = function(model, newdata){
  
  preds = SuperLearner::predict.SuperLearner(model, 
                                             newdata = newdata, 
                                             onlySL = TRUE)
  return(data.frame("response" = as.numeric(preds$pred),
                    "nonresponse" = 1-as.numeric(preds$pred)))
  
}


# Function to extract descriptives
skimReport = function(data, round = FALSE){
  
  require(skimr)
  require(dplyr)
  require(purrr)
  x = skim(data)
  N = nrow(data)
  
  # Summarize Factors
  with(x, {
    skim_type == "factor" -> fac.mask
    vars = skim_variable[fac.mask]
    n.unique = factor.n_unique[fac.mask]
    strsplit(factor.top_counts, "\\, |\\:")[fac.mask] %>% 
      purrr::map(function(x){
        as.numeric(x) -> x
        data.frame(category = x[seq(1, length(x), by = 2)],
                   count = x[seq(2, length(x), by = 2)]) %>% 
          dplyr::mutate(percent = count/N)
      }) %>% 
      {names(.) = vars;.} %>% 
      map_df(~as.data.frame(.), .id = "variable")
  }) -> factors
  
  # Summarize Numeric
  with(x, {
    skim_type == "numeric" -> num.mask
    data.frame(variable = skim_variable[num.mask],
               mean = numeric.mean[num.mask],
               sd = numeric.sd[num.mask],
               n = N-x$n_missing[num.mask],
               n.missing = n_missing[num.mask],
               perc.missing = n_missing[num.mask]/N)
  }) -> numerics
  
  if (round == TRUE){
    within(factors, {
      percent = round(percent*100, 2)
    }) -> factors
    within(numerics, {
      mean = round(mean, 2)
      sd = round(sd, 2)
      perc.missing = round(perc.missing*100, 2)
    }) -> numerics
  }
  
  
  dat = list(factors = factors,
             numerics = numerics)
  class(dat) = c("list", "skimReport")
  
  return(dat)
  
}

print.skimReport = function(x){
  cat("Categorial Variables \n ------------------ \n")
  cat("\n")
  print(x$factors)
  cat("\n")
  cat("Numeric Variables \n ------------------ \n")
  cat("\n")
  print(x$numerics)
} 


# Find unidentified cases and structural zeros (imputation)
find.unidentified = function(label, data, use.glm=TRUE, strict.glm=use.glm, tol=1e-5){
  # label = name of variable in data (character)
  # data = data set (data.frame)
  
  data <- data[!is.na(data[,label]),]
  preds <- setdiff(colnames(data),label)
  is.cat <- is.factor(data[,label])
  
  # structural zeros
  z <- colSums(data[,preds]) == 0
  
  fml <- formula(paste(label, "~", paste(preds, collapse="+")))
  if(use.glm){
    fit <- glm(fml, data=data, family=if(is.cat) binomial() else gaussian())
  }else{
    if(is.cat) data[,label] <- as.numeric(data[,label])-1
    fit <- lm(fml, data=data)
  }
  # linear depencies without zeros
  na <- is.na(coef(fit)[-1])
  ld <- na & !z
  if(is.cat && strict.glm){ # filter cor. effects in GLMs not showing in LM
    crm <- cov2cor(vcov(fit)[-1,-1][!na,!na])
    ind <- lower.tri(crm) & (abs(crm) > 1-tol)
    ld.cat <- rowSums(ind) > 0
  }else{
    ld.cat <- NULL
  }
  if(any(ld.cat)) ld[which(ld.cat)] <- TRUE
  
  # get names
  out <- names(z)[z|ld]
  if(any(z)) attr(out, "struct.zero") <- names(z)[z]
  if(any(ld.cat)) attr(out, "perfect.cor.glm") <- names(ld.cat)[ld.cat]
  return(out)
  
}


extract_percs = function(x){
  
  x = x
  
  # Define total Ns
  ngroup = c(implist[[1]][implist[[1]]$ipd_group == 1,] %>% nrow() %>% rep(2),
             implist[[1]][implist[[1]]$ipd_group == 2,] %>% nrow() %>% rep(2))
  ntotal = implist[[1]] %>% nrow() %>% rep(4)
  
  # Extract Ns & Calulcate Percentages
  implist %>% 
    map(~ group_by(., ipd_group, !!sym(x)) %>% 
          dplyr::summarise(n = n()) %>%
          as.data.frame() %>% 
          mutate_if(is.factor, as.numeric)) %>% 
    purrr::reduce(`+`)/length(implist) -> distribution
  
  distribution %>% 
    mutate(n.perc.group = n/ngroup,
           n.perc.total = n/ntotal) -> distribution
  
  return(distribution)
  
}


calculateFp = function(y1, y0) {
  {y1 <= y0*0.5} %>% as.numeric()
}

calculateCsf = function(y1) {
  {y1 < 5} %>% as.numeric()
}


# fitPoisson ------------------------------------------------------------------

fitPoisson = function(f, rnd, data){
  
  require(stringr)
  
  f = str_replace_all(as.character(f), "group", "T")
  rnd = str_replace_all(as.character(rnd), "group", "T")[2]
  
  frml = as.formula(paste(f[2], "~", f[3], "+", rnd))
  lapply(data, function(x){
    within(x, {T = group})
  }) -> data
  data = as.mitml.list(data)
  
  lapply(data, function(x){
    lme4::glmer(frml, family = poisson(link = "log"),
                data = x,
                control=glmerControl(optimizer="bobyqa", optCtrl=list(maxfun=100000))) # added because of convergence problem
  }) -> res
  
  class(res) = c("mitml.result", "list")
  return(res)
}


fitPoissonMAP = function(f, rnd, data){
  
  require(stringr)
  require(blme)
  
  f = str_replace_all(as.character(f), "group", "T")
  rnd = str_replace_all(as.character(rnd), "group", "T")[2]
  frml = as.formula(paste(f[2], "~", f[3], "+", rnd))
  lapply(data, function(x){
    within(x, {T = group})
  }) -> data
  data = as.mitml.list(data)
  
  lapply(data, function(x){
    blme::bglmer(frml, family = poisson(link = "log"),
                 data = x,
                 control=glmerControl(optimizer="bobyqa", optCtrl=list(maxfun=1000000))) # added because of convergence problem; check.conv.grad=.makeCC("warning", tol = 3e-3, relTol = NULL)
  }) -> res
  
  class(res) = c("mitml.result", "list")
  return(res)
}


fitPoissonCoef = function(m){
  
  require(mitml)
  require(marginaleffects)
  require(mice)
  require(purrr)
  
  te = testEstimates(m, extra.pars = TRUE)
  coefs = as.list(te$estimates["T",])
  coefs["tau2"] = te$extra.pars[2,]
  
  # Calculate marginal IRR
  message("Calculate marginal IRR...")
  lapply(m, function(x){
    avg_comparisons(x, comparison = "lnratio",
                    variables = list(T = 0:1))
  }) -> gcomp
  
  pool.scalar(Q = gcomp %>% map_dbl(~.$estimate),
              U = gcomp %>% map_dbl(~.$std.error^2)) -> irr.pool
  coefs["gamma.FMI"] = irr.pool$fmi
  coefs["r.RIV"] = irr.pool$r
  coefs["k"] = nrow(lme4::ranef(m[[1]])[[1]])
  coefs["irr"] = exp(irr.pool$qbar)
  coefs["log_irr.se"] = sqrt(irr.pool$t)
  coefs["irr.lo"] = exp(with(irr.pool, qbar - qnorm(.975) * sqrt(t)))
  coefs["irr.hi"] = exp(with(irr.pool, qbar + qnorm(.975) * sqrt(t)))
  
  # Calculate marginal incidences
  message("Calculate marginal incidences...")
  lapply(m, function(x){
    list(lo = avg_comparisons(x, comparison = 
                                function(hi, lo) lo)[1,c("estimate", "std.error")],
         hi = avg_comparisons(x, comparison = 
                                function(hi, lo) hi)[1,c("estimate", "std.error")])
  }) -> gcomp.risk
  
  pool.scalar(Q = gcomp.risk %>% map_dbl(~.$lo$estimate),
              U = gcomp.risk %>% map_dbl(~.$lo$std.error^2)) -> risk.lo.pool
  pool.scalar(Q = gcomp.risk %>% map_dbl(~.$hi$estimate),
              U = gcomp.risk %>% map_dbl(~.$hi$std.error^2)) -> risk.hi.pool
  
  coefs["lo"] = risk.lo.pool$qbar
  coefs["lo.lo"] = with(risk.lo.pool, qbar - qnorm(.975) * sqrt(t))
  coefs["lo.hi"] = with(risk.lo.pool, qbar + qnorm(.975) * sqrt(t))
  
  coefs["hi"] = risk.hi.pool$qbar
  coefs["hi.lo"] = with(risk.hi.pool, qbar - qnorm(.975) * sqrt(t))
  coefs["hi.hi"] = with(risk.hi.pool, qbar + qnorm(.975) * sqrt(t))
  
  # Calculate risk difference, NNT
  message("Calculate NNTs...")
  lapply(m, function(x){
    avg_comparisons(x, 
                    comparison = "difference")[1,c("estimate", "std.error")]
  }) -> gcomp.rd
  
  pool.scalar(Q = gcomp.rd %>% map_dbl(~.$estimate),
              U = gcomp.rd %>% map_dbl(~.$std.error^2)) -> rd.pool
  
  coefs["nnt"] = abs(rd.pool$qbar)^-1
  coefs["nnt.lo"] = abs(with(rd.pool, qbar - qnorm(.975) * sqrt(t)))^-1
  coefs["nnt.hi"] = abs(with(rd.pool, qbar + qnorm(.975) * sqrt(t)))^-1
  
  return(as.data.frame(coefs))
  
}


calcIrr = function(data, variable){
  
  require(purrr)
  require(mice)
  require(dplyr)
  
  data %>% 
    map(function(x){
      x %>%   
        split(.$id.study) %>% 
        map(function(y){
          tab = as.data.frame(
            table(group = y$group, status = y[[variable]])+0.5)
          n.cg = tab %>% filter(group==0) %>% {sum(.$Freq)}
          n.ig = tab %>% filter(group==1) %>% {sum(.$Freq)}
          e.cg = tab %>% filter(group==0, status == 1) %>% {sum(.$Freq)}
          e.ig = tab %>% filter(group==1, status == 1) %>% {sum(.$Freq)}
          if (e.cg == 0 | e.ig == 0){e.cg = 0.5; e.ig = 0.5}
          data.frame(log_rr = log((e.ig/n.ig)/(e.cg/n.cg)),
                     se = sqrt(1/e.ig + 1/e.cg + 1/n.ig + 1/n.cg))
        }) %>% bind_rows()
    }) -> list.2stage
  
  suppressMessages({list.2stage %>% map(~.[,"log_rr"]) %>% bind_cols()}) -> rrs
  suppressMessages({list.2stage %>% map(~.[,"se"]) %>% bind_cols()}) -> ses
  
  retlist = list()
  for (i in 1:nrow(rrs)){
    ps = pool.scalar(Q = as.numeric(rrs[i,]), 
                     U = as.numeric(ses[i,]^2))[c("qbar", "t")] %>% 
      as.data.frame()
    ps$id.study = unique(data[[1]]$id.study)[i]
    retlist[[i]] = ps
  }
  bind_rows(retlist) %>% 
    dplyr::mutate(id.study, yi = qbar, vi = t) %>% 
    dplyr::select(id.study, yi, vi) %>% return()
  
}


fitIrr2Stage = function(x, knha = FALSE){
  
  meta::metagen(TE = yi, seTE = sqrt(vi), hakn = knha,
                method.tau = "REML", data = x, 
                common = FALSE, sm="RR") %>% 
    {data.frame(k = .$k,
                irr = exp(.$TE.random),
                irr.lo = exp(.$lower.random),
                irr.hi = exp(.$upper.random),
                pred.lo = exp(.$lower.predict),
                pred.hi = exp(.$upper.predict),
                tau2 = .$tau2, i2 = .$I2,
                logirr = .$TE.random,
                se_logirr = .$seTE.random,
                w.s = sum(.$w.common),
                w.ssq = sum(.$w.common^2))} -> res
  res$v.tilde = ((res$k-1)*res$w.s)/(res$w.s^2-res$w.ssq)
  return(res)
  
}

# g-computation for JAGS model ------------------------------------------------
# https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5790647/

gComputation = function(data, mcmc) {
  
  if (!is.null(data$k.agg[1])) {
    agg.list = list()
    for (i in 1:data$k.agg) {
      rbind(
        cbind(y = rep(1,data$event.c[i]), group = rep(0,data$event.c[i])),
        cbind(y = rep(0,data$n.c[i]-data$event.c[i]), group = rep(0,data$n.c[i]-data$event.c[i])),
        cbind(y = rep(1,data$event.e[i]), group = rep(1,data$event.e[i])),
        cbind(y = rep(0,data$n.e[i]-data$event.e[i]), group = rep(1,data$n.e[i]-data$event.e[i]))
      ) %>% as.data.frame() %>% mutate(cov = 0, study = data$k.ipd+i) -> agg.list[[i]]
    }
    rbind(
      with(data, {data.frame(y = y, group = treat, 
                             cov = cov, study = study)}),
      do.call(rbind, agg.list)
    ) -> dat.expanded
  } else {
    with(data, {data.frame(y = y, group = treat, 
                           cov = cov, study = study)}
    ) -> dat.expanded
  }
  
  p1 = matrix(NA, nrow = 10000, ncol = nrow(dat.expanded))
  p0 = matrix(NA, nrow = 10000, ncol = nrow(dat.expanded))
  message("Performing g-computation:")
  for (i in 1:nrow(dat.expanded)) {
    message(paste0("iter ", i))
    samples = sample(1:nrow(mcmc),10000)
    p1[,i] = {(mcmc[samples,paste0("alpha[", dat.expanded$study[i],"]")]) + 
        ((mcmc[samples,"mu"] + rnorm(length(samples),0,mcmc[samples,"tau"]))) +
        (ifelse(i<=data$n.ipd,
                mcmc[samples,paste0("beta0[", dat.expanded$study[i],"]")] * dat.expanded$cov[i],0))} %>% 
      plogis()
    p0[,i] = {(mcmc[samples,paste0("alpha[", dat.expanded$study[i],"]")]) + 
        (ifelse(i<=data$n.ipd,
                mcmc[samples,paste0("beta0[", dat.expanded$study[i],"]")] * dat.expanded$cov[i],0))} %>% 
      plogis()
  }
  message("DONE")
  
  mean.risk.1 = rowMeans(p1)
  mean.risk.0 = rowMeans(p0)
  rr = mean.risk.1 / mean.risk.0
  
  return(data.frame(mean.risk.1 = mean.risk.1, 
                    mean.risk.0 = mean.risk.0,
                    rr = rr))
}





# fitLme ----------------------------------------------------------------------

fitLme = function(f, rnd, weights, data){
  
  require(purrr)
  require(furrr)
  require(stringr)
  require(nlme)
  
  f = as.character(f)
  rnd = as.character(rnd)
  
  f = paste(str_replace_all(f, "group", "T")[2], "~",
            str_replace_all(f, "group", "T")[3]) %>% as.formula()
  rnd = paste("~", str_replace_all(rnd, "group", "T")[2]) %>% as.formula()
  lapply(data, function(x){
    within(x, {T = group; group = NULL})
  }) -> data
  data = as.mitml.list(data)
  
  message("Fitting models...")
  plan(multisession, workers = 3)
  data %>% 
    future_map(function(x){
      nlme::lme(fixed = f, random = rnd, weights = weights,
                control = lmeControl(msMaxIter = 2e4),
                data = x)
    }, .progress = TRUE) -> res
  
  class(res) = c("mitml.result", "list")
  return(res)
  
}

fitLmeCoef = function(m, cer){
  
  require(purrr)
  require(nlme)
  
  # Calculate sd.pooled
  message("Calculating SDpooled...")
  names(attr(m[[1]]$terms, 
             "dataClasses")[1]) -> outcome
  
  m %>% 
    map(~.$data) %>% 
    map_dbl(function(x){
      x %>% 
        split(.$id.study) %>% 
        map(function(y){
          y[c("T", outcome)] %>% 
            split(.$T) %>% 
            lapply(function(z){
              data.frame(n = nrow(z),
                         sd = sd(z[,2], na.rm=T))
            }) %>% bind_rows()
        }) %>% bind_rows() %>% 
        {w.geomean(x=.$sd, w=.$n)} 
    }) %>% 
    {exp(mean(log(.)))} -> sd.pooled
  
  # Extract coefs
  message("Extract coefficients...")
  coefs = as.list(testEstimates(m)$estimates["T",])
  coefs["tau2"] = m %>% 
    map_dbl(~as.numeric(nlme::VarCorr(.)["T", "Variance"])) %>% 
    {exp(mean(log(.)))}
  coefs["gamma.FMI"] = coefs$FMI
  coefs["r.RIV"] = coefs$RIV
  
  coefs["k"] = nrow(ranef(m[[1]]))
  coefs["g"] = coefs$Estimate/sd.pooled
  coefs["g.lo"] = confint(testEstimates(m))["T",1]/sd.pooled
  coefs["g.hi"] = confint(testEstimates(m))["T",2]/sd.pooled
  coefs["g.tau2"] = (sqrt(coefs$tau2)/sd.pooled)^2
  coefs["sd.pooled"] = sd.pooled
  
  coefs["nnt"] = NNT(abs(coefs[["g"]]), cer)
  coefs["nnt.lo"] = NNT(abs(coefs[["g.lo"]]), cer)
  coefs["nnt.hi"] = NNT(abs(coefs[["g.hi"]]), cer)
  
  return(as.data.frame(coefs))
  
}



calcG = function(data, f)
{
  require(purrr)
  require(mice)
  
  data %>% 
    map(function(x) {
      split(x, x$id.study) %>% 
        map(function(y){
          sdp = lapply(split(y[[as.character(f)[2]]], y[["group"]]), sd) %>% 
            unlist() %>% mean()
          res = summary(lm(f, data = y))$coefficients["group",]
          res["yi"] = res["Estimate"]/sdp
          res["vi"] = (res["Std. Error"]/sdp)^2; res
        }) %>% bind_rows() %>% select(yi, vi)
    }) -> list.g
  
  suppressMessages({list.g %>% map(~.[,"yi"]) %>% bind_cols()}) -> yis
  suppressMessages({list.g %>% map(~.[,"vi"]) %>% bind_cols()}) -> vis
  
  retlist = list()
  for (i in 1:nrow(yis)){
    ps = pool.scalar(Q = as.numeric(yis[i,]), 
                     U = as.numeric(vis[i,]))[c("qbar", "t")] %>% 
      as.data.frame()
    ps$id.study = unique(data[[1]]$id.study)[i]
    retlist[[i]] = ps
  }
  
  bind_rows(retlist) %>% 
    dplyr::mutate(id.study, yi = qbar, vi = t) %>% 
    dplyr::select(id.study, yi, vi) %>% return()
}



fitG2Stage = function(x, knha = TRUE){
  
  metagen(TE = yi, seTE = sqrt(vi), hakn = knha,
          method.tau = "REML", data = x, 
          common = FALSE, sm="SMD") %>% 
    {data.frame(k = .$k,
                g = .$TE.random,
                g.lo = .$lower.random,
                g.hi = .$upper.random,
                pred.lo = .$lower.predict,
                pred.hi = .$upper.predict,
                tau2 = .$tau2, i2 = .$I2,
                w.s = sum(.$w.common),
                w.ssq = sum(.$w.common^2))} -> res
  res$v.tilde = ((res$k-1)*res$w.s)/(res$w.s^2-res$w.ssq)
  return(res)
  
}


# fitBinomial -----------------------------------------------------------------

fitBinomial = function(f, rnd, data){
  
  require(stringr)
  
  f = str_replace_all(as.character(f), "group", "T")
  rnd = str_replace_all(as.character(rnd), "group", "T")[2]
  frml = as.formula(paste(f[2], "~", f[3], "+", rnd))
  lapply(data, function(x){
    within(x, {T = group})
  }) -> data
  data = as.mitml.list(data)
  
  lapply(data, function(x){
    lme4::glmer(frml, family = binomial(link = "logit"),
                data = x,
                control=glmerControl(optimizer="bobyqa", optCtrl=list(maxfun=100000))) # added because of convergence problem
  }) -> res
  
  class(res) = c("mitml.result", "list")
  return(res)
}

fitBinomialMAP = function(f, rnd, data){
  
  require(stringr)
  require(blme)
  
  f = str_replace_all(as.character(f), "group", "T")
  rnd = str_replace_all(as.character(rnd), "group", "T")[2]
  frml = as.formula(paste(f[2], "~", f[3], "+", rnd))
  lapply(data, function(x){
    within(x, {T = group})
  }) -> data
  data = as.mitml.list(data)
  
  lapply(data, function(x){
    blme::bglmer(frml, family = binomial(link = "logit"),
                 data = x)
  }) -> res
  
  class(res) = c("mitml.result", "list")
  return(res)
}


fitBinomialCoef = function(m){
  
  require(mitml)
  require(marginaleffects)
  require(mice)
  require(purrr)
  
  te = testEstimates(m, extra.pars = TRUE)
  coefs = as.list(te$estimates["T",])
  coefs["tau2"] = te$extra.pars[2,]
  
  # Calculate marginal IRR
  message("Calculate marginal RR...")
  lapply(m, function(x){
    avg_comparisons(x, comparison = "lnratio",
                    variables = list(T = 0:1))
  }) -> gcomp
  
  pool.scalar(Q = gcomp %>% map_dbl(~.$estimate),
              U = gcomp %>% map_dbl(~.$std.error^2)) -> rr.pool
  coefs["gamma.FMI"] = rr.pool$fmi
  coefs["r.RIV"] = rr.pool$r
  coefs["k"] = nrow(ranef(m[[1]])[[1]])
  coefs["rr"] = exp(rr.pool$qbar)
  coefs["log_rr.se"] = sqrt(rr.pool$t)
  coefs["rr.lo"] = exp(with(rr.pool, qbar - qnorm(.975) * sqrt(t)))
  coefs["rr.hi"] = exp(with(rr.pool, qbar + qnorm(.975) * sqrt(t)))
  
  # Calculate marginal incidences
  message("Calculate marginal risks...")
  lapply(m, function(x){
    list(lo = avg_comparisons(x, comparison = 
                                function(hi, lo) lo)[1,c("estimate", "std.error")],
         hi = avg_comparisons(x, comparison = 
                                function(hi, lo) hi)[1,c("estimate", "std.error")])
  }) -> gcomp.risk
  
  pool.scalar(Q = gcomp.risk %>% map_dbl(~.$lo$estimate),
              U = gcomp.risk %>% map_dbl(~.$lo$std.error^2)) -> risk.lo.pool
  pool.scalar(Q = gcomp.risk %>% map_dbl(~.$hi$estimate),
              U = gcomp.risk %>% map_dbl(~.$hi$std.error^2)) -> risk.hi.pool
  
  coefs["lo"] = risk.lo.pool$qbar
  coefs["lo.lo"] = with(risk.lo.pool, qbar - qnorm(.975) * sqrt(t))
  coefs["lo.hi"] = with(risk.lo.pool, qbar + qnorm(.975) * sqrt(t))
  
  coefs["hi"] = risk.hi.pool$qbar
  coefs["hi.lo"] = with(risk.hi.pool, qbar - qnorm(.975) * sqrt(t))
  coefs["hi.hi"] = with(risk.hi.pool, qbar + qnorm(.975) * sqrt(t))
  
  # Calculate risk difference, NNT
  message("Calculate NNTs...")
  lapply(m, function(x){
    avg_comparisons(x, 
                    comparison = "difference")[1,c("estimate", "std.error")]
  }) -> gcomp.rd
  
  pool.scalar(Q = gcomp.rd %>% map_dbl(~.$estimate),
              U = gcomp.rd %>% map_dbl(~.$std.error^2)) -> rd.pool
  
  coefs["nnt"] = abs(rd.pool$qbar)^-1
  coefs["nnt.lo"] = abs(with(rd.pool, qbar - qnorm(.975) * sqrt(t)))^-1
  coefs["nnt.hi"] = abs(with(rd.pool, qbar + qnorm(.975) * sqrt(t)))^-1
  
  return(as.data.frame(coefs))
  
}



calcRr = function(data, variable){
  
  require(purrr)
  require(mice)
  require(dplyr)
  
  data %>% 
    map(function(x){
      x %>%   
        split(.$id.study) %>% 
        map(function(y){
          tab = as.data.frame(
            table(group = y$group, status = y[[variable]])+0.5)
          n.cg = tab %>% filter(group==0) %>% {sum(.$Freq)}
          n.ig = tab %>% filter(group==1) %>% {sum(.$Freq)}
          e.cg = tab %>% filter(group==0, status == 1) %>% {sum(.$Freq)}
          e.ig = tab %>% filter(group==1, status == 1) %>% {sum(.$Freq)}
          if (e.cg == 0 | e.ig == 0){e.cg = 0.5; e.ig = 0.5}
          data.frame(log_rr = log((e.ig/n.ig)/(e.cg/n.cg)),
                     se = sqrt(1/e.ig + 1/e.cg + 1/n.ig + 1/n.cg))
        }) %>% bind_rows()
    }) -> list.2stage
  
  suppressMessages({list.2stage %>% map(~.[,"log_rr"]) %>% bind_cols()}) -> rrs
  suppressMessages({list.2stage %>% map(~.[,"se"]) %>% bind_cols()}) -> ses
  
  retlist = list()
  for (i in 1:nrow(rrs)){
    ps = pool.scalar(Q = as.numeric(rrs[i,]), 
                     U = as.numeric(ses[i,]^2))[c("qbar", "t")] %>% 
      as.data.frame()
    ps$id.study = unique(data[[1]]$id.study)[i]
    retlist[[i]] = ps
  }
  bind_rows(retlist) %>% 
    dplyr::mutate(id.study, yi = qbar, vi = t) %>% 
    dplyr::select(id.study, yi, vi) %>% return()
  
}


calcRRPhq = function(dat, study = NA) {
  
  require(marginaleffects)
  within(dat, {grp = group}) -> dat
  glm(dep.status.mod ~ grp * scale(cm.dep.0), 
      data = dat, family = binomial()) -> m
  rbind(c(phq = 5, 
          comparisons(m, newdata = datagrid(cm.dep.0 = f.phq(5)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(phq = 10, 
          comparisons(m, newdata = datagrid(cm.dep.0 = f.phq(10)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(phq = 15, 
          comparisons(m, newdata = datagrid(cm.dep.0 = f.phq(15)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")])) %>% 
    as.data.frame() %>% {.$study = study;.}
}

calcRRPhq.firth = function(dat, study = NA) {
  
  require(marginaleffects)
  within(dat, {grp = group}) -> dat
  logistf(dep.status.mod ~ grp * scale(cm.dep.0), 
          data = dat, control=logistf.control(maxit=100)) -> m
  rbind(c(phq = 5, 
          comparisons(m, newdata = datagrid(cm.dep.0 = f.phq(5)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(phq = 10, 
          comparisons(m, newdata = datagrid(cm.dep.0 = f.phq(10)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(phq = 15, 
          comparisons(m, newdata = datagrid(cm.dep.0 = f.phq(15)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")])) %>% 
    as.data.frame() %>% {.$study = study;.}
}


calcRRGad = function(dat, study = NA) {
  
  require(marginaleffects)
  within(dat, {grp = group}) -> dat
  glm(dep.status.mod ~ grp * scale(cm.anx.0), 
      data = dat, family = binomial()) -> m
  rbind(c(gad = 5, 
          comparisons(m, newdata = datagrid(cm.anx.0 = f.gad(5)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(gad = 10, 
          comparisons(m, newdata = datagrid(cm.anx.0 = f.gad(10)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(gad = 15, 
          comparisons(m, newdata = datagrid(cm.anx.0 = f.gad(15)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")])) %>% 
    as.data.frame() %>% {.$study = study;.}
}

calcRRGad.firth = function(dat, study = NA) {
  
  require(marginaleffects)
  within(dat, {grp = group}) -> dat
  logistf(dep.status.mod ~ grp * scale(cm.anx.0), 
          data = dat, control=logistf.control(maxit=100)) -> m
  rbind(c(gad = 5, 
          comparisons(m, newdata = datagrid(cm.anx.0 = f.gad(5)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(gad = 10, 
          comparisons(m, newdata = datagrid(cm.anx.0 = f.gad(10)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(gad = 15, 
          comparisons(m, newdata = datagrid(cm.anx.0 = f.gad(15)),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")])) %>% 
    as.data.frame() %>% {.$study = study;.}
}

calcRRPrevpsychoth.firth = function(dat, study = NA) {
  
  require(marginaleffects)
  within(dat, {grp = group; prevpsychoth = as.numeric(prevpsychoth!=0)}) -> dat
  logistf(dep.status.mod ~ grp * prevpsychoth, 
          data = dat, control=logistf.control(maxit=100)) -> m
  rbind(c(prevpsychoth = 1, 
          comparisons(m, newdata = datagrid(prevpsychoth = 1),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")]),
        c(prevpsychoth = 0, 
          comparisons(m, newdata = datagrid(prevpsychoth = 0),
                      variables = list(grp = 1:0),
                      comparison = "lnratio")[c("estimate", "std.error")])) %>% 
    as.data.frame() %>% {.$study = study;.}
}


fitRr2Stage = function(x, knha = FALSE){
  
  metagen(TE = yi, seTE = sqrt(vi), hakn = knha,
          method.tau = "REML", data = x, 
          common = FALSE, sm="RR", control = list(stepadj=0.5, maxit = 1000)) %>%  #increased max it and reduced steps because of convergence problem
    {data.frame(k = .$k,
                rr = exp(.$TE.random),
                rr.lo = exp(.$lower.random),
                rr.hi = exp(.$upper.random),
                pred.lo = exp(.$lower.predict),
                pred.hi = exp(.$upper.predict),
                tau2 = .$tau2, i2 = .$I2,
                logrr = .$TE.random,
                se_logrr = .$seTE.random,
                w.s = sum(.$w.common),
                w.ssq = sum(.$w.common^2))} -> res
  res$v.tilde = ((res$k-1)*res$w.s)/(res$w.s^2-res$w.ssq)
  return(res)
  
}


matcher = function(a1, a2, a3, data = NULL, ...) {
  
  mC = match.call()
  has.data = !is.null(data)
  if (has.data && !typeof(data) %in% c("data.frame", "list"))
    stop("'data' must be a data.frame.", call. = FALSE)
  dataEnv = as.environment(as.list(data))
  
  a1 = try(eval(mC[["a1"]], if (has.data) dataEnv else parent.frame()))
  a2 = try(eval(mC[["a2"]], if (has.data) dataEnv else parent.frame()))
  a3 = try(eval(mC[["a3"]], if (has.data) dataEnv else parent.frame()))
  return(list(a1, a2, a3))
  
}



# blipModel --------------------------------------------------------------------

require(R6)

blipModel = R6Class(
  "blipModel",
  public = list(
    predictors = NULL, 
    model.list = NULL,
    initialize = function(model = NA, predictors = NULL) {
      self$model.list = model
      self$predictors = predictors
    },
    predict = function(newdat = NULL) {
      require(marginaleffects)
      require(mice)
      require(purrr)
      
      if (is.null(newdat)){
        
        purrr::map_dfr(self$model.list, function(x){
          marginaleffects::avg_predictions(x)[c("estimate", "std.error")]
        }) %>% 
          {mice::pool.scalar(Q = .$estimate, U = .$std.error)} %>% 
          {data.frame(yi = .[["qbar"]], vi = .[["t"]])} -> res
        return(res)
        
      } else {
        
        if (!is.null(self$predictors)){
          newdat = newdat[names(newdat) %in% rownames(self$predictors)[-1]]
          if (length(newdat)[1] == 0) {
            message("Predictor(s) not found in model: returning marginal effect.")
            purrr::map_dfr(self$model.list, function(x){
              marginaleffects::avg_predictions(x)[c("estimate", "std.error")]
            }) %>% 
              {mice::pool.scalar(Q = .$estimate, U = .$std.error)} %>% 
              {data.frame(yi = .[["qbar"]], vi = .[["t"]])} -> res
            return(res)
          }
        }
        purrr::map_dfr(self$model.list, function(x){
          newdat$model = x
          data.grid = do.call(datagridcf, newdat)
          marginaleffects::avg_predictions(
            x, newdata = data.grid)[c("estimate", "std.error")]
        }) %>% 
          {mice::pool.scalar(Q = .$estimate, U = .$std.error)} %>% 
          {data.frame(yi = .[["qbar"]], vi = .[["t"]])} -> res
        return(res)
        
      }
    },
    getCoefficients = function(x) {
      return(self$predictors)
    }
  )
)


# getTestEstimates -------------------------------------------------------------

getTestEstimates = function(m) {
  if (!identical(class(m), c("mitml.result", "list"))) {
    stop("class must be 'mitml.list'")
  }
  testEstimates(m) %>%
    {.$estimates} %>% return()
}



# loadModels -------------------------------------------------------------------

loadModels = function(base.dir) {
  
  require(stringr)
  reslist = list()
  dirs = list.dirs(base.dir)[-1]
  
  for (i in 1:length(dirs)) {
    models = list.files(dirs[i])
    files = paste0(dirs[i], "/", models)
    for (j in 1:length(models)) {
      models[j] = stringr::str_remove_all(models[j], ".rda")
      message("Loading ", models[j])
      load(files[j])
      expr = paste0("reslist$", models[j], "<-", models[j])
      eval(parse(text = expr))
    }
  }
  return(reslist)
}



# predictMarginalEffect --------------------------------------------------------

predictMarginalEffect = function(model.list, params = NULL, 
                                 method = "REML", test = "knha",
                                 rho = 0.5){
  
  if (!is.null(params) && class(params)[1] != "list")
    stop("params must be either NULL or list.")
  
  require(purrr)
  require(metafor)
  require(dplyr)
  
  # Get study data
  model.list %>% 
    map_dfr(function(x) {x$predict(newdat = params)}) %>% 
    {.$study = names(model.list);.} %>% 
    dplyr::mutate(ci.lb = yi + qnorm(.975)*sqrt(vi),
                  ci.ub = yi - qnorm(.975)*sqrt(vi)) %>% 
    dplyr::select(study, yi, vi, ci.lb, ci.ub) -> dat.ma
  dat.ma$clus = strsplit(dat.ma$study, "\\.") %>% map_chr(~.[3])
  
  # Aggregate es data & run meta-analysis
  metafor::escalc(slab = study, yi = yi, vi = vi, data = dat.ma) %>% 
    metafor::aggregate.escalc(cluster = clus, rho = rho) %>% 
    metafor::rma(yi = yi, vi = vi, data = ., test = test, 
                 method = method) -> m.ma
  
  return(list(data = dat.ma, model = m.ma))
  
}



# transformer ------------------------------------------------------------------

transformer = R6Class(
  "transformer",
  public = list(
    hamd2bdi = NULL,
    bdi2hamd = NULL,
    bdi2cm = NULL,
    cm2bdi = NULL,
    phq2cm = NULL,
    cm2phq = NULL,
    initialize = function(hamd2bdi = NULL,
                          bdi2hamd = NULL,
                          bdi2cm = NULL,
                          cm2bdi = NULL,
                          phq2cm = NULL,
                          cm2phq = NULL) {
      self$hamd2bdi = hamd2bdi
      self$bdi2hamd = bdi2hamd
      self$bdi2cm = bdi2cm
      self$cm2bdi = cm2bdi
      self$phq2cm = phq2cm
      self$cm2phq = cm2phq
    },
    hamd2phq = function(v) {
      self$cm2phq(
        self$bdi2cm(
          self$hamd2bdi(v)
        )
      )
    },
    phq2hamd = function(v) {
      self$bdi2hamd(
        self$cm2bdi(
          self$phq2cm(v)
        )
      )
    },
    bdi2phq = function(v) {
      self$cm2phq(
        self$bdi2cm(v) 
      )
    },
    phq2bdi = function(v) {
      self$cm2bdi(
        self$phq2cm(v) 
      )
    }
  )
)


# twoStageModel ----------------------------------------------------------------


twoStageModel = function(es, se, studlab, samples = 25000, 
                         tau.prior = "half-cauchy", s=0.5) {
  
  
  if (length(studlab) != length(es)) {
    stop("'studlab' and 'es' do not match.")
  }
  
  data = list(k = length(es), es = es, se = se)
  
  
  # Define model formulas
  M.hc = paste0("
  model{
  
  for(i in 1:k) { 
    delta[i] ~ dnorm(mu, inv.tau.sq) 
    es[i] ~ dnorm(delta[i], prec.sigma[i])
    prec.sigma[i] <- inverse(pow(se[i],2))
  }
  
  mu ~ dnorm(0.0,1.0E-6)
  tau ~ dt(0, pow(",s,",-2), 1)T(0,)
  inv.tau.sq <- inverse(pow(tau,2))
  }
  
  #monitor# delta, mu, tau
  ")
  
  # Define model formulas
  M.unif = "
  model{
  
  for(i in 1:k) { 
    delta[i] ~ dnorm(mu, inv.tau.sq) 
    es[i] ~ dnorm(delta[i], prec.sigma[i])
    prec.sigma[i] <- inverse(pow(se[i],2))
  }
  
  mu ~ dnorm(0.0,1.0E-6)
  tau ~ dunif(0,10)
  inv.tau.sq <- inverse(pow(tau,2))
  }
  
  #monitor# delta, mu, tau
  "
  
  if (tau.prior[1] == "half-cauchy") { M <- M.hc } else { M <- M.unif }
  
  # MCMC Simulation
  run.jags(M,
           data = data,
           n.chains = 4,
           burnin = 1000,
           summarise = FALSE,
           sample = samples) -> fit
  
  return(fit)
  
}


extract2Stage = function(fits) {
  
  require(coda)
  require(purrr)
  
  message("Extracting ...")
  
  # Get samples
  if (class(fits)[1]=="runjags") {
    coda::as.mcmc(fits) -> mcmc
  } else {
    fits %>% map(~coda::as.mcmc(.)) %>% do.call(rbind,.) -> mcmc 
  }
  
  # Extract samples and pivot to long
  mcmc[,c("mu", "tau")] %>%
    as_tibble() %>%
    pivot_longer(everything()) %>%
    arrange(name) %>%
    group_by(name) %>%
    mutate(median = median(value),) %>%
    ungroup() %>%
    mutate(name = fct_reorder(name, median, .desc = TRUE)) -> mcmc.dat
  
  # Get point estimates and HDIs
  mcmc.dat %>%
    group_by(name) %>%
    summarise(lo = hdi(value)[1],
              md = median(value),
              hi = hdi(value)[2],
              sd = sd(value)) %>% 
    pivot_wider(names_from = name, values_from = everything()) %>% 
    {
      .$pi_lo = .$md_mu-(qt(.975,(colnames(mcmc) %>% {length(.)-2})-2)*sqrt(.$sd_mu^2 + as.numeric(.$md_tau)^2));
      .$pi_hi = .$md_mu+(qt(.975,(colnames(mcmc) %>% {length(.)-2})-2)*sqrt(.$sd_mu^2 + as.numeric(.$md_tau)^2));.
    } %>% 
    mutate(effect = paste0(sprintf("%.2f", md_mu), " [", 
                           sprintf("%.2f", lo_mu), "; ", 
                           sprintf("%.2f", hi_mu), "]"),
           pi = paste0("[", sprintf("%.2f", pi_lo), "; ", 
                       sprintf("%.2f", pi_hi), "]"),
           tau = paste0(sprintf("%.2f", md_tau), " [", 
                        sprintf("%.2f", lo_tau), "; ", 
                        sprintf("%.2f", hi_tau), "]"),
           k = colnames(mcmc) %>% {length(.)-2}) %>% 
    select(k, effect, pi, tau, effect.se = sd_mu) -> res
  
  return(res)
}



# studycounter ------------------------------------------------------------------

studycounter = R6Class(
  "studycounter",
  public = list(
    filters = NULL,
    vars.matx = NULL,
    initialize = function(filters = NULL,
                          vars.matrix = NULL) {
      self$filters = filters
      self$vars.matrix = vars.matrix
    },
    filter = function(newfilters = NULL) {
      
      require(dplyr)
      
      filters = self$filters
      if (!is.null(newfilters)[1]) {
        filters[names(newfilters)] = newfilters
      }
      
      self$vars.matrix %>% 
        dplyr::filter(dep %in% filters$dep) %>% 
        dplyr::filter(age %in% filters$age) %>% 
        dplyr::filter(sex %in% filters$sex | sex.all.0 %in% filters$sexfemale) %>% 
        dplyr::filter(employment %in% filters$employment | employ.all.1 %in% filters$employmentyes) %>% 
        dplyr::filter(degree %in% filters$degree) %>% 
        dplyr::filter(chronmed %in% filters$chronmed | chron.med.all.1 %in% filters$chronmedyes) %>% 
        dplyr::filter(depmed %in% filters$depmed) %>% 
        dplyr::filter(intervention %in% filters$intervention) %>% 
        dplyr::filter(control %in% filters$control) %>% 
        dplyr::filter(timecat %in% filters$timecat) -> res
      
      return(list(nstudies = length(unique(res$study)),
                  models = res$model,
                  studies = unique(res$study),
                  filter = filters))
    }
  )
)


make.components = function(data, comp.name, 
                           study.name = "study",
                           arm.name = "arm") {
  require(tidyr)
  data[,comp.name] = as.integer(data[,comp.name] == 1)
  data[,c(study.name, arm.name, comp.name)] -> df
  tidyr::spread(df, arm.name, comp.name) -> res
  return(as.matrix(res[,-1]))
}



all.interactions = function(combination, Nc, components) {
  
  # Helper function to calculate position
  place = function(i, j) {
    if (i > j) {
      temp = i
      i = j
      j = temp
    }
    return((2 * Nc - i) * (i - 1) / 2 + j - i)
  }
  
  # Generate interaction indices
  combination = sort(combination)
  nc = length(combination)
  interactions = c()
  labels = c()
  
  for (i in 1:(nc - 1)) {
    for (j in (i + 1):nc) {
      idx = place(combination[i], combination[j])
      label = paste(components[combination[i]], "-", components[combination[j]])
      interactions = c(interactions, idx)
      labels = c(labels, label)
    }
  }
  
  return(data.frame(index = interactions, label = labels))
}

calc.interaction.number <- function(Nc) {
  return(Nc * (Nc - 1) / 2)
}

make.interactions <- function(Ns, Nc, Ninter, na, component_matrices) {
  interactions <- array(NA, dim = c(Ns, 4, Ninter))
  for (i in 1:Ns) {
    for (j in 1:na[i]) {
      place1 <- 0
      for (k in 1:(Nc - 1)) {
        for (l in (k + 1):Nc) {
          place1 <- place1 + 1
          interactions[i, j, place1] <- 
            component_matrices[[k]][i, j] * component_matrices[[l]][i, j]
        }
      }
    }
  }
  return(interactions)
}





