# ========================================================
# DATA CLEANING AND DESCRIPTIVE STATISTICS IN BASE R
# Canadian Social Connection Survey (CSCS), public data
# ========================================================
# Purpose: prepare a set of variables for analysis and describe them.
# Every recode uses the same base R pattern:
#   data$new_variable <- NA
#   data$new_variable[condition] <- "Category"
# ========================================================


# ========================================================
# SECTION 1: LOAD THE DATA
# ========================================================

# The .RData file holds R objects saved from an earlier R session.
github_url <- "https://raw.githubusercontent.com/jorgeandr3s/heal/main/cscs/public_data/CSCS2025_full_cleaned_deidentified_data_and_metadata.RData"
load(url(github_url))

# Number of rows (survey responses) and columns (variables)
dim(data)


# ========================================================
# SECTION 2: KEEP THE ROWS AND COLUMNS WE NEED
# ========================================================

# Square brackets select parts of a data frame: data[rows, columns]

# --- 2.1 Keep one survey wave
# The UCLA Loneliness Scale was asked of every participant in 2021.
# Keeping the 2021 wave also gives one row per person.
data <- data[data$SURVEY_collection_year == 2021, ]
dim(data)

# --- 2.2 Keep the variables used in this analysis
keep_vars <- c("LONELY_ucla_loneliness_scale_score",
               "DEMO_age",
               "DEMO_gender",
               "DEMO_relationship_status",
               "GEO_housing_household_size",
               "WELLNESS_self_rated_mental_health")
data <- data[, keep_vars]
dim(data)

# --- 2.3 Look at the structure of the data frame
# str() lists each variable, its type, and its first few values
str(data)

# --- 2.4 Count the missing values in each variable
colSums(is.na(data))


# ========================================================
# SECTION 3: EXPLORE EACH VARIABLE BEFORE CHANGING IT
# ========================================================

# --- 3.1 The outcome: UCLA Loneliness Scale score (3 items, range 3 to 9)
class(data$LONELY_ucla_loneliness_scale_score)
table(data$LONELY_ucla_loneliness_scale_score, useNA = "ifany")
round(prop.table(table(data$LONELY_ucla_loneliness_scale_score)), 3)

summary(data$LONELY_ucla_loneliness_scale_score)
sd(data$LONELY_ucla_loneliness_scale_score, na.rm = TRUE)
IQR(data$LONELY_ucla_loneliness_scale_score, na.rm = TRUE)

# Shape of the distribution (install.packages("moments") once)
library(moments)

# Skewness: 0 = symmetric; negative = longer left tail;
# positive = longer right tail. Values beyond +/-1 suggest marked skew.
skewness(data$LONELY_ucla_loneliness_scale_score, na.rm = TRUE)

# Kurtosis from the moments package: a normal distribution has a value
# of 3. Above 3 = heavier tails; below 3 = lighter tails, flatter shape.
kurtosis(data$LONELY_ucla_loneliness_scale_score, na.rm = TRUE)

# --- 3.2 Age (a continuous variable): check that the range is plausible
class(data$DEMO_age)
summary(data$DEMO_age)

# --- 3.3 Gender, relationship status and mental health (factors)
table(data$DEMO_gender, useNA = "ifany")
table(data$DEMO_relationship_status, useNA = "ifany")
table(data$WELLNESS_self_rated_mental_health, useNA = "ifany")

# --- 3.4 Household size (number of other people in the household)
summary(data$GEO_housing_household_size)
table(data$GEO_housing_household_size, useNA = "ifany")


# ========================================================
# SECTION 4: RECODE AND CREATE NEW VARIABLES
# ========================================================

# --- 4.1 Two categories from a continuous score (cut point of 6)
data$ucla_high_low <- NA
data$ucla_high_low[data$LONELY_ucla_loneliness_scale_score < 6] <- "Low"
data$ucla_high_low[data$LONELY_ucla_loneliness_scale_score >= 6] <- "High"

# Factor levels set the order of the categories in tables and plots
data$ucla_high_low <- factor(data$ucla_high_low, levels = c("Low", "High"))

# Check: every original score should fall in exactly one new category
table(data$LONELY_ucla_loneliness_scale_score, data$ucla_high_low,
      useNA = "ifany")

# --- 4.2 Several categories from a continuous variable (age groups)
data$age_group <- NA
data$age_group[data$DEMO_age >= 16 & data$DEMO_age <= 29] <- "16-29"
data$age_group[data$DEMO_age >= 30 & data$DEMO_age <= 44] <- "30-44"
data$age_group[data$DEMO_age >= 45 & data$DEMO_age <= 64] <- "45-64"
data$age_group[data$DEMO_age >= 65] <- "65+"
data$age_group <- factor(data$age_group,
                         levels = c("16-29", "30-44", "45-64", "65+"))

# Check: the youngest and oldest age in each group
table(data$age_group, useNA = "ifany")
tapply(data$DEMO_age, data$age_group, min)
tapply(data$DEMO_age, data$age_group, max)

# --- 4.3 Treat a non-answer as missing (gender)
data$gender_3cat <- NA
data$gender_3cat[data$DEMO_gender == "Man"] <- "Man"
data$gender_3cat[data$DEMO_gender == "Woman"] <- "Woman"
data$gender_3cat[data$DEMO_gender == "Non-binary"] <- "Non-binary"
data$gender_3cat <- factor(data$gender_3cat,
                           levels = c("Man", "Woman", "Non-binary"))
table(data$DEMO_gender, data$gender_3cat, useNA = "ifany")

# --- 4.4 Collapse five categories into three (self-rated mental health)
# %in% asks whether each value matches any value in a list
data$mental_health_3cat <- NA
data$mental_health_3cat[data$WELLNESS_self_rated_mental_health %in%
                          c("Poor", "Fair")] <- "Poor or fair"
data$mental_health_3cat[data$WELLNESS_self_rated_mental_health ==
                          "Good"] <- "Good"
data$mental_health_3cat[data$WELLNESS_self_rated_mental_health %in%
                          c("Very good", "Excellent")] <- "Very good or excellent"
data$mental_health_3cat <- factor(data$mental_health_3cat,
                                  levels = c("Poor or fair", "Good",
                                             "Very good or excellent"))
table(data$WELLNESS_self_rated_mental_health, data$mental_health_3cat,
      useNA = "ifany")

# --- 4.5 Collapse three categories into two (relationship status)
data$partnered <- NA
data$partnered[data$DEMO_relationship_status %in%
                 c("Single and not dating", "Single and dating")] <- "Single"
data$partnered[data$DEMO_relationship_status ==
                 "In a relationship"] <- "In a relationship"
data$partnered <- factor(data$partnered,
                         levels = c("Single", "In a relationship"))
table(data$DEMO_relationship_status, data$partnered, useNA = "ifany")

# --- 4.6 Two categories from a count (living alone)
# Household size counts the OTHER people in the household, so 0 = alone
data$lives_alone <- NA
data$lives_alone[data$GEO_housing_household_size == 0] <- "Lives alone"
data$lives_alone[data$GEO_housing_household_size >= 1] <- "Lives with others"
data$lives_alone <- factor(data$lives_alone,
                           levels = c("Lives with others", "Lives alone"))
table(data$lives_alone, useNA = "ifany")

# --- 4.7 One variable from two variables (relationship and living)
# Each person gets the one category that matches BOTH of their answers
data$partner_living <- NA
data$partner_living[data$partnered == "In a relationship" &
                      data$lives_alone == "Lives with others"] <- "Partnered, lives with others"
data$partner_living[data$partnered == "In a relationship" &
                      data$lives_alone == "Lives alone"] <- "Partnered, lives alone"
data$partner_living[data$partnered == "Single" &
                      data$lives_alone == "Lives with others"] <- "Single, lives with others"
data$partner_living[data$partnered == "Single" &
                      data$lives_alone == "Lives alone"] <- "Single, lives alone"
data$partner_living <- factor(data$partner_living,
                              levels = c("Partnered, lives with others",
                                         "Partnered, lives alone",
                                         "Single, lives with others",
                                         "Single, lives alone"))

# Check: the four counts should match the cells of the two-way table
table(data$partnered, data$lives_alone)
table(data$partner_living, useNA = "ifany")


# ========================================================
# SECTION 5: DESCRIPTIVE STATISTICS
# ========================================================

# --- 5.1 A continuous variable: centre and spread
mean(data$DEMO_age)
median(data$DEMO_age)
sd(data$DEMO_age)
quantile(data$DEMO_age, probs = c(0.25, 0.50, 0.75))

# --- 5.2 A categorical variable: counts and percentages
table(data$age_group)
round(100 * prop.table(table(data$age_group)), 1)

# --- 5.3 A summary statistic within groups
tapply(data$LONELY_ucla_loneliness_scale_score, data$age_group,
       mean, na.rm = TRUE)
tapply(data$LONELY_ucla_loneliness_scale_score, data$gender_3cat,
       median, na.rm = TRUE)


# ========================================================
# SECTION 6: DESCRIPTIVE PLOTS IN BASE R
# ========================================================
# main = plot title, xlab = x-axis label, ylab = y-axis label

# --- 6.1 Histogram: the distribution of a continuous variable
hist(data$DEMO_age,
     main = "Age of participants",
     xlab = "Age (years)",
     ylab = "Number of participants",
     col = "grey80")

# --- 6.2 Bar plot: counts of each value of the loneliness score
barplot(table(data$LONELY_ucla_loneliness_scale_score),
        main = "UCLA Loneliness Scale scores",
        xlab = "UCLA score (3 = least lonely, 9 = most lonely)",
        ylab = "Number of participants",
        col = "grey80")

# --- 6.3 Bar plot: counts of a categorical variable
barplot(table(data$mental_health_3cat),
        main = "Self-rated mental health",
        xlab = "Self-rated mental health",
        ylab = "Number of participants",
        col = "grey80")

# --- 6.4 Box plot: a continuous variable across groups
boxplot(LONELY_ucla_loneliness_scale_score ~ age_group, data = data,
        main = "Loneliness score by age group",
        xlab = "Age group (years)",
        ylab = "UCLA Loneliness Scale score",
        col = "grey80")

# --- 6.5 Scatter plot: two continuous variables
# jitter() adds a little random noise so stacked points can be seen
plot(jitter(data$DEMO_age), jitter(data$LONELY_ucla_loneliness_scale_score),
     main = "Loneliness score by age",
     xlab = "Age (years)",
     ylab = "UCLA Loneliness Scale score",
     pch = 16, col = rgb(0, 0, 0, 0.15))


# ========================================================
# SECTION 7: STRATIFIED TABLES (CROSS-TABULATIONS)
# ========================================================

# --- 7.1 Counts: rows = age group, columns = loneliness category
table(data$age_group, data$ucla_high_low)

# addmargins() adds row and column totals
addmargins(table(data$age_group, data$ucla_high_low))

# --- 7.2 Row percentages (margin = 1): each ROW adds to 100%
# "Within each age group, what percentage scored High?"
round(100 * prop.table(table(data$age_group, data$ucla_high_low),
                       margin = 1), 1)

# --- 7.3 Column percentages (margin = 2): each COLUMN adds to 100%
# "Within the High group, what percentage were in each age group?"
round(100 * prop.table(table(data$age_group, data$ucla_high_low),
                       margin = 2), 1)

# --- 7.4 A grouped bar plot of the row percentages
age_lonely_pct <- round(100 * prop.table(table(data$ucla_high_low,
                                               data$age_group),
                                         margin = 2), 1)
barplot(age_lonely_pct, beside = TRUE,
        main = "Loneliness category by age group",
        xlab = "Age group (years)",
        ylab = "Percentage of age group",
        ylim = c(0, 70),
        legend.text = TRUE, col = c("grey85", "grey35"))

# --- 7.5 The combined variable against loneliness
round(100 * prop.table(table(data$partner_living, data$ucla_high_low),
                       margin = 1), 1)

# --- 7.6 Stratify by a third variable: mean score by age AND gender
round(tapply(data$LONELY_ucla_loneliness_scale_score,
             list(data$age_group, data$gender_3cat),
             mean, na.rm = TRUE), 2)


# ========================================================
# SECTION 8: SUBSETS AND THE ANALYSIS DATASET
# ========================================================

# --- 8.1 Keep only the rows in one group
# Comparing with NA gives NA, and data[NA, ] returns an empty row
data_high <- data[data$ucla_high_low == "High", ]
nrow(data_high)

# which() returns only the row numbers where the condition is TRUE
data_high <- data[which(data$ucla_high_low == "High"), ]
nrow(data_high)

# --- 8.2 The analysis dataset: cleaned variables, complete rows only
analysis_vars <- c("LONELY_ucla_loneliness_scale_score", "ucla_high_low",
                   "DEMO_age", "age_group", "gender_3cat",
                   "mental_health_3cat", "partner_living")
analysis_data <- data[, analysis_vars]
nrow(analysis_data)
analysis_data <- na.omit(analysis_data)
nrow(analysis_data)


# ========================================================
# SECTION 9: TABLE 1
# ========================================================

# install.packages("tableone") once
library(tableone)

# --- 9.1 Describe the whole analysis sample
table1_vars <- c("DEMO_age", "age_group", "gender_3cat",
                 "mental_health_3cat", "partner_living")
table1 <- CreateTableOne(vars = table1_vars, data = analysis_data)
print(table1, showAllLevels = TRUE)

# --- 9.2 Describe the sample separately for each loneliness category
table1_by_ucla <- CreateTableOne(vars = table1_vars,
                                 strata = "ucla_high_low",
                                 data = analysis_data,
                                 test = FALSE)
print(table1_by_ucla, showAllLevels = TRUE)

# --- 9.3 Save the stratified table as a CSV file in the working directory
table1_print <- print(table1_by_ucla, showAllLevels = TRUE,
                      printToggle = FALSE)
write.csv(table1_print, "table1.csv")