# ========================================================
# DATA VISUALIZATION IN BASE R AND GGPLOT2
# Canadian Social Connection Survey (CSCS), 2021 wave
# ========================================================
# Question: what do loneliness, social support and mental
# health look like in the 2021 CSCS data, and how do they
# differ by age group and gender? The script looks at the
# data with plots before any model is fitted.
# ========================================================


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

# install.packages(c("ggplot2", "ggExtra", "patchwork"))   # run once
library(ggplot2)     # the grammar of graphics
library(ggExtra)     # ggMarginal() adds marginal histograms
library(patchwork)   # combines several ggplots into one figure

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 Six numeric variables
viz <- data.frame(
  loneliness = data$LONELY_ucla_loneliness_scale_score,  # 3 to 9
  support    = data$PSYCH_zimet_multidimensional_social_support_scale_score,  # 1 to 7
  depression = data$WELLNESS_phq_score,                  # PHQ-2, 0 to 6
  anxiety    = data$WELLNESS_gad_score,                  # GAD-2, 0 to 6
  age        = data$DEMO_age,
  household  = data$GEO_housing_household_size)          # other people at home
summary(viz)

# --- 2.2 Gender: treat a non-answer as missing
viz$gender <- NA
viz$gender[data$DEMO_gender == "Man"] <- "Man"
viz$gender[data$DEMO_gender == "Woman"] <- "Woman"
viz$gender[data$DEMO_gender == "Non-binary"] <- "Non-binary"
viz$gender <- factor(viz$gender, levels = c("Man", "Woman", "Non-binary"))

# --- 2.3 Three age groups
viz$age_group <- NA
viz$age_group[viz$age < 30] <- "16 to 29"
viz$age_group[viz$age >= 30 & viz$age < 50] <- "30 to 49"
viz$age_group[viz$age >= 50] <- "50 and over"
viz$age_group <- factor(viz$age_group)

# --- 2.4 Self-rated mental health, in its natural order
# Answers that are not one of the five levels become NA
viz$mental_health <- factor(data$WELLNESS_self_rated_mental_health,
  levels = c("Poor", "Fair", "Good", "Very good", "Excellent"))

# --- 2.5 Keep the people with no missing values, so that
# every plot shows the same people
viz <- na.omit(viz)
nrow(viz)


# ========================================================
# SECTION 3: QUICK LOOKS IN BASE R
# ========================================================

# --- 3.1 One numeric variable: a histogram
hist(viz$loneliness)

# One bar for each possible score, from 3 to 9
hist(viz$loneliness, breaks = seq(2.5, 9.5, by = 1),
     main = "UCLA loneliness score", xlab = "Score (3 to 9)",
     col = "grey80")

# --- 3.2 A problem that summary() did not reveal
hist(viz$household, breaks = seq(-0.5, 20.5, by = 1),
     main = "Other people in the household",
     xlab = "Number of other people", col = "grey80")
table(viz$household)

# --- 3.3 One categorical variable: a bar chart of counts
mh_counts <- table(viz$mental_health)
mh_counts
barplot(mh_counts, main = "Self-rated mental health",
        ylab = "Number of people", cex.names = 0.85)

# --- 3.4 A numeric variable by group: boxplots
boxplot(loneliness ~ age_group, data = viz,
        main = "Loneliness by age group",
        xlab = "Age group", ylab = "UCLA loneliness score")

# --- 3.5 Two numeric variables: a scatterplot
plot(viz$support, viz$loneliness)

# jitter() adds a little random noise so that people with
# the same scores no longer sit on exactly the same spot
set.seed(2021)    # makes the random noise repeatable
plot(jitter(viz$support), jitter(viz$loneliness),
     pch = 16, col = rgb(0, 0, 0, 0.15),
     xlab = "Social support (1 to 7)",
     ylab = "UCLA loneliness score (3 to 9)")

# --- 3.6 Every pair of four variables at once
pairs(viz[, c("loneliness", "support", "depression", "anxiety")])

# --- 3.7 Two base R plots side by side, on the same y-axis
par(mfrow = c(1, 2))   # 1 row and 2 columns
hist(viz$depression, breaks = seq(-0.5, 6.5, by = 1), ylim = c(0, 750),
     main = "Depression (PHQ-2)", xlab = "Score (0 to 6)", col = "grey80")
hist(viz$anxiety, breaks = seq(-0.5, 6.5, by = 1), ylim = c(0, 750),
     main = "Anxiety (GAD-2)", xlab = "Score (0 to 6)", col = "grey80")
par(mfrow = c(1, 1))   # back to one plot per window


# ========================================================
# SECTION 4: THE GRAMMAR OF GRAPHICS WITH GGPLOT2
# ========================================================

# --- 4.1 Data and aesthetics, then a geom
ggplot(viz, aes(x = loneliness))

p_hist <- ggplot(viz, aes(x = loneliness)) +
  geom_histogram(binwidth = 1, fill = "grey60", colour = "white")
p_hist

# --- 4.2 geom_bar() counts the rows for you
ggplot(viz, aes(x = mental_health)) +
  geom_bar()

# geom_col() draws values that you have already calculated
mean_lonely <- aggregate(loneliness ~ age_group, data = viz, FUN = mean)
mean_lonely
ggplot(mean_lonely, aes(x = age_group, y = loneliness)) +
  geom_col()

# --- 4.3 A boxplot with every person shown as a point
p_box <- ggplot(viz, aes(x = age_group, y = loneliness)) +
  geom_boxplot(outlier.shape = NA) +
  geom_jitter(width = 0.2, height = 0.2, alpha = 0.1)
p_box

# --- 4.4 A scatterplot with a straight-line smoother
p_scatter <- ggplot(viz, aes(x = support, y = loneliness)) +
  geom_jitter(width = 0.1, height = 0.25, alpha = 0.15) +
  geom_smooth(method = "lm")
p_scatter

# --- 4.5 Facets: one small panel for each group
p_scatter + facet_wrap(~ gender)
p_scatter + facet_grid(gender ~ age_group)

# --- 4.6 Colour, labels and a colour-blind-safe palette
okabe_ito <- c("#E69F00", "#56B4E9", "#009E73")   # Okabe-Ito colours
p_colour <- ggplot(viz, aes(x = support, y = loneliness,
                            colour = age_group)) +
  geom_jitter(width = 0.1, height = 0.25, alpha = 0.25) +
  geom_smooth(method = "lm", se = FALSE) +
  scale_colour_manual(values = okabe_ito) +
  labs(title = "Loneliness and social support, by age group",
       x = "Social support (MSPSS, 1 to 7)",
       y = "Loneliness (UCLA, 3 to 9)",
       colour = "Age group") +
  theme_minimal()
p_colour

# --- 4.7 Save a plot to a file
ggsave("loneliness_support.png", p_colour,
       width = 7, height = 5, dpi = 300)


# ========================================================
# SECTION 5: MANY VARIABLES AND COMBINED FIGURES
# ========================================================

# --- 5.1 A correlation matrix
cor_mat <- cor(viz[, c("loneliness", "support", "depression",
                       "anxiety", "age")])
round(cor_mat, 2)

# --- 5.2 The same matrix as a heatmap
# ggplot2 needs one row per cell: variable 1, variable 2, r
cor_long <- as.data.frame(as.table(cor_mat))
names(cor_long) <- c("var1", "var2", "r")
head(cor_long)
p_heat <- ggplot(cor_long, aes(x = var1, y = var2, fill = r)) +
  geom_tile(colour = "white") +
  geom_text(aes(label = round(r, 2))) +
  scale_fill_gradient2(low = "#2166AC", mid = "white",
                       high = "#B2182B", limits = c(-1, 1)) +
  labs(x = NULL, y = NULL, fill = "r")
p_heat

# --- 5.3 A scatterplot with marginal histograms
ggMarginal(p_scatter, type = "histogram", binwidth = 0.5)

# --- 5.4 Several ggplots in one figure with patchwork
p_hist | p_box
(p_hist | p_box) / p_colour


# ========================================================
# SECTION 6: A PUBLICATION-READY FIGURE
# ========================================================

# --- 6.1 Every axis states the measure and its range
fig_a <- p_hist + labs(x = "Loneliness (UCLA, 3 to 9)",
                       y = "Number of people")
fig_b <- p_box + coord_flip() +   # long group labels read best on the side
  labs(x = "Age group", y = "Loneliness (UCLA, 3 to 9)")
fig_c <- p_colour + labs(title = NULL)

# --- 6.2 Panel letters and one theme for every panel
figure1 <- (fig_a | fig_b) / fig_c +
  plot_annotation(tag_levels = "A") &
  theme_classic(base_size = 9)
figure1

# --- 6.3 Save at the size and resolution a journal expects
ggsave("figure1.png", figure1, width = 7, height = 7, dpi = 300)