#' Performs survival analysis
#' @description Performs survival analysis by grouping samples from maf based on mutation status of given gene(s) or manual grouping of samples.
#' @details This function takes MAF file and groups them based on mutation status associated with given gene(s) and performs survival analysis. Requires dataframe containing survival status and time to event.
#' Make sure sample names match to Tumor Sample Barcodes from MAF file.
#' @param maf an \code{\link{MAF}} object generated by \code{\link{read.maf}}
#' @param genes gene names for which survival analysis needs to be performed.
#' @param samples samples to group by. Genes and samples are mutually exclusive.
#' @param clinicalData dataframe containing events and time to events. Default looks for clinical data in annotation slot of \code{\link{MAF}}.
#' @param time column name contining time in \code{clinicalData}
#' @param Status column name containing status of patients in \code{clinicalData}. must be logical or numeric. e.g, TRUE or FALSE, 1 or 0.
#' @param groupNames names for groups. Should be of length two. Default c("Mutant", "WT")
#' @param col colors for plotting.
#' @param isTCGA FALSE. Is data is from TCGA.
#' @param showConfInt TRUE. Whether to show confidence interval in KM plot.
#' @param addInfo TRUE. Whether to show survival info in the plot.
#' @param textSize Text size for surv table. Default 7.
#' @param fn NULL. If provided saves pdf plot with basename fn.
#' @param width width of plot to be saved. Default 6
#' @param height height of plot to be saved. Default 6
#' @import survival
#' @importFrom gridExtra tableGrob ttheme_minimal
#' @return Survival plot
#' @examples
#' laml.maf <- system.file("extdata", "tcga_laml.maf.gz", package = "maftools")
#' laml.clin <- system.file("extdata", "tcga_laml_annot.tsv", package = "maftools")
#' laml <- read.maf(maf = laml.maf,  clinicalData = laml.clin)
#' mafSurvival(maf = laml, genes = 'DNMT3A', time = 'days_to_last_followup', Status = 'Overall_Survival_Status', isTCGA = TRUE)
#'
#'@export

mafSurvival = function(maf, genes = NULL, samples = NULL, clinicalData = NULL, time = "Time",
                       Status = "Status", groupNames = c("Mutant", "WT"), showConfInt = TRUE, addInfo = TRUE,
                       col = c('maroon', 'royalblue'), isTCGA = FALSE,
                       textSize = 12, fn = NULL, width = 6, height = 6){

  if(is.null(genes) & is.null(samples)){
    stop("Either provide Gene names or Sample names to group by.")
  }

  if(!is.null(genes) & !is.null(samples)){
    stop("Either provide Gene names or Sample names to group by. Not both!")
  }

  if(is.null(clinicalData)){
    message("Looking for clinical data in annoatation slot of MAF..")
    clinicalData = getClinicalData(x = maf)
    clinicalData = data.table::setDT(clinicalData)
  }else{
    clinicalData = data.table::setDT(clinicalData)
  }

  if(!"Tumor_Sample_Barcode" %in% colnames(clinicalData)){
    print(colnames(clinicalData))
    stop("Column Tumo_Sample_Barcode not found in clinical data. Check column names and rename it to Tumo_Sample_Barcode if necessary.")
  }

  #clinicalData$Tumor_Sample_Barcode = gsub(pattern = '-', replacement = '.', x = clinicalData$Tumor_Sample_Barcode)
  if(isTCGA){
    clinicalData$Tumor_Sample_Barcode = substr(x = clinicalData$Tumor_Sample_Barcode, start = 1, stop = 12)
  }

  if(length(colnames(clinicalData)[colnames(clinicalData) %in% time]) == 0){
    print(colnames(clinicalData))
    stop(paste0(time, " not found in clinicalData. Use argument time to povide column name containing time to event."))
  }else{
    colnames(clinicalData)[colnames(clinicalData) %in% time] = 'Time'
  }

  if(length(colnames(clinicalData)[colnames(clinicalData) %in% Status]) == 0){
    print(colnames(clinicalData))
    stop(paste0(Status, " not found in clinicalData. Use argument Status to povide column name containing events (Dead or Alive)."))
  }else{
    colnames(clinicalData)[colnames(clinicalData) %in% Status] = 'Status'
  }

  if(!is.null(genes)){
    #genesTSB = unique(as.character(subsetMaf(maf = maf, includeSyn = FALSE, genes = genes)[,Tumor_Sample_Barcode]))
    genesTSB = genesToBarcodes(maf = maf, genes = genes, justNames = TRUE)
    #genesTSB = lapply(X = genes, FUN = function(x) unique(as.character(subsetMaf(maf = maf, includeSyn = FALSE, genes = x)[,Tumor_Sample_Barcode])))
    #names(genesTSB) = genes
    genesTSB = genesTSB[sapply(genesTSB, FUN = function(x) length(x) != 0)]
    message("Number of mutated samples for given genes: ")
    print(sapply(genesTSB, FUN = length))

    genesMissing = genes[!genes %in% names(genesTSB)]
    if(length(genesMissing) > 0){
      genes = genes[!genes %in% genesMissing]
      genesMissing = paste(genesMissing, collapse = ', ')
      message(paste0("genes ", genesMissing, " does not seeem to be mutated. Removing them."))
    }

    if(length(genes) == 0){
      stop('None of the given genes are mutated!')
    }else{
      genes = paste(genes, collapse = ', ')
    }

    genesTSB = unique(as.character(unlist(genesTSB)))
  }else{
    #genesTSB = gsub(pattern = '-', replacement = '.', x = samples)
    genesTSB = samples
  }

  #clinicalData = clinicalData[!is.na(Time)]
  data.table::setDT(clinicalData)

  clinicalData$Time = suppressWarnings( as.numeric(as.character(clinicalData$Time)) )
  clinicalData$Status = suppressWarnings( as.integer(as.character(clinicalData$Status)) )
  clinicalData$Group = ifelse(test = clinicalData$Tumor_Sample_Barcode %in% genesTSB, yes = groupNames[1], no = groupNames[2])
  clin.mut.dat = clinicalData[,.(medianTime = median(Time, na.rm = TRUE),N = .N), Group][order(Group)]
  message("Median survival..")
  print(clin.mut.dat)

  clinicalData$Time = ifelse(test = is.infinite(clinicalData$Time), yes = 0, no = clinicalData$Time)

  surv.km = survival::survfit(formula = survival::Surv(time = Time, event = Status) ~ Group, data = clinicalData, conf.type = "log-log")
  res = summary(surv.km)

  surv.diff = survival::survdiff(formula = survival::Surv(time = Time, event = Status) ~ Group, data = clinicalData)

  surv.diff.pval = round(1 - pchisq(surv.diff$chisq, length(surv.diff$n) - 1), digits = 5)


  surv.dat = data.frame(Group = res$strata, Time = res$time, survProb = res$surv, survUp = res$upper, survLower = res$lower)
  surv.dat$Group = gsub(pattern = 'Group=', replacement = '', x = surv.dat$Group)

  surv.gg = ggplot(data = surv.dat, aes(x = Time, y = survProb, group = Group))+geom_line(aes(color = Group))+
    geom_point(aes(color = Group))+
    cowplot::theme_cowplot(font_size = 12, line_size = 1)+
    theme(legend.position = 'bottom', plot.title = element_text(size = 14, face = "bold"), axis.text.x = element_text(face = "bold"), axis.text.y = element_text(face = "bold"), axis.title.x = element_text(face = "bold"), axis.title.y = element_text(face = "bold"),
          plot.subtitle = element_text(size = 12, face = "bold", colour = ifelse(surv.diff.pval < 0.05, yes = 'red', no = 'black')),
          legend.title = element_blank())+
    ggtitle(label = paste0(groupNames[1], " v/s ", groupNames[2]), subtitle = paste0("P-value: ", surv.diff.pval))+
    xlab('Time')+ylab('Survival')+
    scale_color_manual(values = col)

  if(addInfo){
    ypos = as.numeric(as.character(quantile(surv.dat$Time, probs = seq(0, 1, 0.1))["90%"]))
    surv.gg = surv.gg+
      annotation_custom(grob = gridExtra::tableGrob(d = clin.mut.dat, rows = NULL,
                                         cols = c('Group', 'Median', '#Cases'),
                                         theme = gridExtra::ttheme_minimal(base_size = textSize)),
                                         ymin = 0.75, ymax = 0.95,
                                         xmin = ypos, xmax = max(surv.dat$Time, na.rm = TRUE))
  }
  ypos = as.numeric(quantile(surv.dat$Time, probs = seq(0, 1, 0.1))["70%"])

  if(showConfInt){
    surv.gg = surv.gg+geom_ribbon(aes(ymin = survLower, ymax = survUp), alpha = 0.2, fill = "grey70")
  }

  if(!is.null(fn)){
    cowplot::save_plot(filename = paste0(fn, '.pdf'), plot = surv.gg, base_height = height, base_width = width, bg = 'white')
  }

  surv.gg
}
