import_datasets.R

# Double-checking and cleaning new datasets
# See also https://github.com/RaphaelS1/proba_benchmark/issues/28 for discussion on individual datasets
library(dplyr)

save_data = function(x, path = here::here("datasets/")) {
  xname = deparse(substitute(x))
  path_file_rds = paste0(path, xname, ".rds")
  path_file_csv = paste0(path, xname, ".csv")
  message("Saving ", xname, " at ", path_file_rds, " and ", path_file_csv)

  checkmate::assert_integerish(x[["status"]], lower = 0, upper = 1)
  checkmate::assert_numeric(x[["time"]], lower = 0)

  x_nomiss = na.omit(x)

  if (nrow(x) - nrow(x_nomiss) > 0) {
    message("Dropping ", nrow(x) - nrow(x_nomiss), " rows due to missings")
  }

  time_zeros = which(x_nomiss$time == 0)
  if (length(time_zeros) > 0) {
    message("Dropping ", length(time_zeros), " rows due to time = 0")
    x_nomiss = x_nomiss[-time_zeros, ]
  }

  saveRDS(x_nomiss, file = path_file_rds)
  write.csv(x_nomiss, file = path_file_csv, row.names = FALSE)
}


# dynpred::ova --------------------------------------------------------------------------------

ova = mlr3misc::load_dataset("ova", "dynpred") |>
  mutate(
    time = tyears,
    status = as.integer(d),
    Broders = factor(Broders),
    FIGOIII = as.integer(FIGO == "III"),
    Ascites = factor(Ascites),
    # Remove > < symbols as they can cause issues
    Diam = stringr::str_replace(Diam, ">", "gt"),
    Diam = stringr::str_replace(Diam, "<", "lt"),
    Diam = factor(Diam)
  ) |>
  select(-tyears, -d, -FIGO, -id)

save_data(ova)


# dynpred::wbc1 -------------------------------------------------------------------------------

wbc1 = mlr3misc::load_dataset("wbc1", "dynpred") |>
  mutate(
    time = as.numeric(tyears),
    status = as.integer(d)
  ) %>%
  select(-tyears, -d, -patnr)

save_data(wbc1)


# eha::child --------------------------------------------------------------
# Children born in Skellefteå, Sweden, 1850-1884, are followed fifteen years
# or until death or out-migration

child = eha::child |>
  select(
    -enter, # = 0 for all obs
    -id, -m.id, # ID vars
    -birthdate, # YYYY-MM-DD variable, reasonable to exclude?
    time = exit, status = event
  )

save_data(child)

# frailtyHL::bladder0 -----------------------------------------------------
# Bladder0 is a subset of 410 patients from a full data set with bladder cancer
# from 21 centers that participated in the EORTC trial (Sylvester et al., 2006).
# Time to event is the duration of the disease free interval (DFI), which is
# defined as time from randomization to the date of the first recurrence.

bladder0 = mlr3misc::load_dataset("bladder0", "frailtyHL") |>
  rename(time = Surtime, status = Status) |>
  mutate(
    Center = factor(Center) # Group ID was integer
    # time = ifelse(time == 0, 0.001, time)
  )

save_data(bladder0)

# frailtySurv::hdfail ------------------------------------------------------
# This dataset contains the observed follow-up times and SMART statistics of
# 52k unique hard drives.

hdfail = frailtySurv::hdfail |>
  select(-serial) # ID column

summary(hdfail$time)

# Coarsened version to reduce number of unique time points
# Values 0 < t < 1 will be rounded to time = 0, save_data() drops these
hdfail = hdfail |>
  dplyr::mutate(
    time = floor(time)
  )

save_data(hdfail)

# JM::aids.id -------------------------------------------------------------
# A randomized clinical trial in which both longitudinal and survival data were
# collected to compare the efficacy and safety of two antiretroviral drugs in
# treating patients who had failed or were intolerant of zidovudine (AZT) therapy.

# First measurements of JM::aids, therefore obstime == 0 here
aids.id = JM::aids.id |>
  rename(time = Time, status = death) |>
  select(
    -patient, # ID var
    -obstime, # 0 for all in this version, first measurement,
    -start, -stop, -event # not relevant here (per AB)
  )

save_data(aids.id)

# joineR::heart.valve -----------------------------------------------------

# Excluded because too few events after baseline filtering

# Per AB: long format with most features collected at follow up.
# Could be used if we only keep baseline covariates (which then are only few however).
# For time-to-event and status variables use last obs. of each subject (num)

# # Making sure to arrange by num/time for filtering
# heart.valve = joineR::heart.valve |>
#   arrange(num, time)
#
# # time to event: last obs per subject
# heart.valve_tte = heart.valve |>
#   select(num, time, status) |>
#   group_by(num) |>
#   slice_tail(n = 1) |>
#   ungroup()
#
# # Covariates: baseline, e.g. first per subject
# heart.valve_covars = heart.valve |>
#   select(-time, -status) |>
#   group_by(num) |>
#   slice_head(n = 1) |>
#   ungroup()
#
# # Double check no subject num was lost for whatever reason
# all(unique(heart.valve_covars$num) == unique(heart.valve_tte$num))
# nrow(heart.valve_covars) == nrow(heart.valve_tte)
#
# # Join by num explicitly just in case
# heart.valve = left_join(
#   heart.valve_tte,
#   heart.valve_covars,
#   by = "num") |>
#   select(-num)
#
# # Sadly only 54 events afterwards
# table(heart.valve)
#
# save_data(heart.valve)

# joineR::liver -----------------------------------------------------------
# Liver cirrhosis drug trial data

# Per AB: joineR::liver: can be used, but only first row per subject (id).
# Column survival is our event time, column time is the time at which liver
# function was measured (feature prothrombin). Keep the value at time = 0.

liver = joineR::liver |>
  filter(time == 0) |>
  select(
    -id, # ID var
    -time # Now irrelevant after filtering
  ) |>
  rename(time = survival, status = cens)

save_data(liver)

# KMSurv::channing ----------------------------------------------------------------------------


channing = mlr3misc::load_dataset("channing", "KMsurv") |>
  mutate(
    time = as.numeric(time),
    status = death,
    genderF = as.integer(gender - 1),
    ageentry = as.numeric(ageentry)
  ) %>%
  select(-obs, -age, -death, -gender) |>
  filter(time > 0)

save_data(channing)

# KMsurv::kidtran -----------------------------------------------------------------------------

kidtran = mlr3misc::load_dataset("kidtran", "KMsurv") |>
  mutate(
    time = as.numeric(time),
    status = delta,
    genderF = as.integer(gender - 1),
    raceBlack = as.integer(race - 1),
    age = as.numeric(age)
  ) |>
  select(-obs, -delta, -gender, -race)

save_data(kidtran)

# KMsurv::std -------------------------------------------------------------
# data from Section 1.12 in Klein and Moeschberger (1997)

std = mlr3misc::load_dataset("std", "KMsurv") |>
  rename(
    status = rinfct
  ) |>
  mutate(
    condom = factor(condom, labels = c("always", "sometime", "never")),
    iinfct = factor(iinfct, labels = c("gonorrhea", "chlamydia", "both"))
  ) |>
  select(-obs) # Identifier

save_data(std)

# MASS::Aids2 ---------------------------------------------------------------------------------

aids2 = mlr3misc::load_dataset("Aids2", "MASS") %>%
  mutate(
    time = as.numeric(death - diag),
    sexF = if_else(sex == "F", 1L, 0L),
    status = if_else(status == "D", 1L, 0L),
    age = as.numeric(age)
  ) %>%
  select(-diag, -death, -sex) %>%
  filter(time > 0)

save_data(aids2)


# mlr3proba::grace ----------------------------------------------------------------------------
requireNamespace("mlr3proba", quietly = TRUE)

grace = data.frame(mlr3::tsk("grace")$data())
save_data(grace)

# mlr3proba::whas -----------------------------------------------------------------------------
whas = data.frame(mlr3::tsk("whas")$data())
save_data(whas)

# locfit::livmet ----------------------------------------------------------
# Survival times for 622 patients diagnosed with Liver Metastases.
# "Beware, the censoring variable is coded as 1 = uncensored, so use cens=1-z
# in locfit() calls."
# -> so 0 = censored, no recoding needed. Reading is hard.
# https://ftp.uni-bayreuth.de/math/statlib/S/survcart
# Can't find source or real indication for dataset being real

# livmet = mlr3misc::load_dataset("livmet", "locfit") |>
#   rename(time = t, status = z) |>
#   mutate(
#     # status = 1 - status,
#     # Categorical recoding to factor for tumor TNM
#     tnm = factor(tnm),
#     # (1, 2) -> (0, 1), binary variables
#     sex = sex - 1,
#     pt = pt - 1,
#     lap = lap - 1
#   ) |>
#   filter(time > 0)
#
# save_data(livmet)

# RISCA::dataFTR --------------------------------------------------------
# Data were extracted from the DIVAT cohort. It corresponds to the reference
# sample constituted by first transplant recipients (FTR).

dataFTR = mlr3misc::load_dataset("dataFTR", "RISCA") |>
  rename(
    time = Tps.Evt,
    status = Evt
  )

save_data(dataFTR)

# RISCA::dataSTR --------------------------------------------------------
# Data were extracted from the DIVAT cohort. It corresponds to the relative
# sample constituted by second transplant recipients (STR).

# "Tattente2cl" = Waiting time between consecutive transplants, ok to include?

dataSTR = mlr3misc::load_dataset("dataSTR", "RISCA") |>
  rename(
    time = Tps.Evt,
    status = Evt
  )

save_data(dataSTR)

# pammtools::patient --------------------------------------------------------------------------

patient = pammtools::patient %>%
  mutate(
    time = Survdays,
    status = as.integer(PatientDied),
    sexF = if_else(Gender == "Female", 1L, 0L),
    age = as.numeric(Age)
  ) %>%
  select(
    -CombinedicuID, -CombinedID, -Survdays, -PatientDied, -survhosp,
    -Gender, -Age
  )

save_data(patient)

# pammtools::tumor ----------------------------------------------------------------------------

tumor = pammtools::tumor %>%
  mutate(
    time = days,
    transfusion = if_else(transfusion == "yes", 1L, 0L),
    complications = if_else(complications == "yes", 1L, 0L),
    metastases = if_else(metastases == "yes", 1L, 0L),
    resection = if_else(resection == "yes", 1L, 0L),
    charlson_score = factor(charlson_score),
    sexF = as.integer(sex == "female"),
    age = as.numeric(age)
  ) %>%
  select(-days, -sex) %>%
  as.data.frame()

save_data(tumor)

# pec::cost ---------------------------------------------------------------
# This data set contains a subset of the data from the Copenhagen stroke study

# This seems too good to be true almost
cost = mlr3misc::load_dataset("cost", "pec")

# Check to see if I missed a constant variable maybe?
# all(dim(cost) == dim(janitor::remove_constant(cost)))

save_data(cost)

# quantreg::uis -----------------------------------------------------------
# UIS Drug Treatment study data

# Exclude LEN.T and FRAC to avoid optimistic bias due to future information
uis = mlr3misc::load_dataset("uis", "quantreg") |>
  select(
    -ID, # ID var
    -Y, # Y: log(TIME)
    -LEN.T, -FRAC # Length of stay, compliance, not known at baseline
  ) |>
  rename(
    time = TIME,
    status = CENSOR
  ) |>
  mutate(
    HC = factor(HC), # Categorical, 4 levels
  )

save_data(uis)

# relsurv::rdata ----------------------------------------------------------
# "Survival data."
# Need to check: Pohar M., Stare J. (2006) "Relative survival analysis in R."
# Computer Methods and Programs in Biomedicine, 81: 272-278.

# year: recode to factor date
# age + agegrp can both stay

# Year column in class 'date' DDmonYY format for year extraction
rdata = mlr3misc::load_dataset("rdata", "relsurv") |>
  rename(status = cens) |>
  mutate(
    # Extract 2 digit year from date of diagnosis
    year = stringr::str_extract(date:::as.character.date(year), "\\d{2}$"),
    year = factor(year),
    # recode sex from 1,2 to 0,1 for consistency
    sex = sex - 1
  ) |>
  # Redundant variable, grouped version of age
  select(-agegr)

# Sanity check to ensure year recoding worked
stopifnot(length(unique(rdata$year)) == 5)

save_data(rdata)

# relsurv::colrec --------------------------------------------------------------
# Survival of patients with colon and rectal cancer diagnosed in 1994-2000.

# Removed due to date column by ML
# Re-added with converted date col and further preprocessing
colrec = relsurv::colrec |>
  mutate(
    # Convert date variable to numeric
    diag = as.numeric(diag),
    stage = ifelse(stage == 99, NA_real_, stage),
    sex = factor(sex, levels = 1:2, labels = c("male", "female"))
  ) |>
  rename(status = stat)

save_data(colrec)

# simPH::CarpenterFdaData -------------------------------------------------
# A data set from Carpenter (2002).
# https://www.jstor.org/stable/3088394 Table 1

# No clear doc for censoring variable but assuming normal coding
# based on usage in simPH example code with {survival} etc.

CarpenterFdaData = mlr3misc::load_dataset("CarpenterFdaData", "simPH") |>
  rename(
    time = acttime,
    status = censor
  ) |>
  select(
    -caseid, # ID var
    -X_st, # constant 1
    -X_t0, # constant 0
    -X_t, # identical to acttime (time).
    -X_d # identical to censor (status)
  )

save_data(CarpenterFdaData)

# smcure::e1684 -----------------------------------------------------------
# The melanoma data from the Eastern Cooperative Oncology Group (ECOG) phase
# III clinical trial e1684 which is used for modeling semicure PH mixture
# cure model (Kirkwood, et al., 1996)

e1684 = mlr3misc::load_dataset("e1684", "smcure") |>
  rename(time = FAILTIME, status = FAILCENS)

save_data(e1684)

# survival::flchain ---------------------------------------------------------------------------
flchain = mlr3misc::load_dataset("flchain", "survival") %>%
  mutate(
    mgus = as.factor(as.character(mgus)),
    time = futime,
    flc.grp = as.factor(as.character(flc.grp)),
    status = death
  ) %>%
  select(-futime, -chapter, -death, -creatinine, -chapter) %>%
  filter(time > 0)

save_data(flchain)

# survival::lung ------------------------------------------------------------------------------
lung = mlr3misc::load_dataset("lung", "survival") |>
  mutate(
    sex = factor(if_else(sex == 1, "M", "F")),
    ph.ecog = as.integer(ph.ecog),
    inst = factor(inst),
    status = status - 1
  ) |>
  tibble::remove_rownames()

save_data(lung)

# survival::nafld1 ----------------------------------------------------------------------------
nafld1 = survival::nafld1 |>
  mutate(
    sex = factor(if_else(male == 0, "F", "M")),
    # Convert time (days) to whole weeks to reduce number of unique time points
    time = round(futime / 7)
  ) |>
  select(-male, -case.id, -futime, -id)

save_data(nafld1)

# survival::nwtco -----------------------------------------------------------------------------
nwtco = mlr3misc::load_dataset("nwtco", "survival") %>%
  mutate(
    status = rel,
    time = as.numeric(edrel),
    age = as.numeric(age),
    histol = as.integer(histol - 1),
    stage = factor(stage)
  ) %>%
  select(-seqno, -study, -rel, -edrel, -in.subcohort, -instit)

# survival::transplant ------------------------------------------------------------------------
transplant = mlr3misc::load_dataset("transplant", "survival") %>%
  mutate(
    time = futime,
    status = if_else(event == "death", 1, 0)
  ) %>%
  select(-futime, -event) %>%
  filter(time > 0)

save_data(transplant)

# survival::mgus ------------------------------------------------------------------------------
mgus = mlr3misc::load_dataset("mgus", "survival") %>%
  mutate(
    time = futime,
    status = death
  ) %>%
  select(-futime, -death, -pcdx, -id, -pctime)

save_data(mgus)

# survival::veteran ---------------------------------------------------------------------------
veteran = mlr3misc::load_dataset("veteran", "survival") %>%
  mutate(
    trt = as.factor(as.character(trt)),
    prior = factor(if_else(prior == 0, "0", "1"))
  )

save_data(veteran)

# pycox datasets (gbsg, support) --------------------------------------------------------------

if (requireNamespace("reticulate", quietly = TRUE)) {
  # reticulate::py_install("pycox")
  pycox = reticulate::import("pycox")

  gbsg = pycox$datasets$gbsg$read_df() %>%
    mutate(
      x0 = as.integer(x0),
      x1 = as.integer(x1),
      x2 = as.integer(x2),
      status = event,
      time = as.numeric(duration)
    ) %>%
    select(-event, -duration)

  save_data(gbsg)

  metabric = pycox$datasets$metabric$read_df() %>%
    mutate(
      x4 = as.integer(x4),
      x5 = as.integer(x5),
      x6 = as.integer(x6),
      x7 = as.integer(x7),
      status = event,
      time = as.numeric(duration)
    ) %>%
    select(-event, -duration) %>%
    filter(time > 0)

  save_data(metabric)

  support = pycox$datasets$support$read_df() %>%
    mutate(
      x1 = as.integer(x1),
      x4 = as.integer(x4),
      x5 = as.integer(x5),
      x6 = as.integer(x6),
      status = event,
      time = as.numeric(duration)
    ) %>%
    select(-event, -duration)

  save_data(support)
}