# ==============================================================================
#  R for Psychological Science - custom functions
#
#  Every function defined across the workshops, collected in one place.
#
#  To use these, save this file next to your own script and run:
#
#      source("r4psych-functions.R")
#
#  Each function needs certain packages loaded first. Between them they use
#  dplyr, magrittr, ggplot2 (all part of the tidyverse) and psych:
#
#      library(tidyverse)
#      library(magrittr)
#      library(psych)
#
#  A line-by-line explanation of each function is in the Extras chapter of the
#  book. Nothing here is magic - it is all built from functions the workshops
#  already taught.
# ==============================================================================

# ----------------------------------------------------------------------------
# var.center()  --  Mean-centre a variable
# Introduced in Workshop 2
# ----------------------------------------------------------------------------

var.center <- function(x) {
    scale(x, scale = FALSE)
}

# ----------------------------------------------------------------------------
# gen_comp()  --  Build a composite (mean) score from a set of items
# Introduced in Workshop 2
# ----------------------------------------------------------------------------

gen_comp <- function(data, comp, vector){
   comp <- enquo(comp)
   data %>% 
       rowwise() %>% 
       mutate(!!quo_name(comp) := mean(c(!!!vector), na.rm = TRUE)) %>% 
       ungroup()
}

# ----------------------------------------------------------------------------
# MO_Detection()  --  Find and remove multivariate outliers using Mahalanobis distance
# Introduced in Workshop 5
# ----------------------------------------------------------------------------

MO_Detection = function (CompleteDataset, AnalyzedDataset, alpha = 0.001) {
  Means = colMeans(AnalyzedDataset, na.rm = T)
  Covariance = cov(AnalyzedDataset, use = "pairwise.complete.obs")
  Distances = mahalanobis(AnalyzedDataset, Means, Covariance)
  cutoff = qchisq(1-alpha, ncol(AnalyzedDataset))
  remain = Distances < cutoff
  PlotData = AnalyzedDataset %>% 
    mutate(ID = 1:nrow(AnalyzedDataset), 
           Distance = Distances, 
           # state both levels explicitly, so this still works when every case
           # falls on the same side of the cutoff
           Outlier = factor(remain, levels = c(FALSE, TRUE), labels = c("Yes","No")))
  n = nrow(AnalyzedDataset) - sum(remain)
  report = paste(n, "cases were multivariate outliers")
  print(report)
  p = ggplot(PlotData, aes(ID, Distance, color = Outlier)) + 
    geom_point() + 
    labs(title = "Multivariate Outliers",
         subtitle = report) +
    theme(legend.position = "bottom")
  print(p)

  return(CompleteDataset %>% filter(remain))
}

# ----------------------------------------------------------------------------
# alpha_table()  --  A readable Cronbach's alpha table
# Introduced in Workshop 5
# ----------------------------------------------------------------------------

alpha_table <- function(data, digits = 3) {
  
  a <- psych::alpha(data)   # everything we need is already in here
  
  data.frame(
    Item = rownames(a$item.stats),
    obs  = a$item.stats$n,
    sign = ifelse(a$item.stats$raw.r >= 0, "+", "-"),
    
    # how each item relates to the scale
    `item-test correlation` = round(a$item.stats$r.cor,  digits),
    `item-rest correlation` = round(a$item.stats$r.drop, digits),
    
    # what happens to the scale if we drop that item
    `avg inter-item correlation` = round(a$alpha.drop$average_r,  digits),
    `alpha if dropped`           = round(a$alpha.drop$raw_alpha, digits),
    
    check.names = FALSE, row.names = NULL
  )
}

# ----------------------------------------------------------------------------
# cor_table()  --  An APA-style correlation table with descriptives and significance stars
# Introduced in Workshop 7
# ----------------------------------------------------------------------------

cor_table <- function(data, vars, describe = TRUE, digits = 2) {
  
  d  <- data[, vars, drop = FALSE]          # keep only the variables we asked for
  d  <- d[complete.cases(d), , drop = FALSE] # drop rows with missing values
  ct <- psych::corr.test(d)                  # correlations + p-values in one go
  
  # turn p-values into the stars you see in journal articles
  stars <- ifelse(ct$p < .001, "***",
           ifelse(ct$p < .01,  "**",
           ifelse(ct$p < .05,  "*", "")))
  
  cells <- matrix(paste0(formatC(ct$r, digits = digits, format = "f"), stars),
                  nrow = nrow(ct$r), dimnames = dimnames(ct$r))
  
  cells[upper.tri(cells, diag = TRUE)] <- ""   # keep only the lower triangle
  cells <- cells[, -ncol(cells), drop = FALSE] # the last column is now empty
  colnames(cells) <- seq_len(ncol(cells))
  
  out <- data.frame(Variable = paste0(seq_along(vars), ". ", vars),
                    cells, check.names = FALSE)
  
  if (describe) {
    out <- cbind(out[1],
                 M  = formatC(sapply(d, mean), digits = digits, format = "f"),
                 SD = formatC(sapply(d, sd),   digits = digits, format = "f"),
                 out[-1])
  }
  
  row.names(out) <- NULL
  out
}

# ----------------------------------------------------------------------------
# partial_cor_table()  --  The same, for partial correlations
# Introduced in Workshop 7
# ----------------------------------------------------------------------------

partial_cor_table <- function(data, vars, control.vars, digits = 2) {
  
  d  <- data[, c(vars, control.vars), drop = FALSE]
  d  <- d[complete.cases(d), , drop = FALSE]
  
  pr <- psych::partial.r(d, vars, control.vars)   # the partial correlations
  n  <- nrow(d) - length(control.vars)            # we "spend" a df per control variable
  p  <- psych::corr.p(pr, n = n)$p                # p-values for those correlations
  
  stars <- ifelse(p < .001, "***", ifelse(p < .01, "**", ifelse(p < .05, "*", "")))
  
  cells <- matrix(paste0(formatC(pr, digits = digits, format = "f"), stars),
                  nrow = nrow(pr), dimnames = dimnames(pr))
  
  cells[upper.tri(cells, diag = TRUE)] <- ""
  cells <- cells[, -ncol(cells), drop = FALSE]
  colnames(cells) <- seq_len(ncol(cells))
  
  data.frame(Variable = paste0(seq_along(vars), ". ", vars),
             cells, check.names = FALSE, row.names = NULL)
}