# Extracted from estimation.qmd, first R chunk.
# Original course code; not re-executed for this website refresh.

#| code-fold: true
#| code-summary: "Show the code"
#| eval: false

##################################################################
# This code generates the numerical results in chapter 3         #
##################################################################

# install.packages("rmt") # if not already installed
library(rmt)
library(tidyverse)

##### Read in HF-ACTION DATA ########
data(hfaction)
head(hfaction)
#> Displays the first few rows of the HF-ACTION dataset (patid, time, status, trt_ab)

# TFE: take the first event per patient id
# TFE = Time to First Event (death or hospitalization)
hfaction_TFE <- hfaction |>
  arrange(patid, time) |>
  group_by(patid) |>
  slice_head() |>
  ungroup()

###### Standard RMST analysis #################
library(survRM2)

# -------------------------------------------------------------------------
# Mortality analysis
# -------------------------------------------------------------------------

## get mortality data
hfaction_D <- hfaction |>
  filter(status != 1) # remove hospitalization records (status=1)

## RMST (Restricted Mean Survival Time) analysis for overall survival
# Here we compare arm=1 (training) vs arm=0 (usual care)
rmst_obj <- rmst2(
  time    = hfaction_D$time,
  status  = hfaction_D$status > 0,
  arm     = hfaction_D$trt_ab,
  tau     = 3.97
)
rmst_obj
#> Prints estimates of restricted mean survival (in years), differences, ratios, and p-values
# Between-group contrast 
#                       Est. lower .95 upper .95     p
# RMST (arm=1)-(arm=0) 0.238     0.013     0.464 0.039
# RMST (arm=1)/(arm=0) 1.074     1.003     1.150 0.040
# RMTL (arm=1)/(arm=0) 0.680     0.468     0.988 0.043

# Extract the difference in RMST (in years) and convert to months
rmst <- rmst_obj$unadjusted.result[1, 1] * 12 
rmst_p <- rmst_obj$unadjusted.result[1, 4]

# -------------------------------------------------------------------------
# TFE analysis
# -------------------------------------------------------------------------

## Check how many of the first events are death (status=1) or hospitalization (status=2)
hfaction_TFE |>
  count(status)

## RMST analysis for hospitalization-free survival
rmest_obj <- rmst2(
  time    = hfaction_TFE$time,
  status  = hfaction_TFE$status > 0,
  arm     = hfaction_TFE$trt_ab,
  tau     = 3.97
)
rmest_obj
#> Similar output for restricted mean event-free survival
# Between-group contrast 
#                       Est. lower .95 upper .95     p
# RMST (arm=1)-(arm=0) 0.198    -0.064     0.459 0.139
# RMST (arm=1)/(arm=0) 1.145     0.957     1.370 0.139
# RMTL (arm=1)/(arm=0) 0.924     0.832     1.027 0.141

# Extract the difference in event-free survival (years) and convert to months
rmest <- rmest_obj$unadjusted.result[1, 1] * 12
rmest_p <- rmest_obj$unadjusted.result[1, 4]

# -------------------------------------------------------------------------
# Mortality vs TFE (Kaplan-Meier plots with annotations)
# -------------------------------------------------------------------------

library(ggsurvfit)
library(patchwork)

# Plot for overall survival
pD <- survfit2(Surv(time, status > 0) ~ trt_ab, data = hfaction_D) |>
  ggsurvfit(linewidth = 1) +
  scale_ggsurvfit() +
  annotate(
    "text", x = 4, y = 1, hjust = 1, vjust = 1,
    label = str_c("4y-RMST = ", round(rmst, 2), " months",
                  " (P = ", round(rmst_p, 3), ")")
  ) +
  scale_color_discrete(labels = c("Usual care", "Training")) +
  scale_x_continuous("Time (years)", limits = c(0, 4)) +
  labs(y = "Overall survival")

# Plot for hospitalization-free survival
pTFE <- survfit2(Surv(time, status > 0) ~ trt_ab, data = hfaction_TFE) |>
  ggsurvfit(linewidth = 1) +
  annotate(
    "text", x = 4, y = 1, hjust = 1, vjust = 1,
    label = str_c("4y-RMEST = ", round(rmest, 2), " months",
                  " (P = ", round(rmest_p, 3), ")")
  ) +
  scale_ggsurvfit() +
  scale_color_discrete(labels = c("Usual care", "Training")) +
  scale_x_continuous("Time (years)", limits = c(0, 4)) +
  labs(
    y = "Hospitalization-free survival",
    caption = "RMEST: restricted mean event-free survival time"
  )

# Combine side-by-side
pD + pTFE + plot_layout(guides = "collect") &
  theme(
    legend.position = "top",
    legend.text = element_text(size = 12)
  ) 

# # ggsave("images/est_hfaction_unis.png", width = 8, height = 4.6)
#> Uncomment above to save the combined figure

# -------------------------------------------------------------------------
# RMT-IF analysis
# -------------------------------------------------------------------------

# rmtfit() from the rmt package fits the restricted mean time lost / 
# time free approach for recurrent events & death
obj <- rmtfit(rec(patid, time, status) ~ trt_ab, data = hfaction)

summary(obj, Kmax = 1, tau = 3.97)
#> Summarizes the model up to a maximum follow-up of 3.97 years, 
#> focusing on the first event (Kmax=1) or aggregated events

#############################################################
# Graphical analysis of the HF-ACTION trial to
# evaluate the effect of exercise training.
###########################################################

par(mfrow = c(1, 2))

# bouquet() displays the "bouquet plot" for each k up to Kmax=4
bouquet(
  obj, Kmax = 4, cex.group = 1.0,
  xlab = "Restricted mean win/loss time (years)",
  ylab = "Follow-up time (years)",
  group.label = FALSE, ylim = c(0, 4.2)
)
text(-0.8, 4.15, paste("Usual care"))
text(0.8, 4.15, paste("Exercise training"))

# plot() visualizes the RMT-IF difference as a function of time
plot(
  obj, conf = TRUE, lwd = 2,
  xlab = "Follow-up time (years)",
  ylab = "RMT-IF of training (years)", 
  main = ""
)

par(mfrow = c(1, 1))
#> Reset the graphic device to a single panel

### LaTeX table ###

pval_fmt3 = function(x) {
  if (x < 0.001) {
    return("$<$0.001")
  } else {
    return(round(x, 3))
  }
}

ltable = NULL
# aggregate the results for k=1,...,K
hosp_sum = summary(obj, Kmax = 1, tau = 3.97)$tab
# aggregate the results for k=4,...,K
all_sum = summary(obj, Kmax = 4, tau = 3.97)$tab

ltable = c(
  "&", "&", "&",
  round(12 * hosp_sum[1, 1], 2), "&", round(12 * hosp_sum[1, 2], 2),
  "&", pval_fmt3(hosp_sum[1, 4]), "\\"
)

for (i in 1:6) {
  tmp = c(
    "&", i, "&&",
    round(12 * all_sum[i, 1], 2), "&", round(12 * all_sum[i, 2], 2),
    "&", pval_fmt3(all_sum[i, 4]), "\\"
  )
  ltable = rbind(ltable, tmp)
}

ltable[5, 2] = "4+"
ltable[6:7, 2] = ""

rownames(ltable) = c("Hopitalization", "", "", "", "", "Death", "Overall")

noquote(ltable)
#> Produces a LaTeX-formatted table summarizing the RMT-IF results

#######################################################################
#               WA analysis                                           #
#######################################################################

# install.packages("WA") # if needed
library(WA)

# load the hf-action study data (already loaded, but repeated here)
head(hfaction)

# descriptive analysis

## death & hosp rates by treatment arm
hfaction |>
  group_by(trt_ab, patid) |>
  summarize(
    ND = sum(status == 2),
    NH = sum(status == 1)
  ) |>
  summarize(
    death_rate = mean(ND),
    avgNH = mean(NH),
    sdNH = sd(NH)
  )

# weighted while-alive event rate analysis 
# with death weighted as 2 vs 1 for hospitalization
obj <- LRfit(
  id = hfaction$patid,
  time = hfaction$time,
  status = hfaction$status,
  trt = hfaction$trt_ab,
  Dweight = 2
)
## print some descriptive information 
obj

## summarize the inference results at tau=4 years
summary(obj, tau = 3.97, joint.test = TRUE)

plot(obj)
#> Plots the Weighted-Alive event rates over time for each arm

# -------------------------------------------------------------------------
# unadjusted cumulative mean (Ch 1)
# -------------------------------------------------------------------------

## fit proportional means model with death = 2 x hosp
library(Wcompo)

## change status coding
status <- hfaction$status
status[status != 0] <- 3 - status[status != 0]
#> This transforms status=1->2, status=2->1, effectively swapping them

obj_ML <- CompoML(
  hfaction$patid,
  hfaction$time,
  status,
  hfaction$trt_ab,
  w = c(2, 1)
)
## summary results
t <- obj_ML$t     # time points
mu0 <- obj_ML$y   # baseline cumulative mean events
mu1 <- mu0 * exp(as.numeric(obj_ML$beta))
#> The exponentiated beta shifts the baseline function for the second arm

## plot survival-adjusted cumulative event 
## vs unadjusted under PM model
plot(
  obj,
  ylab = "Cumulative loss",
  xlab = "Time (years)"
)
lines(
  t[t <= 3.97], mu0[t <= 3.97],
  lty = 3, col = "red", lwd = 2
)
lines(
  t[t <= 3.97], mu1[t <= 3.97],
  lty = 3, col = "blue", lwd = 2
)
legend(
  "bottomright",
  col = c("red", "red", "blue", "blue"),
  c("Usual care (WA)", "Usual care (unadj)",
    "Training (WA)", "Training (unadj)"),
  lty = c(1, 3, 1, 3),
  lwd = 2
)
#> Compares Weighted-Alive-based event accumulation with the "unadjusted" PM model