# =============================================================================
# 02_reproduce_figures.R
# The Psychometrician's Handbook, Chapter 1 (Correlations)
#
# Reproduces the data figures from the chapter, and is meant to be taken apart
# and reused. Each block is self-contained: run the setup at the top, then any
# single figure block on its own.
#
# Every block opens with a "SET THESE" section. In most cases swapping your own
# PARAMCD values, or your own data frame, is the only change needed.
#
#   source("01_simulate_coa_data.R")   # or point ADCORR at your own file
#   source("02_reproduce_figures.R")
#
# Figures are written to figures/ as PNG at 200 dpi.
# =============================================================================

# ----------------------------------------------------------------- setup ----
pkgs <- c("dplyr","tidyr","ggplot2","psych","polycor","scales")
for (p in pkgs) if (!requireNamespace(p, quietly = TRUE))
  install.packages(p, repos = "https://cloud.r-project.org")
suppressPackageStartupMessages({
  library(dplyr); library(tidyr); library(ggplot2); library(psych); library(polycor)
})

ADCORR  <- Sys.getenv("CH1_DATA", "adcorr.csv")   # <- your own long dataset here
OUTDIR  <- "figures"
dir.create(OUTDIR, showWarnings = FALSE)

adcorr <- if (exists("adcorr")) get("adcorr") else read.csv(ADCORR, stringsAsFactors = FALSE)

# Handbook palette. Swap these for your own brand and every figure follows.
VIOLET <- "#453065"; MAGENTA <- "#AA378B"; CORAL <- "#EC663B"
EMERALD <- "#019939"; TEAL <- "#069999"
GRAPHITE <- "#33323A"; MUTED <- "#6B6B72"; GRID <- "#E7E5EA"

# The house font is Calibri. If it is not registered with R, ggplot silently
# falls back to the default sans, which is fine. Set BASE_FAMILY <- "" to force it.
BASE_FAMILY <- if ("Calibri" %in% systemfonts::system_fonts()$family) "Calibri" else ""

theme_handbook <- function(base_size = 12) {
  theme_minimal(base_size = base_size, base_family = BASE_FAMILY) +
    theme(
      text             = element_text(colour = GRAPHITE),
      axis.text        = element_text(colour = MUTED, size = rel(0.9)),
      axis.title       = element_text(colour = GRAPHITE, size = rel(0.95)),
      panel.grid.minor = element_blank(),
      panel.grid.major = element_line(colour = GRID, linewidth = 0.4),
      panel.border     = element_blank(),
      axis.line.x      = element_line(colour = "#C9C7CD", linewidth = 0.5),
      axis.line.y      = element_line(colour = "#C9C7CD", linewidth = 0.5),
      plot.title       = element_text(face = "bold", size = rel(1.05), colour = GRAPHITE),
      plot.margin      = margin(10, 14, 8, 10),
      strip.text       = element_text(colour = GRAPHITE, face = "bold")
    )
}
save_fig <- function(plot, name, w = 6.2, h = 4.6) {
  ggsave(file.path(OUTDIR, paste0(name, ".png")), plot,
         width = w, height = h, dpi = 200, bg = "white")
  message("  wrote ", file.path(OUTDIR, paste0(name, ".png")))
}

# CDISC BDS is long. Most figures want one row per subject, so pivot once.
wide <- adcorr %>%
  filter(AVISIT == "Baseline") %>%
  pivot_wider(id_cols = USUBJID, names_from = PARAMCD, values_from = AVAL)

message("Reproducing Chapter 1 figures into ", OUTDIR, "/")

# ============================================================ FIGURE 1 ======
# Scatter with a linear fit.                        Deck slide 59.
# --- SET THESE ---------------------------------------------------------------
X <- "TARGTOT"; Y <- "RELATEDA"
XLAB <- "Target PRO total score"; YLAB <- "Related measure A"
# -----------------------------------------------------------------------------
r1 <- cor(wide[[X]], wide[[Y]])
p1 <- ggplot(wide, aes(.data[[X]], .data[[Y]])) +
  geom_point(colour = TEAL, alpha = 0.65, size = 2.1) +
  geom_smooth(method = "lm", formula = y ~ x, se = TRUE,
              colour = VIOLET, fill = VIOLET, alpha = 0.12, linewidth = 0.9) +
  annotate("text", x = -Inf, y = Inf, hjust = -0.25, vjust = 1.6,
           label = sprintf("r = %.2f", r1), colour = CORAL,
           fontface = "bold", size = 5.2, family = BASE_FAMILY) +
  labs(x = XLAB, y = YLAB) + theme_handbook()
save_fig(p1, "fig01_scatter", 5.4, 4.1)

# ============================================================ FIGURE 2 ======
# Correlation matrix.                               Deck slide 60.
# --- SET THESE ---------------------------------------------------------------
VARS <- c("TARGTOT","RELATEDA","RELATEDB","UNRELATD")
# -----------------------------------------------------------------------------
m <- cor(wide[, VARS], use = "pairwise.complete.obs")
m[upper.tri(m)] <- NA                       # show one triangle only
p2 <- as.data.frame(as.table(m)) %>%
  filter(!is.na(Freq)) %>%
  ggplot(aes(Var2, forcats::fct_rev(Var1), fill = Freq)) +
  geom_tile(colour = "white", linewidth = 1.4) +
  geom_text(aes(label = sprintf("%.2f", Freq)),
            colour = ifelse(abs(as.vector(na.omit(as.vector(m)))) > 0.6, "white", GRAPHITE),
            size = 4.1, fontface = "bold", family = BASE_FAMILY) +
  scale_fill_gradient2(low = CORAL, mid = "white", high = VIOLET,
                       midpoint = 0, limits = c(-1, 1), name = "r") +
  labs(x = NULL, y = NULL) + coord_equal() +
  theme_handbook() + theme(panel.grid.major = element_blank(),
                           axis.line = element_blank(),
                           axis.text.x = element_text(angle = 30, hjust = 1))
save_fig(p2, "fig02_correlation_matrix", 5.4, 4.4)

# ============================================================ FIGURE 3 ======
# Bland-Altman: difference against the mean.        Deck slide 61.
# Plot the difference against the MEAN, never against one of the two measures:
# doing the latter produces a slope even when the measures are unrelated
# (Clifton & Clifton, 2019).
# --- SET THESE ---------------------------------------------------------------
BA_PARAM <- "TARGTOT"; BA_V1 <- "Baseline"; BA_V2 <- "Retest"
# -----------------------------------------------------------------------------
ba <- adcorr %>%
  filter(PARAMCD == BA_PARAM, AVISIT %in% c(BA_V1, BA_V2)) %>%
  pivot_wider(id_cols = USUBJID, names_from = AVISIT, values_from = AVAL) %>%
  mutate(avg = (.data[[BA_V1]] + .data[[BA_V2]]) / 2,
         dif = .data[[BA_V2]] - .data[[BA_V1]])
bias <- mean(ba$dif); loa <- 1.96 * sd(ba$dif)
p3 <- ggplot(ba, aes(avg, dif)) +
  geom_hline(yintercept = 0, colour = GRID, linewidth = 0.6) +
  geom_point(colour = TEAL, alpha = 0.65, size = 2.1) +
  geom_hline(yintercept = bias, colour = VIOLET, linewidth = 0.9) +
  geom_hline(yintercept = bias + c(-loa, loa), colour = CORAL,
             linetype = "dashed", linewidth = 0.8) +
  annotate("text", x = Inf, y = bias + loa, hjust = 1.02, vjust = -0.6,
           label = sprintf("upper LoA %.1f", bias + loa),
           colour = CORAL, size = 3.4, family = BASE_FAMILY) +
  annotate("text", x = Inf, y = bias - loa, hjust = 1.02, vjust = 1.5,
           label = sprintf("lower LoA %.1f", bias - loa),
           colour = CORAL, size = 3.4, family = BASE_FAMILY) +
  annotate("text", x = Inf, y = bias, hjust = 1.02, vjust = -0.6,
           label = sprintf("bias %.2f", bias),
           colour = VIOLET, size = 3.4, fontface = "bold", family = BASE_FAMILY) +
  labs(x = sprintf("mean of %s and %s", BA_V1, BA_V2),
       y = sprintf("%s minus %s", BA_V2, BA_V1)) +
  theme_handbook()
save_fig(p3, "fig03_bland_altman", 5.4, 4.1)

# ============================================================ FIGURE 4 ======
# Spaghetti plot with the group mean.               Deck slide 62.
# --- SET THESE ---------------------------------------------------------------
SP_PARAM <- "SYMBUR"; SP_XLAB <- "Month"
# -----------------------------------------------------------------------------
sp <- adcorr %>% filter(PARAMCD == SP_PARAM)
# With a few hundred subjects the trajectories merge into a solid block. Drawing
# a random subset keeps the shape readable; raise SP_SHOW to draw everyone.
SP_SHOW <- 60
set.seed(2026)
keep_ids <- sample(unique(sp$USUBJID), min(SP_SHOW, dplyr::n_distinct(sp$USUBJID)))
p4 <- ggplot(sp, aes(AVISITN, AVAL)) +
  geom_line(data = subset(sp, USUBJID %in% keep_ids), aes(group = USUBJID),
            colour = TEAL, alpha = 0.30, linewidth = 0.4) +
  stat_summary(fun = mean, geom = "line", colour = "white", linewidth = 2.6) +
  stat_summary(fun = mean, geom = "line", colour = VIOLET, linewidth = 1.5) +
  stat_summary(fun = mean, geom = "point", colour = VIOLET, size = 3) +
  scale_x_continuous(breaks = sort(unique(sp$AVISITN))) +
  labs(x = SP_XLAB, y = SP_PARAM) + theme_handbook()
save_fig(p4, "fig04_spaghetti", 5.4, 4.1)

# ============================================================ FIGURE 5 ======
# Three coefficients on the same ordinal pair.      Deck slide 51.
# --- SET THESE ---------------------------------------------------------------
I1 <- "ITEM02"; I2 <- "ITEM03"; TRUTH <- 0.713   # NA if you have no known truth
# -----------------------------------------------------------------------------
co <- data.frame(
  coefficient = factor(c("Pearson","Spearman","Polychoric"),
                       levels = c("Pearson","Spearman","Polychoric")),
  value = c(cor(wide[[I1]], wide[[I2]]),
            cor(wide[[I1]], wide[[I2]], method = "spearman"),
            suppressWarnings(polychoric(wide[, c(I1, I2)], progress = FALSE)$rho[1, 2])))
p5 <- ggplot(co, aes(coefficient, value, fill = coefficient)) +
  geom_col(width = 0.62, show.legend = FALSE) +
  geom_text(aes(label = sprintf("%.2f", value)), vjust = -0.55,
            fontface = "bold", size = 5, colour = GRAPHITE, family = BASE_FAMILY) +
  { if (!is.na(TRUTH)) geom_hline(yintercept = TRUTH, linetype = "dashed",
                                  colour = GRAPHITE, linewidth = 0.7) } +
  { if (!is.na(TRUTH)) annotate("text", x = 3.35, y = TRUTH,
                                label = sprintf("true r\n%.2f", TRUTH), hjust = 0,
                                fontface = "italic", size = 3.4, colour = GRAPHITE,
                                family = BASE_FAMILY) } +
  scale_fill_manual(values = c(Pearson = VIOLET, Spearman = TEAL, Polychoric = MAGENTA)) +
  scale_y_continuous(limits = c(0, 0.95), expand = expansion(mult = c(0, 0))) +
  coord_cartesian(xlim = c(0.5, 4.1), clip = "off") +
  labs(x = NULL, y = "estimated correlation") + theme_handbook()
save_fig(p5, "fig05_coefficient_bars", 5.6, 4.4)

# ============================================================ FIGURE 6 ======
# How often each response category was used.        Deck slide 51 inset.
# --- SET THESE ---------------------------------------------------------------
DIST_ITEM <- "ITEM02"; N_CAT <- 5
# -----------------------------------------------------------------------------
dd <- data.frame(cat = factor(1:N_CAT),
                 pct = as.numeric(prop.table(table(factor(wide[[DIST_ITEM]],
                                                          levels = 1:N_CAT)))) * 100)
p6 <- ggplot(dd, aes(cat, pct)) +
  geom_col(fill = MAGENTA, width = 0.66) +
  labs(x = "ordinal response category", y = NULL) +
  scale_y_continuous(expand = expansion(mult = c(0, 0.12))) +
  theme_handbook(11) + theme(axis.text.y = element_blank(),
                             panel.grid.major.y = element_blank())
save_fig(p6, "fig06_response_distribution", 4.0, 2.3)

# ============================================================ FIGURE 7 ======
# Monotonic but curved: Pearson understates it.     Deck slide 50.
# --- SET THESE ---------------------------------------------------------------
CX <- "FATVAS"; CY <- "WALK6MIN"
# -----------------------------------------------------------------------------
rp <- cor(wide[[CX]], wide[[CY]]); rs <- cor(wide[[CX]], wide[[CY]], method = "spearman")
p7 <- ggplot(wide, aes(.data[[CX]], .data[[CY]])) +
  geom_point(colour = TEAL, alpha = 0.6, size = 2) +
  geom_smooth(method = "loess", formula = y ~ x, se = FALSE,
              colour = VIOLET, linewidth = 1.1) +
  geom_smooth(method = "lm", formula = y ~ x, se = FALSE,
              colour = CORAL, linetype = "dashed", linewidth = 0.9) +
  annotate("text", x = Inf, y = Inf, hjust = 1.05, vjust = 1.5,
           label = sprintf("Pearson %.2f\nSpearman %.2f", rp, rs),
           colour = GRAPHITE, size = 4, family = BASE_FAMILY) +
  labs(x = "Fatigue VAS", y = "Six-minute walk distance (m)") + theme_handbook()
save_fig(p7, "fig07_curved_monotonic", 5.4, 4.1)

# ============================================================ FIGURE 8 ======
# Within against between subjects.                  Deck slide 31.
# --- SET THESE ---------------------------------------------------------------
WX <- "FINTOX"; WY <- "SYMBUR"; N_SHOW <- 8       # subjects to draw
# -----------------------------------------------------------------------------
wb <- adcorr %>% filter(PARAMCD %in% c(WX, WY)) %>%
  pivot_wider(id_cols = c(USUBJID, AVISITN), names_from = PARAMCD, values_from = AVAL)
pick <- wb %>% group_by(USUBJID) %>% summarise(m = mean(.data[[WX]]), .groups = "drop") %>%
  arrange(m) %>% slice(round(seq(1, n(), length.out = N_SHOW))) %>% pull(USUBJID)
sub <- wb %>% filter(USUBJID %in% pick)
pm  <- sub %>% group_by(USUBJID) %>%
  summarise(x = mean(.data[[WX]]), y = mean(.data[[WY]]), .groups = "drop")
r_w <- with(wb %>% group_by(USUBJID) %>%
              mutate(xw = .data[[WX]] - mean(.data[[WX]]),
                     yw = .data[[WY]] - mean(.data[[WY]])) %>% ungroup(), cor(xw, yw))
r_b <- cor(pm$x, pm$y)
subj_cols <- setNames(
  colorRampPalette(c(VIOLET, TEAL, EMERALD, CORAL, MAGENTA))(length(pick)), pick)
p8 <- ggplot(sub, aes(.data[[WX]], .data[[WY]])) +
  geom_smooth(data = pm, aes(x, y), method = "lm", formula = y ~ x, se = FALSE,
              colour = GRAPHITE, linetype = "dashed", linewidth = 1.1, inherit.aes = FALSE) +
  geom_point(aes(colour = USUBJID), size = 2.2, show.legend = FALSE) +
  geom_smooth(aes(colour = USUBJID), method = "lm", formula = y ~ x, se = FALSE,
              linewidth = 0.8, show.legend = FALSE) +
  geom_point(data = pm, aes(x, y), shape = 21, size = 4.6, stroke = 1.2,
             colour = GRAPHITE, fill = NA, inherit.aes = FALSE) +
  scale_colour_manual(values = subj_cols) +
  annotate("text", x = -Inf, y = Inf, hjust = -0.06, vjust = 1.5,
           label = sprintf("within a subject: %.2f", r_w),
           colour = TEAL, fontface = "bold", size = 4, family = BASE_FAMILY) +
  annotate("text", x = Inf, y = -Inf, hjust = 1.04, vjust = -1.0,
           label = sprintf("between subjects: %.2f", r_b),
           colour = GRAPHITE, fontface = "bold", size = 4, family = BASE_FAMILY) +
  labs(x = WX, y = WY) + theme_handbook()
save_fig(p8, "fig08_within_between", 5.8, 4.4)

# ============================================================ FIGURE 9 ======
# How the confidence interval narrows with n.       Deck slide 63.
# Analytic, via the Fisher z transform. No data needed, so it is a useful
# planning figure: set R_TRUE to the correlation you expect.
# --- SET THESE ---------------------------------------------------------------
R_TRUE <- 0.50; N_GRID <- c(20, 40, 80, 160, 320)
# -----------------------------------------------------------------------------
z <- atanh(R_TRUE)
ci <- data.frame(n = N_GRID) %>%
  mutate(se = 1 / sqrt(n - 3),
         lo = tanh(z - 1.96 * se), hi = tanh(z + 1.96 * se),
         lab = sprintf("[%.2f, %.2f]", lo, hi),
         nlab = factor(sprintf("n = %d", n), levels = rev(sprintf("n = %d", N_GRID))))
p9 <- ggplot(ci, aes(y = nlab)) +
  geom_vline(xintercept = R_TRUE, linetype = "dashed", colour = "#C9C7CD") +
  geom_segment(aes(x = lo, xend = hi, yend = nlab), colour = TEAL, linewidth = 2.4) +
  geom_point(aes(x = R_TRUE), colour = VIOLET, size = 3.6) +
  geom_text(aes(x = hi, label = lab), hjust = -0.12, size = 3.6,
            colour = GRAPHITE, family = BASE_FAMILY) +
  scale_x_continuous(limits = c(0, 1.22), breaks = seq(0, 1, 0.25)) +
  labs(x = "correlation r (95% CI)", y = NULL) +
  theme_handbook() + theme(panel.grid.major.y = element_blank())
save_fig(p9, "fig09_ci_by_sample_size", 5.6, 4.0)

# =========================================================== FIGURE 10 ======
# Sampling variability: the same truth, two sample sizes.  Deck slide 72.
# --- SET THESE ---------------------------------------------------------------
SV_TRUE <- 0.40; SV_N <- c(25, 200); SV_REPS <- 170
# -----------------------------------------------------------------------------
set.seed(2026)
sv <- do.call(rbind, lapply(SV_N, function(nn) {
  r <- replicate(SV_REPS, {
    x <- rnorm(nn); y <- SV_TRUE * x + sqrt(1 - SV_TRUE^2) * rnorm(nn); cor(x, y) })
  data.frame(n = nn, r = r)
})) %>% mutate(nlab = factor(sprintf("n = %d", n), levels = sprintf("n = %d", SV_N)))
p10 <- ggplot(sv, aes(r, nlab, colour = nlab)) +
  geom_vline(xintercept = SV_TRUE, linetype = "dashed", colour = VIOLET, linewidth = 0.9) +
  geom_jitter(height = 0.22, alpha = 0.55, size = 1.9, show.legend = FALSE) +
  annotate("text", x = SV_TRUE, y = 2.45, hjust = -0.08,
           label = sprintf("true r = %.1f", SV_TRUE), colour = VIOLET,
           fontface = "bold", size = 4, family = BASE_FAMILY) +
  scale_colour_manual(values = setNames(c(CORAL, TEAL)[seq_along(SV_N)],
                                        sprintf("n = %d", SV_N))) +
  coord_cartesian(ylim = c(0.5, 2.7)) +
  labs(x = sprintf("sample correlation r (%d samples each)", SV_REPS), y = NULL) +
  theme_handbook() + theme(panel.grid.major.y = element_blank())
save_fig(p10, "fig10_sampling_variability", 5.8, 4.0)

# =========================================================== FIGURE 11 ======
# Same r, four very different pictures.             Deck slide 58.
# Anscombe's quartet, which ships with R. Plot first, always.
q <- do.call(rbind, lapply(1:4, function(i)
  data.frame(set = paste("Dataset", i),
             x = anscombe[[paste0("x", i)]], y = anscombe[[paste0("y", i)]])))
labs11 <- q %>% group_by(set) %>% summarise(r = cor(x, y), .groups = "drop") %>%
  mutate(lab = sprintf("r = %.2f", r))
p11 <- ggplot(q, aes(x, y)) +
  geom_smooth(method = "lm", formula = y ~ x, se = FALSE, colour = VIOLET, linewidth = 0.9) +
  geom_point(colour = TEAL, size = 2.4) +
  geom_text(data = labs11, aes(x = -Inf, y = Inf, label = lab), hjust = -0.25, vjust = 1.6,
            colour = CORAL, fontface = "bold", size = 4.4, family = BASE_FAMILY,
            inherit.aes = FALSE) +
  facet_wrap(~ set) + labs(x = NULL, y = NULL) + theme_handbook()
save_fig(p11, "fig11_same_r_four_pictures", 6.2, 4.6)

# =========================================================== FIGURE 12 ======
# Restricted range pulls r down.                    Deck slide 74.
# --- SET THESE ---------------------------------------------------------------
RR_PARAM_X <- "TARGTOT"; RR_PARAM_Y <- "RELATEDA"; RR_KEEP_ABOVE <- 55
# -----------------------------------------------------------------------------
full <- wide %>% select(x = all_of(RR_PARAM_X), y = all_of(RR_PARAM_Y))
rest <- full %>% filter(x >= RR_KEEP_ABOVE)
rr <- rbind(transform(full, grp = sprintf("full range   r = %.2f", cor(full$x, full$y))),
            transform(rest, grp = sprintf("restricted   r = %.2f", cor(rest$x, rest$y))))
p12 <- ggplot(rr, aes(x, y, colour = grp)) +
  geom_point(alpha = 0.6, size = 2) +
  geom_smooth(method = "lm", formula = y ~ x, se = FALSE, linewidth = 1) +
  scale_colour_manual(values = setNames(c(TEAL, CORAL), unique(rr$grp)), name = NULL) +
  labs(x = RR_PARAM_X, y = RR_PARAM_Y) +
  theme_handbook() + theme(legend.position = "top")
save_fig(p12, "fig12_restricted_range", 5.8, 4.4)

message("done. ", length(list.files(OUTDIR, pattern = "png$")), " figures in ", OUTDIR, "/")