# libraries
library(cmprsk)
library(survminer)
library(ggplot2)

set.seed(2)
df <- data.frame(del=rexp(100)*5, event=sample(0:2,100,replace=TRUE), groupe=sample(c("A", "B"),100,replace=TRUE))
df$event_surv <- ifelse(df$event==0, 0, 1)

fit <- cuminc(df$del, df$event)
fit_surv <- survfit(Surv(del, event_surv)~1, data=df)

fit_gp <- cuminc(df$del, df$event, df$groupe)
fit_surv_gp <- survfit(Surv(del, event_surv)~groupe, data=df)

# Cumulative Incidence Curves for Competing Risks function

ggcompetingrisks1 <- function(fit, gnames = NULL, gsep=" ", event_suppr=NULL, 
                              multiple_panels = FALSE,
                              ggtheme = theme_survminer(),
                              coef = 1.96, conf.int = FALSE, ...) {
  # stopifnot(any(class(fit) %in% c("cuminc", "survfitms")))
  
  if (any(class(fit) == "cuminc")) {
    pl <- ggcompetingrisks.cuminc1(fit = fit, gnames=gnames, 
                                   gsep=gsep, multiple_panels=multiple_panels,
                                   coef = coef, conf.int = conf.int, event_suppr=event_suppr)
  }
  
  if (any(class(fit) == "list") & any(class(fit[[1]]) == "cuminc")) {
    pl <- ggcompetingrisks.cuminc.combine(fit = fit, gnames=gnames, 
                                   gsep=gsep, multiple_panels=multiple_panels,
                                   coef = coef, conf.int = conf.int, event_suppr=event_suppr)
  }
  
  if (any(class(fit) == "survfitms")) {
    pl <- ggcompetingrisks.survfitms(fit = fit)
  }
  
  pl <- pl + ggtheme +
    ylab("Probability of an event") + xlab("Time") +
    ggtitle("Cumulative incidence")
  ggpubr::ggpar(pl, ...)
}


ggcompetingrisks.cuminc1 <- function(fit, gnames = NULL, gsep=" ", 
                                     multiple_panels = FALSE, coef = 1.96, conf.int = FALSE, event_suppr=NULL) {
  if (!is.null(fit$Tests))
    fit <- fit[names(fit) != "Tests"]
  fit2 <- lapply(fit, `[`, 1:3)
  if (is.null(gnames)) gnames <- names(fit2)
  fit2_list <- lapply(seq_along(gnames), function(ind) {
    df <- as.data.frame(fit2[[ind]])
    df$name <- gnames[ind]
    df
  })
  time <- est <- event <- group <- NULL
  df <- do.call(rbind, fit2_list)
  df$event <- sapply(strsplit(df$name, split=gsep), `[`, 2)
  df$group <- sapply(strsplit(df$name, split=gsep), `[`, 1)
  df$std <- std <- sqrt(df$var)
  
  if(!is.null(event_suppr)){
    df=df[!df$event %in% event_suppr, ]
    if(unique(df$event)==1) {df$event=NULL}
  }
  
  if(length(unique(df$group))==1){
    pl <- ggplot(df, aes(time, est, color=event), legend="none")
  } else {
    if (multiple_panels) {
      pl <- ggplot(df, aes(time, est, color=event)) + facet_wrap(~group)
    } else {
      pl <- ggplot(df, aes(time, est, color=group, linetype=event))
    }
  }
  
  
  if (conf.int) {
    pl <- pl + geom_ribbon(aes(ymin = est - coef*std, ymax=est + coef*std), alpha = 0.2, linetype=0)
  }
  
pl + geom_line()

}



ggcompetingrisks.cuminc.combine <- function(fit, gnames = NULL, gsep=" ", 
                                     multiple_panels = TRUE, coef = 1.96, conf.int = FALSE, event_suppr=NULL) {
  # if (!is.null(fit$Tests))
  #   fit <- fit[names(fit) != "Tests"]
  
  fit2=fit
  for(i in 1:length(fit)){
    fit2[[i]]=lapply(fit2[[i]], `[`, 1:3)
  }
  
  fit2_list=fit2
  
  fit2_list=lapply(1:length(fit2_list), function(i) {
    gnames <- names(fit2_list[[i]])
    fit2_list[[i]] <- lapply(seq_along(gnames), function(ind) {
      df <- as.data.frame(fit2[[i]][[ind]])
      df$name <- gnames[ind]
      df$event <- sapply(strsplit(df$name, split=gsep), `[`, 2)
      df$group <- i
      df
    })
    fit2_list[[i]]=do.call(rbind, fit2_list[[i]])
  })
  
  time <- est <- event <- group <- NULL
  df <- do.call(rbind, fit2_list)
  
  
  df$std <- std <- sqrt(df$var)
  
  df$group=factor(df$group)
  
  if(!is.null(event_suppr)){
    df=df[!df$event %in% event_suppr, ]
    pl <- ggplot(df, aes(time, est, color=group)) 
  }
  
  if(is.null(event_suppr)){
    pl <- ggplot(df, aes(time, est, color=group, linetype=event)) 
  }
  
  # if (multiple_panels) {
  #   pl <- ggplot(df, aes(time, est, color=event + facet_wrap(~group)))
  # } else {
  #   pl <- ggplot(df, aes(time, est, color=group, linetype=event))
  # }
  if (conf.int) {
    pl <- pl + geom_ribbon(aes(ymin = est - coef*std, ymax=est + coef*std), alpha = 0.2, linetype=0)
  }
  
  pl + geom_line()
  
}
ggcompetingrisks.survfitms <- function(fit) {
  times <- fit$time
  psta <- as.data.frame(fit$pstate)
  colnames(psta) <- fit$states
  if (is.null(fit$strata)) {
    psta$strata <- "all"
  } else {
    psta$strata <- rep(names(fit$strata), fit$strata)
  }
  psta$times <- times
  psta <- .rename_empty_colname(
    psta, newname = "."
  )
  
  event <- value <- strata <- NULL
  pstal <- gather(psta, event, value, -strata, -times)
  ggplot(pstal, aes(times, value, fill=event)) +
    geom_area() + facet_wrap(~strata)
  
}

.rename_empty_colname <- function(df, newname = "."){
  empty.col <- colnames(df) == ""
  empty.col.exist <- length(empty.col) > 0
  if(empty.col.exist)
    colnames(df)[empty.col] <- newname
  df
}