# load libraries --------------------------------------------------------------------

library(survival)
library(data.table)
library(magrittr)
library(ggplot2)

options(stringsAsFactors = FALSE, scipen = 10)


# load data (from package) ----------------------------------------------------------

cancer2 <-
    copy(cancer) %>%
    setDT() # dataset stored in the survival package for testing
cancer3 <- cancer2[complete.cases(cancer2)] # taking complete cases to avoid NA observations on fitted model


# run cox model ---------------------------------------------------------------------

x <- coxph(Surv(time, status) ~ age + factor(sex) + factor(ph.ecog), data = cancer3) # cox model



# function to generate mean survival by factor --------------------------------------

adj_curves <- function(dt, obj, fac){
    if(!is.factor(dt[[fac]])) {
        dt[[fac]] <- factor(dt[[fac]])
    }
    groups <- levels(dt[[fac]])
    pred <- survfit(obj, dt)
    timepoints <- c(0, pred$time)
    timenames <- paste0("time", timepoints)
    adj_surv <- as.data.frame(t(pred$surv))
    adj_surv <- cbind(time0 = 1, adj_surv)
    names(adj_surv) <- timenames
    dt <- cbind(dt, adj_surv)
    dt2 <- dt[, lapply(.SD, mean), .SDcols = timenames, by = fac]
    graph <- as.data.table(cbind(timepoints, t(dt2[, -1])))
    names(graph) <- c("Time", groups)
    graph <- melt.data.table(data = graph, id.vars = "Time", variable.name = fac, value.name = "Survival")
    return(graph)
}


# run function and plot data --------------------------------------------------------

gr <- adj_curves(cancer3, x, "sex")

ggplot(gr, aes(x = Time, y = Survival, color = sex)) +
    geom_step() +
    scale_y_continuous(limits = c(0, 1)) +
    theme_bw()