# ========================================================
# MISSING DATA AND MULTIPLE IMPUTATION IN R
# Canadian Social Connection Survey (CSCS), 2021 wave
# ========================================================
# Question: is social support associated with loneliness,
# after adjusting for age, gender and household income?
# Many people left at least one of these questions blank.
# The script describes the missing values, then compares
# complete-case analysis, mean imputation and multiple
# imputation.
# ========================================================


# ========================================================
# SECTION 1: PACKAGES AND DATA
# ========================================================

# install.packages("mice")   # run once
library(mice)   # multiple imputation by chained equations

github_url <- "https://raw.githubusercontent.com/jorgeandr3s/heal/main/cscs/public_data/CSCS2025_full_cleaned_deidentified_data_and_metadata.RData"
load(url(github_url))

# Keep the 2021 wave, which gives one row per person
data <- data[data$SURVEY_collection_year == 2021, ]
dim(data)


# ========================================================
# SECTION 2: PREPARE THE VARIABLES
# ========================================================

# --- 2.1 Gender: a non-answer becomes missing (NA)
data$gender <- NA
data$gender[data$DEMO_gender == "Man"] <- "Man"
data$gender[data$DEMO_gender == "Woman"] <- "Woman"
data$gender[data$DEMO_gender == "Non-binary"] <- "Non-binary"
data$gender <- factor(data$gender,
                      levels = c("Man", "Woman", "Non-binary"))
table(data$DEMO_gender, data$gender, useNA = "ifany")

# --- 2.2 Household income: 26 brackets into 3 groups
# as.integer() gives each bracket's position in the list of
# levels: 1 = "Under $5,000" up to 26 = "$200,000 or more".
# Level 27, "Presented but no response", stays missing.
income_code <- as.integer(data$DEMO_household_income)
data$income <- NA
data$income[income_code <= 8] <- "Under $40,000"
data$income[income_code >= 9 & income_code <= 13] <- "$40,000 to $79,999"
data$income[income_code >= 14 & income_code <= 26] <- "$80,000 or more"
data$income <- factor(data$income,
                      levels = c("Under $40,000", "$40,000 to $79,999",
                                 "$80,000 or more"))
table(data$DEMO_household_income == "Presented but no response",
      useNA = "ifany")
table(data$income, useNA = "ifany")

# --- 2.3 One data frame for the analysis
# loneliness: UCLA 3-item Loneliness Scale, 3 to 9 (outcome)
# support:    Multidimensional Scale of Perceived Social
#             Support, 1 to 7 (exposure)
# age, gender, income: adjustment variables
# depression: PHQ-2 depressive symptoms, 0 to 6 (auxiliary)
# finished:   TRUE if the person reached the end of the
#             survey (auxiliary)
mi_data <- data.frame(
  loneliness = data$LONELY_ucla_loneliness_scale_score,
  support    = data$PSYCH_zimet_multidimensional_social_support_scale_score,
  age        = data$DEMO_age,
  gender     = data$gender,
  income     = data$income,
  depression = data$WELLNESS_phq_score,
  finished   = data$SURVEY_finished)
summary(mi_data)


# ========================================================
# SECTION 3: HOW MUCH IS MISSING, AND WHERE?
# ========================================================

# --- 3.1 Missing values in each variable
colSums(is.na(mi_data))
round(100 * colMeans(is.na(mi_data)), 1)

barplot(100 * colMeans(is.na(mi_data)),
        main = "Missing values in each variable",
        ylab = "Percent missing", ylim = c(0, 30),
        col = "steelblue", las = 2, cex.names = 0.8)

# --- 3.2 Complete cases: rows with no missing value
sum(complete.cases(mi_data))
round(100 * mean(complete.cases(mi_data)), 1)

# --- 3.3 Missing-data patterns (1 = observed, 0 = missing)
md.pattern(mi_data, plot = FALSE)


# ========================================================
# SECTION 4: WHO IS MISSING?
# ========================================================

# --- 4.1 Did people who stopped early leave more gaps?
table(finished = mi_data$finished,
      support_missing = is.na(mi_data$support))

# --- 4.2 Among people who finished, who skipped income?
finishers <- mi_data[mi_data$finished == "TRUE", ]
income_missing <- is.na(finishers$income)
table(income_missing)
round(100 * prop.table(table(finishers$gender, income_missing),
                       margin = 1), 1)
round(tapply(finishers$age, income_missing, mean), 1)
round(tapply(finishers$loneliness, income_missing, mean,
             na.rm = TRUE), 2)

# --- 4.3 An analysis sample: the people with an outcome
# Many analyses keep only the people whose outcome was
# measured. A report then compares them with the people
# who were left out.
has_outcome <- !is.na(mi_data$loneliness)
table(has_outcome)
outcome_data <- mi_data[has_outcome, ]
dim(outcome_data)

# Variables with no gaps, or almost none
round(tapply(mi_data$age, has_outcome, mean), 1)
round(100 * prop.table(table(mi_data$gender, has_outcome),
                       margin = 2), 1)
round(100 * prop.table(table(mi_data$finished, has_outcome),
                       margin = 2), 1)

# The exposure and the other incomplete variables
round(tapply(mi_data$support, has_outcome, mean, na.rm = TRUE), 2)
round(tapply(mi_data$depression, has_outcome, mean,
             na.rm = TRUE), 2)
round(100 * prop.table(table(mi_data$income, has_outcome),
                       margin = 2), 1)

# Percent missing in each variable, for each group
round(100 * colMeans(is.na(outcome_data)), 1)
round(100 * colMeans(is.na(mi_data[!has_outcome, ])), 1)


# ========================================================
# SECTION 5: COMPLETE-CASE ANALYSIS
# ========================================================
# lm() drops every row with a missing value in any of the
# variables in the model, and reports how many it dropped.
cc_model <- lm(loneliness ~ support + age + gender + income,
               data = mi_data)
summary(cc_model)
nobs(cc_model)


# ========================================================
# SECTION 6: SINGLE IMPUTATION, AND WHY IT MISLEADS
# ========================================================

# --- 6.1 Fill each gap with one value: the mean for scores,
# and the most common category (the mode) for categories
mean_data <- mi_data
mean_data$loneliness[is.na(mean_data$loneliness)] <-
  mean(mi_data$loneliness, na.rm = TRUE)
mean_data$support[is.na(mean_data$support)] <-
  mean(mi_data$support, na.rm = TRUE)
mean_data$gender[is.na(mean_data$gender)] <- "Woman"
mean_data$income[is.na(mean_data$income)] <- "Under $40,000"

# --- 6.2 What single imputation does to the data
sd(mi_data$support, na.rm = TRUE)
sd(mean_data$support)
cor(mi_data$support, mi_data$loneliness, use = "complete.obs")
cor(mean_data$support, mean_data$loneliness)
table(mean_data$income)
hist(mean_data$support, breaks = 24,
     main = "Social support after mean imputation",
     xlab = "Social support (1 to 7)", col = "grey80")

# --- 6.3 The same model on the filled-in data
mean_model <- lm(loneliness ~ support + age + gender + income,
                 data = mean_data)
nobs(mean_model)
round(summary(mean_model)$coefficients["support", ], 4)
round(summary(cc_model)$coefficients["support", ], 4)


# ========================================================
# SECTION 7: MULTIPLE IMPUTATION WITH mice()
# ========================================================

# --- 7.1 A dry run shows the default settings
dry_run <- mice(mi_data, maxit = 0)
dry_run$method
dry_run$predictorMatrix

# --- 7.2 Create 40 imputed datasets
# m = 40 datasets, maxit = 10 rounds of the algorithm,
# seed makes the random draws repeatable
imp <- mice(mi_data, m = 40, maxit = 10, seed = 2021,
            printFlag = FALSE)
imp

# --- 7.3 Look at one completed dataset
completed_1 <- complete(imp, 1)
colSums(is.na(completed_1))

# --- 7.4 Do the imputed values look plausible?
observed_support <- mi_data$support[!is.na(mi_data$support)]
imputed_support <- completed_1$support[is.na(mi_data$support)]
round(c(observed = mean(observed_support),
        imputed = mean(imputed_support)), 2)
plot(density(observed_support, bw = 0.3), lwd = 3,
     main = "Observed and imputed support (dataset 1)",
     xlab = "Social support (1 to 7)", ylim = c(0, 0.42))
lines(density(imputed_support, bw = 0.3), lwd = 3, lty = 2,
      col = "firebrick")
legend("topleft", legend = c("Observed", "Imputed"),
       col = c("black", "firebrick"), lwd = 3, lty = c(1, 2))

# --- 7.5 Did the algorithm settle down? (convergence)
plot(imp, c("loneliness", "support"), layout = c(2, 2))


# ========================================================
# SECTION 8: ANALYZE EACH DATASET AND POOL THE RESULTS
# ========================================================

# --- 8.1 Fit the model in each of the 40 datasets
mi_fits <- with(imp, lm(loneliness ~ support + age + gender + income))

# --- 8.2 Pool the 40 results with Rubin's rules
mi_pooled <- pool(mi_fits)
mi_summary <- summary(mi_pooled, conf.int = TRUE)
print(mi_summary[, c("term", "estimate", "std.error", "p.value",
                     "conf.low", "conf.high")], digits = 3)

# --- 8.3 Rubin's rules by hand for the support coefficient
# ubar = average within-dataset variance (squared SE)
# b    = variance of the 40 estimates between datasets
# t    = total variance = ubar + (1 + 1/m) * b
support_row <- mi_pooled$pooled[mi_pooled$pooled$term == "support", ]
support_row[, c("estimate", "ubar", "b", "t", "fmi")]
support_row$ubar + (1 + 1/40) * support_row$b
sqrt(support_row$t)    # the pooled standard error
sqrt(support_row$ubar) # the standard error if imputed values were known


# ========================================================
# SECTION 9: COMPARE THE THREE APPROACHES
# ========================================================
comparison <- data.frame(
  method   = c("Complete cases", "Mean imputation",
               "Multiple imputation"),
  n        = c(nobs(cc_model), nobs(mean_model), nrow(mi_data)),
  estimate = c(coef(cc_model)["support"], coef(mean_model)["support"],
               mi_summary$estimate[2]),
  std_error = c(summary(cc_model)$coefficients["support", 2],
                summary(mean_model)$coefficients["support", 2],
                mi_summary$std.error[2]))
print(comparison, digits = 3)