#' @include utilities.R surv_summary.R
NULL
#'Plot Survival Tables
#'
#'@description Plot survival tables:
#'\itemize{
#' \item \code{ggrisktable()}: Plot the number at risk table.
#' \item \code{ggcumevents()}: Plot the cumulative number of events table.
#' \item \code{ggcumcensor()}: Plot the cumulative number of censored subjects, the number of subjects who
#'  exit the risk set, without an event, at time t. Normally, users don't need
#'  to use this function directly.
#' \item \code{ggsurvtable()}: Generic function to plot any survival tables.
#'}
#'Normally, users don't need to use this function directly. Internally used by the function
#'  \code{\link{ggsurvplot}}.
#'
#'
#'@inheritParams ggsurvplot_arguments
#'@param fit an object of class survfit. Can be a list containing two
#'  components: 1) time: time variable used in survfit; 2) table: survival table
#'  as generated by the internal function .get_timepoints_survsummary(). Can be
#'  also a simple data frame.
#'@param survtable a character string specifying the type of survival table to plot.
#'@param risk.table.type risk table type. Allowed values include: "absolute" or
#'  "percentage": to show the \bold{absolute number} and the \bold{percentage}
#'  of subjects at risk by time, respectively. Use "abs_pct" to show both
#'  absolute number and percentage. Used only when survtable = "risk.table".
#'@param title the title of the plot.
#'@param xlog logical value. If TRUE, x axis is tansformed into log scale.
#'@param y.text logical. Default is TRUE. If FALSE, the table y axis tick
#'  labels will be hidden.
#'@param y.text.col logical. Default value is FALSE. If TRUE, the table tick
#'  labels will be colored by strata.
#'@param fontsize text font size.
#'@param font.family character vector specifying text element font family, e.g.: font.family = "Courier New".
#'@param ... other arguments passed to the function \code{\link{ggsurvtable}} and \code{\link[ggpubr]{ggpar}}.
#'@return a ggplot.
#'@author Alboukadel Kassambara, \email{alboukadel.kassambara@@gmail.com}
#' @examples
#' # Fit survival curves
#' #:::::::::::::::::::::::::::::::::::::::::::::::
#'require("survival")
#'fit<- survfit(Surv(time, status) ~ sex, data = lung)
#'
#'# Survival tables
#' #:::::::::::::::::::::::::::::::::::::::::::::::
#' tables <- ggsurvtable(fit, data = lung, color = "strata",
#'   y.text = FALSE)
#'
#' # Risk table
#' tables$risk.table
#'
#' # Number of cumulative events
#' tables$cumevents
#'
#' # Number of cumulative censoring
#' tables$cumcensor

#' @describeIn ggsurvtable Plot the number at risk table.
#' @export
ggrisktable <- function (fit, data = NULL,
                         risk.table.type = c("absolute", "percentage", "abs_pct", "nrisk_cumcensor", "nrisk_cumevents"),
                         ...)
{
  ggsurvtable(fit, data, survtable = "risk.table", ...)
}

#' @describeIn ggsurvtable Plot the cumulative number of events table
#' @export
ggcumevents <- function (fit, data = NULL, ...)
{
  ggsurvtable(fit, data, survtable = "cumevents", ...)
}

#' @describeIn ggsurvtable Plot the cumulative number of censor table
#' @export
ggcumcensor <- function (fit, data = NULL, ...)
{
  ggsurvtable(fit, data, survtable = "cumcensor", ...)
}

#' @describeIn ggsurvtable Generic function to plot survival tables: risk.table, cumevents and cumcensor
#' @export
ggsurvtable <- function (fit, data = NULL, survtable = c("cumevents",  "cumcensor", "risk.table"),
                         risk.table.type = c("absolute", "percentage", "abs_pct", "nrisk_cumcensor", "nrisk_cumevents"),
                         title = NULL, risk.table.title = NULL, cumevents.title = title, cumcensor.title = title,
                         color = "black", palette = NULL, break.time.by = NULL,  xlim = NULL,
                         xscale = 1, xlab = "Time", ylab = "Strata",
                         xlog = FALSE, legend = "top",
                         legend.title = "Strata", legend.labs = NULL,
                         y.text = TRUE, y.text.col = TRUE,
                         fontsize = 4.5, font.family = "",
                         axes.offset = TRUE,
                         ggtheme = theme_survminer(),
                         tables.theme = ggtheme, ...)
{

  if(is.data.frame(fit)){}
  else if(.is_list(fit)){
    if(!all(c("time", "table") %in% names(fit)))
      stop("fit should contain the following component: time and table")
  }
  else if(!.is_survfit(fit))
    stop("Can't handle an object of class ", class(fit))

  # Define time axis breaks
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  xmin <- ifelse(xlog, min(c(1, fit$time)), 0)
  if(is.null(xlim)) xlim <- c(xmin, max(fit$time))
  times <- .get_default_breaks(fit$time, .log = xlog)
  if(!is.null(break.time.by) &!xlog) times <- seq(0, max(c(fit$time, xlim)), by = break.time.by)



  # Surv summary at specific time points
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  if(.is_survfit(fit)){
    data <- .get_data(fit, data = data)
    survsummary <- .get_timepoints_survsummary(fit, data, times)
  }
  else if(.is_list(fit)){
    survsummary <- fit$table
  }
  else if(inherits(fit, "data.frame")){
    survsummary <- fit
  }

  opts <- list(
    survsummary = survsummary, times = times,
    survtable = survtable, risk.table.type = risk.table.type,  color = color, palette = palette,
    xlim = xlim, xscale = xscale,
    title = title, xlab = xlab, ylab = ylab, xlog = xlog,
    legend = legend, legend.title = legend.title, legend.labs = legend.labs,
    y.text = y.text, y.text.col = y.text.col,
    fontsize = fontsize, font.family = font.family,
    axes.offset = axes.offset,
    ggtheme = ggtheme, tables.theme = tables.theme,...)

  res <- list()
  time <- strata <- label <- n.event <- cum.n.event <- NULL

  # Ploting the cumulative number of events table
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

  if("cumevents" %in% survtable){
    opts$survtable = "cumevents"
    opts$title <- ifelse(is.null(cumevents.title),
                         "Cumulative number of events", cumevents.title)
    res$cumevents <- do.call(.plot_survtable, opts)

  }

  if("cumcensor" %in% survtable){
    opts$survtable = "cumcensor"
    opts$title <- ifelse(is.null(cumcensor.title),
                         "Cumulative number of events", cumcensor.title)
    res$cumcensor <- do.call(.plot_survtable, opts)

  }
  if("risk.table" %in% survtable){
    opts$survtable = "risk.table"
    if(is.null(risk.table.title)) opts$title <- NULL
    else opts$title <- risk.table.title
    res$risk.table <- do.call(.plot_survtable, opts)
  }


  if(length(res) == 1) res <- res[[1]]
  res
}








# Helper function to plot a specific survival table
.plot_survtable <- function (survsummary, times, survtable = c("cumevents", "risk.table", "cumcensor"),
                             risk.table.type = c("absolute", "percentage", "abs_pct", "nrisk_cumcensor", "nrisk_cumevents"),
                         color = "black", palette = NULL, xlim = NULL,
                         xscale = 1,
                         title = NULL, xlab = "Time", ylab = "Strata",
                         xlog = FALSE, legend = "top",
                         legend.title = "Strata", legend.labs = NULL,
                         y.text = TRUE, y.text.col = TRUE, fontsize = 4.5,
                         font.family = "",
                         axes.offset = TRUE,
                         ggtheme = theme_survminer(), tables.theme = ggtheme,
                          ...)
{

  survtable <- match.arg(survtable)
  risk.table.type <- match.arg(risk.table.type)

  # Defining plot title
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  if(is.null(title)){

    if(survtable == "risk.table"){
      risk.table.type <- match.arg(risk.table.type)
      title <- switch(risk.table.type,
                      absolute = "Number at risk",
                      percentage = "Percentage at risk",
                      abs_pct = "Number at risk: n (%)",
                      nrisk_cumcensor = "Number at risk (number censored)",
                      nrisk_cumevents = "Number at risk (number of events)",
                      "Number at risk")

    }
    else
    title <- switch(survtable,
                    cumevents = "Cumulative number of events",
                    cumcensor = "Number of censored subjects"
                    )
  }

  # Legend labels
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  if(is.null(color))
    color <- .strata.var <- "strata"
  else if(color %in% colnames(survsummary))
    .strata.var <- color
  else
    .strata.var <- "strata"

  # Number of strata and strata names
  .strata <- survsummary[, .strata.var]
  strata_names <- .levels(.strata)
  n.strata <- length(strata_names)

  # Check legend labels and title
  if(!is.null(legend.labs)){
    if(n.strata != length(legend.labs))
      warning("The length of legend.labs should be ", n.strata )
    else survsummary$strata <- factor(survsummary$strata, labels = legend.labs)
  }
  else if(is.null(legend.labs))
    legend.labs <- strata_names



  # Adjust table y axis tick labels in case of long strata
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  yticklabs <- rev(levels(survsummary$strata))
  n_strata <- length(levels(survsummary$strata))
  if(!y.text) yticklabs <- rep("-", n_strata)

  time <- strata <- label <- n.event <- cum.n.event  <- cum.n.censor<- NULL

  # Ploting the cumulative number of events table
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  if(survtable == "cumevents"){
    mapping <- aes(x = time, y = rev(strata),
                   label = cum.n.event, shape = rev(strata))
  }
  else if (survtable == "cumcensor"){
    mapping <- aes(x = time, y = rev(strata),
                   label = cum.n.censor, shape = rev(strata))

  }
  else if (survtable == "risk.table"){
    # risk table labels depending on the type argument
    pct.risk <- abs_pct.risk <- n.risk <- NULL
    llabels <- switch(risk.table.type,
                      percentage = round(survsummary$n.risk*100/survsummary$strata_size),
                      abs_pct = paste0(survsummary$n.risk, " (", survsummary$pct.risk, ")"),
                      nrisk_cumcensor = paste0(survsummary$n.risk, " (", survsummary$cum.n.censor, ")"),
                      nrisk_cumevents = paste0(survsummary$n.risk, " (", survsummary$cum.n.event, ")"),
                      survsummary$n.risk
    )
    survsummary$llabels <- llabels
    mapping <- aes(x = time, y = rev(strata),
                   label = llabels, shape = rev(strata))

  }


  # Plotting survival table
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  .expand <- ggplot2::waiver()
  # Tables labels Offset from origing
  if(!axes.offset){
    .expand <- c(0,0)
    offset <- max(xlim)/30
    survsummary <- survsummary %>%
      dplyr::mutate(time = ifelse(time == 0, offset, time))
  }

  p <- ggplot(survsummary, mapping) +
    scale_shape_manual(values = 1:length(levels(survsummary$strata)))+
    ggpubr::geom_exec(geom_text, data = survsummary, size = fontsize, color = color, family = font.family) +
    ggtheme +
    scale_y_discrete(breaks = as.character(levels(survsummary$strata)),labels = yticklabs ) +
    coord_cartesian(xlim = xlim) +
    labs(title = title, x = xlab, y = ylab, color = legend.title, shape = legend.title)

  if (survtable == "risk.table")
    p <- .set_risktable_gpar(p, ...) # For backward compatibility

  p <- ggpubr::ggpar(p, legend = legend, palette = palette,...)

  # Customize axis ticks
  #:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  xticklabels <- .format_xticklabels(labels = times, xscale = xscale)
  if(!xlog) p <- p + ggplot2::scale_x_continuous(breaks = times, labels = xticklabels, expand = .expand)
  else p <- p + ggplot2::scale_x_continuous(breaks = times,
                                            trans = "log10", labels = xticklabels)

  p <- p + tables.theme

  if(!y.text) {
    p <- .set_large_dash_as_ytext(p)
  }

  # Color table tick labels by strata
  if(is.logical(y.text.col) & y.text.col[1] == TRUE){
    cols <- .extract_ggplot_colors(p, grp.levels = legend.labs)
    p <- p + theme(axis.text.y = element_text(colour = rev(cols)))
  }
  else if(is.character(y.text.col))
    p <- p + theme(axis.text.y = element_text(colour = rev(y.text.col)))

  p

}



# For backward compatibility
# Specific graphical params to risk.table
.set_risktable_gpar <- function(p,  ...){
  extra.params <- list(...)
  ggpubr:::.labs(p,
                 font.main = extra.params$font.risk.table.title,
                 font.x = extra.params$font.risk.table.x,
                 font.y = extra.params$font.risk.table.y,
                 submain = extra.params$risk.table.subtitle,
                 caption = extra.params$risk.table.caption,
                 font.submain = extra.params$font.risk.table.subtitle,
                 font.caption = extra.params$font.risk.table.caption)
}


