#' Subset MAF
#'
#' @description Subsets MAF based on given conditions.
#' @param maf an MAF object generated by \code{\link{read.maf}}
#' @param tsb subset by these samples (Tumor Sample Barcodes)
#' @param genes subset by these genes
#' @param fields include only these fields along with necessary fields in the output
#' @param query query string. e.g, "Variant_Classification == 'Missense_Mutation'" returns only Missense variants.
#' @param mafObj returns output as MAF class \code{\link{MAF-class}}. Default FALSE
#' @param includeSyn Default TRUE, only applicable when mafObj = FALSE. If mafObj = TRUE, synonymous variants will be stored in a seperate slot of MAF object.
#' @param isTCGA Is input MAF file from TCGA source.
#' @param restrictTo restrict subset operations to these. Can be 'all', 'cnv', or 'mutations'. Default 'all'. If 'cnv' or 'mutations', subset operations will only be applied on copy-number or mutation data respectively, while retaining other parts as is.
#' @return subset table or an object of class \code{\link{MAF-class}}
#' @seealso \code{\link{getFields}}
#' @examples
#' laml.maf <- system.file("extdata", "tcga_laml.maf.gz", package = "maftools")
#' laml <- read.maf(maf = laml.maf)
#' ##Select all Splice_Site mutations from DNMT3A and NPM1
#' subsetMaf(maf = laml, genes = c('DNMT3A', 'NPM1'),
#' query = "Variant_Classification == 'Splice_Site'")
#' ##Select all variants with VAF above 30%
#' subsetMaf(maf = laml, query = "i_TumorVAF_WU > 30")
#' ##Extract data for samples 'TCGA.AB.3009' and 'TCGA.AB.2933' but only include vaf filed.
#' subsetMaf(maf = laml, tsb = c('TCGA.AB.3009', 'TCGA.AB.2933'), fields = 'i_TumorVAF_WU')
#' @export

subsetMaf = function(maf, tsb = NULL, genes = NULL, fields = NULL, query = NULL, mafObj = FALSE, includeSyn = TRUE, isTCGA = FALSE, restrictTo = 'all'){


  if(length(restrictTo) > 1){
    stop("restrictTo can only be 'all', 'cnv', or 'mutations'")
  }

  if(!restrictTo %in% c("all", "cnv", "mutations")){
    stop("restrictTo can only be 'all', 'cnv', or 'mutations'")
  }


  #Synonymous variants
  maf.silent <- maf@maf.silent
  #Main data
  maf.dat <- maf@data
  #Annotations
  maf.anno <- maf@clinical.data


  if(restrictTo == "cnv"){
    maf.silent.rest = maf.silent[!Variant_Type %in% 'CNV']
    maf.silent = maf.silent[Variant_Type %in% 'CNV']

    maf.dat.rest = maf.dat[!Variant_Type %in% 'CNV']
    maf.dat = maf.dat[Variant_Type %in% 'CNV']
  }else if(restrictTo == "mutations"){
    maf.silent.rest = maf.silent[Variant_Type %in% 'CNV']
    maf.silent = maf.silent[!Variant_Type %in% 'CNV']

    maf.dat.rest = maf.dat[Variant_Type %in% 'CNV']
    maf.dat = maf.dat[!Variant_Type %in% 'CNV']
  }


  #Select
  if(!is.null(tsb)){
    tsb = as.character(tsb)
    if(isTCGA){
      tsb = substr(x = tsb, start = 1, stop = 12)
    }
    maf.dat = maf.dat[Tumor_Sample_Barcode %in% tsb,]
    maf.silent = maf.silent[Tumor_Sample_Barcode %in% tsb,]
  }

  if(!is.null(genes)){
    genes = as.character(genes)
    maf.dat = maf.dat[Hugo_Symbol %in% genes, ]
    maf.silent = maf.silent[Hugo_Symbol %in% genes, ]
  }

  if(!is.null(query)){
    maf.dat = maf.dat[eval(parse(text=query))]
    maf.silent = maf.silent[eval(parse(text=query))]
  }

  default.fields = c('Hugo_Symbol', 'Chromosome', 'Start_Position', 'End_Position', 'Reference_Allele', 'Tumor_Seq_Allele2','Variant_Classification', 'Variant_Type', 'Tumor_Sample_Barcode') #necessary fields.

  if(!is.null(fields)){
    default.fields = unique(c(default.fields, fields))

    if(length(default.fields[!default.fields %in% colnames(maf.dat)]) > 0){
      message("Missing fields. Ignoring them.. ")
      print(default.fields[!default.fields %in% colnames(maf.dat)])
      default.fields = default.fields[default.fields %in% colnames(maf.dat)]
    }

    maf.dat = maf.dat[,default.fields, with = FALSE]
    maf.silent = maf.silent[,default.fields, with = FALSE]

    if(restrictTo != "all"){
      maf.dat.rest = maf.dat.rest[,default.fields, with = FALSE]
      maf.silent.rest = maf.silent.rest[,default.fields, with = FALSE]
    }
  }


  if(restrictTo != "all"){
    maf.dat = rbind(maf.dat, maf.dat.rest, fill = TRUE, use.names = TRUE)
    maf.silent = rbind(maf.silent, maf.silent.rest, fill = TRUE, use.names = TRUE)
  }


  if(mafObj){

    maf.silent = droplevels.data.frame(maf.silent)
    maf.dat = droplevels.data.frame(maf.dat)
    maf.anno = droplevels.data.frame(maf.anno)

    mafSummary = summarizeMaf(maf.dat, chatty = FALSE, anno = maf.anno)

    m = MAF(data = maf.dat, variants.per.sample = mafSummary$variants.per.sample, variant.type.summary = mafSummary$variant.type.summary,
              variant.classification.summary = mafSummary$variant.classification.summary, gene.summary = mafSummary$gene.summary,
              summary = mafSummary$summary, maf.silent = maf.silent, clinical.data = droplevels(mafSummary$sample.anno))

    return(m)
  }else{
    if(includeSyn){
      maf.dat = rbind(maf.dat, maf.silent, use.names = TRUE, fill = TRUE)
      maf.dat = droplevels.data.frame(x = maf.dat)
      return(maf.dat)
    }else{
      maf.dat = droplevels.data.frame(x = maf.dat)
      return(maf.dat)
    }
  }
}
