# Chapter 1: code in reading order.

library(rmt)      # HF-ACTION teaching dataset
library(dplyr)    # Data preparation
library(survival) # Kaplan-Meier and Cox methods

data("hfaction", package = "rmt")
hf <- hfaction

# Inspect event records: a participant can occupy several rows.
head(hf)

# Count each randomized participant once, not once per event.
hf |>
  distinct(patid, trt_ab) |>
  count(trt_ab, name = "patients")

# Retain death or censoring; admissions do not end survival follow-up.
hf_death <- hf |>
  filter(status %in% c(0, 2)) |>
  mutate(death = status == 2)

# Order records within each patient; retain the earliest event or censoring.
hf_first <- hf |>
  arrange(patid, time, desc(status)) |>
  group_by(patid) |>
  slice_head(n = 1) |>
  ungroup() |>
  mutate(event = status != 0)

# Mortality model: time is in years; death marks observed deaths.
fit_death <- coxph(
  Surv(time, death) ~ trt_ab,
  data = hf_death
)

# Composite model: event marks the first admission or death.
fit_first <- coxph(
  Surv(time, event) ~ trt_ab,
  data = hf_first
)

# exp(coef) compares training (1) with usual care (0).
summary(fit_death)
summary(fit_first)

library(ggsurvfit)
library(patchwork)

# Shared appearance; muted red = usual care, blue = training.
book_theme <- theme_minimal(base_size = 12, base_family = "Georgia") +
  theme(panel.grid.minor = element_blank(),
        panel.grid.major = element_line(color = "#ded5da", linewidth = 0.3),
        text = element_text(color = "#332d34"),
        axis.text = element_text(size = 12, color = "#332d34"),
        axis.title = element_text(size = 12),
        legend.text = element_text(size = 12),
        plot.background = element_rect(fill = "#fffefd", color = NA),
        legend.position = "top", legend.title = element_blank())

# Mortality follow-up continues after hospitalization.
p_death <- survfit2(Surv(time, death) ~ trt_ab, data = hf_death) |>
  ggsurvfit(linewidth = 0.7) + scale_ggsurvfit() +
  scale_color_manual(values = c("#a85d60", "#50758a"),
                     labels = c("Usual care", "Training")) +
  scale_x_continuous("Time (years)", breaks = 0:4) +
  coord_cartesian(xlim = c(0, 4)) +
  labs(y = "Overall survival")

# Here the first hospitalization or death ends event-free follow-up.
p_first <- survfit2(Surv(time, event) ~ trt_ab, data = hf_first) |>
  ggsurvfit(linewidth = 0.7) + scale_ggsurvfit() +
  scale_color_manual(values = c("#a85d60", "#50758a"),
                     labels = c("Usual care", "Training")) +
  scale_x_continuous("Time (years)", breaks = 0:4) +
  coord_cartesian(xlim = c(0, 4)) +
  labs(y = "Hospitalization-free survival")

# Stack the two endpoints and share one treatment legend.
(p_death / p_first + plot_layout(guides = "collect")) & book_theme


library(Wcompo)

# Recode event types for Wcompo without changing the original status.
hf_weighted <- hf |>
  mutate(status_w = case_when(
    status == 0 ~ 0L, # Censoring
    status == 2 ~ 1L, # Death
    status == 1 ~ 2L  # Hospitalization
  ))

# Use all event records; patid links repeated admissions to a patient.
fit_weighted <- CompoML(
  id = hf_weighted$patid,
  time = hf_weighted$time,
  status = hf_weighted$status_w,
  Z = as.matrix(hf_weighted["trt_ab"]),
  w = c(2, 1) # Weights in status_w order: death, hospitalization
)

fit_weighted

# Match the book typography; these settings do not change estimates.
par(family = "Georgia", mar = c(3.6, 3.8, 0.8, 0.6),
    cex.axis = 1, cex.lab = 1,
    fg = "#332d34", col.axis = "#332d34", col.lab = "#332d34", las = 1,
    mgp = c(2.5, 0.7, 0), tcl = -0.25, bty = "l")
# z is treatment assignment: 0 = usual care, 1 = training.
plot(fit_weighted, z = 0, col = "#a85d60", lwd = 2,
     ylim = c(0, 5), xlim = c(0, 4),
     xlab = "Years since randomization", ylab = "Mean cumulative weighted count")
plot(fit_weighted, z = 1, add = TRUE, col = "#50758a", lwd = 2)
legend("topleft", c("Usual care", "Exercise training"),
       col = c("#a85d60", "#50758a"), lwd = 2, bty = "n")