library(getopt)

# get options first
spec <- matrix(c(
  'help', 'h', 0, "logical", "help",
  'models_path', 'm', 1, "character", "list of filepaths (comma separated) to gapseq-Draft-Models (RDS format) to reconstruct the pan-Model;
                                    Input can also be provided using wildcards or path to a directory",
  'rxn.weights.files_path', 'c', 2, "character", "list of filepaths (comma separated) to reaction weight tables generated by gapseq function \"generate_GSdraft.R\" (RDS format);
                                    Input can also be provided using wildcards or path to a directory",
  'rxnXgene.tables_path', 'g', 2, "character", "list of paths to tables with gene-X-reaction associations as generated by the \"generate_GSdraft.R\" (RDS format);
                                    Input can also be provided using wildcards or path to a directory",
  'pathways.table.path', 'w', 2, "character", "list of paths to all-Pathways tables as generated by the \"gapseq_find.sh\" (tbl format)",
  'min.rxn.freq.in.mods', 't', 2, "numeric", "minimum reaction frequency (mrf) to include the reactions in the pan-Draft. Default: 0.06",
  'only.binary.rxn.tbl','b', 2, "logical", "perform only models comparison to get a binary table summarizing rxn presence/absence. Default: FALSE",
  'output.dir', 'f', 2, "character", "path to directory, where output files will be saved (default: current directory)",
  'sbml.no.output', 's', 2, "logical", "Do not save model as sbml file. Default: Save as SBML"
  # 'verbose', 'v', 2, "logical", "verbose output and printing of debug messages. Default: FALSE"
), ncol = 5, byrow = T)

# Manually handle wildcard expansion for input files
# all input files belonging to the same option will be paste together separated by a comma
# if the parsed arguments strat with "-" are considered options, whereas everyting in between are input to the same option
args <- commandArgs(trailingOnly = TRUE)
args_mod <- vector()
parm_idx <- grep("^-[a-z]|^--[a-z]", args)
for(i in parm_idx){
  parm_nr <- which(i==parm_idx)
  if(parm_nr < length(parm_idx)){
    end <- parm_idx[parm_nr+1]-1
  } else{ 
    end <- length(args) 
  }
  parm_items <- args[i:end]
  if(length(parm_items) == 1){
    args_mod <- c(args_mod, parm_items)
  }else{
    args_mod <- c(args_mod, parm_items[1], paste0(parm_items[-1],collapse=","))
  }
}
opt <- getopt(spec, args_mod)

# Help Screen
if ( !is.null(opt$help) | is.null(opt$models_path)){
  cat(getopt(spec, usage=TRUE))
  q(status=1)
}

# get current script path
if (!is.na(Sys.getenv("RSTUDIO", unset = NA))) {
    # RStudio specific code
    script.dir    <- dirname(rstudioapi::getSourceEditorContext()$path)
} else{
    initial.options <- commandArgs(trailingOnly = FALSE)
    script.name <- sub("--file=", "", initial.options[grep("--file=", initial.options)])
    script.dir  <- dirname(script.name)
}

# Load R-packages
suppressMessages(library(cobrar))
suppressMessages(library(data.table)); setDTthreads(1)
suppressMessages(library(methods))

# Little helpers
source(paste0(script.dir,"/pan-draft_functions.R"))

# Setting defaults if required
if ( is.null(opt$output.dir) ) { opt$output.dir = "." }
if ( is.null(opt$min.rxn.freq.in.mods) ) { opt$min.rxn.freq.in.mods = 0.06 }
# if ( is.null(opt$verbose) ) { opt$verbose = F }
if ( is.null(opt$sbml.no.output) ) { opt$sbml.no.output = F } else { opt$sbml.no.output = T }
if ( is.null(opt$only.binary.rxn.tbl) ) { opt$only.binary.rxn.tbl = F }
if ( !opt$only.binary.rxn.tbl & ( is.null(opt$rxn.weights.files_path) | is.null(opt$rxnXgene.tables_path))) {
  stop("MISSING INPUT, if 'only.binary.rxn.tbl' is FALSE you should define also 'rxn.weights.files_path' and 'rxnXgene.tables_path'", call. = FALSE)
}

# Arguments:
mod.path              <- opt$models_path
rxn.weights.path      <- opt$rxn.weights.files_path
rxnXgene.table.path   <- opt$rxnXgene.tables_path
pathways.table.path   <- opt$pathways.table.path
min.rxn.freq.in.mods  <- opt$min.rxn.freq.in.mods
output.dir            <- opt$output.dir
only.binary.rxn.tbl   <- opt$only.binary.rxn.tbl
# verbose               <- opt$verbose

# Parameters:
dummy.weight <- 100
core.th <- 0.95
shell.th <- 0.05

# create output directory if not already there
dir.create(output.dir, recursive = TRUE, showWarnings = FALSE)
if (!dir.exists(output.dir) || file.access(output.dir, mode = 2) == -1)
  stop(paste("Output directory",output.dir,"cannot be created or is not writable."))

### Read the text file containing paths to full model rxn weigths and genes
cat("\nLoading input data, be patient... \n")
mod.path_opt <- parse_input_options(mod.path, "-draft.RDS")
model_list <- load_files_from_paths_for_RDS(mod.path_opt, "-draft.RDS")
if ( !only.binary.rxn.tbl ){
  weights_opt <- parse_input_options(rxn.weights.path, "-rxnWeights.RDS")
  weights_list <- load_files_from_paths_for_RDS(weights_opt, "-rxnWeights.RDS")
  xgenes_opt <- parse_input_options(rxnXgene.table.path, "-rxnXgenes.RDS")
  xgenes_list <- load_files_from_paths_for_RDS(xgenes_opt, "-rxnXgenes.RDS")
  pathways_opt <- parse_input_options(pathways.table.path, "-all-Pathways.tbl")
  pathways_list <- load_files_from_paths_for_tbl(pathways_opt, "-all-Pathways.tbl")
}
cat("loading completed \n\n")

# check input size
if ( !only.binary.rxn.tbl ){
  if(length(model_list)==length(weights_list) & length(model_list) == length(xgenes_list) & length(model_list) == length(pathways_list)){
    cat("The sizes of input lists are consistent, the number of loaded model is", length(model_list), "\n\n")
  } else {
    stop("Check you input, the number of loaded files per data type is differ")  
  } 

  ### Update rxnWeights table scores 
  num.mod <- length(weights_list)
  weights_dt <- rbindlist(weights_list, idcol = "model_id")
  weights_dt[, num.pan := .N, by = .(seed)] # Add lines for seed in order to obtain corrected median  
  # Calculate custom median of "weight" by grouping "seed"
  # alternative: "median(weight)" do not consider missing RXN and compute only the once present
  #              "custom_quartile_weight(weight, num.pan, num.mod, min.rxn.freq.in.mods)" consider missing RXN and compute the updated weight based on the total number input models and the threshold passed (e.g. 0.06)
  #                 has been tested the effect of changing th on gapfilling, it didn't showed considerable effects.  
  weights_dt[, weight.pan := custom_median(weight, num.pan, num.mod), by = .(seed)]
  weights_dt[, num.pan := NULL] # drop the colum num.pan
  weights_dt <- weights_dt[order(seed, weight)] # alternative: "abs(weight - weigth.pan)" 
  weights_dt <- weights_dt[!duplicated(seed)] # keep the entry with the highest score.
  weights_dt[,c("weight", "weight.pan")] <- weights_dt[,c("weight.pan", "weight")] # Swap the values of "weight" and "weight.pan" columns
  colnames(weights_dt)[colnames(weights_dt)=="weight.pan"] <- "weight.old" # Rename column
  setkeyv(weights_dt, c("model_id", "seed")) # set data.table keys
  saveRDS(weights_dt, file.path(output.dir,"panModel-rxnWeigths.RDS"))

  ### Update rxnXGenes table  
  xgenes_dt <- rbindlist(xgenes_list, idcol = "model_id")
  setkeyv(xgenes_dt, c("model_id", "seed"))
  xgenes_dt <- xgenes_dt[weights_dt] # extract only genes corresponding to the ref reactions
  saveRDS(xgenes_dt, file.path(output.dir,"panModel-rxnXgenes.RDS"))

  ### Save the list of detected Pathways
  pathways_dt <- do.call(rbind, pathways_list)
  pathways_dt <- pathways_dt[, .(Prediction = sum(Prediction)), by=c("ID", "Name")]
  pathways_dt$Prediction <- pathways_dt$Prediction > num.mod*min.rxn.freq.in.mods
  fwrite(pathways_dt, file = file.path(output.dir,"panModel-tmp-Pathways.tbl"), sep = "\t", quote = FALSE)

  vars_to_remove <- c("xgenes_list", "xgenes_dt", "weights_list", "weights_dt", "pathways_dt")
  rm(list = vars_to_remove) # Remove some variables
}

### Build dataset for RXN & MET frequency in mod
# look for duplicated model id, if "mod_id" was not specified when generating the Draft-models
mod_ids <- list()
for (mod_idx in names(model_list)) {
  mod <- model_list[[mod_idx]]
  if (mod@mod_id %in% mod_ids){ # if duplicated rename it
    n_name <- paste0(mod_idx, "_", mod@mod_id)
    cat(paste("\nModel ID:", mod@mod_id, "is duplicated\n"))
    cat(paste("renames as:", n_name))
    mod@mod_id <- n_name
    mod_ids <- c(mod_ids, n_name)
    model_list[[mod_idx]] <- mod
  } else {
    mod_ids <- c(mod_ids, mod@mod_id)
  }
  # Add annotation column to model attributes if not already there
  mod <- add_annotation_column_to_attributes(mod) 
}
mod_ids <- unlist(mod_ids)

# Build the data.table of presence/absence reaction in a list of models
res <- build_rxn2mod_dt(model_list) 
rxn2mod_dt <- res[[1]]
mod_id2mod_dict <- res[[2]]
first_mod_desc <- model_list[[1]]@mod_desc
rm(list = c("model_list", "res")) # Remove some variables
fwrite(rxn2mod_dt, file = file.path(output.dir,"rxnXmod.tsv"), sep = "\t", quote = FALSE)

if ( !only.binary.rxn.tbl ){
  # Core reactome size 
  strict_core_rxn_df <- subset(rxn2mod_dt, rowSums(rxn2mod_dt[, ..mod_ids]) == num.mod) # STRICT CORE
  core_rxn_df <- subset(rxn2mod_dt, rowSums(rxn2mod_dt[, ..mod_ids]) >= round(num.mod*core.th)) # CORE
  shell_rxn_df <- subset(rxn2mod_dt, rowSums(rxn2mod_dt[, ..mod_ids]) >= round(num.mod*shell.th)) # SHELL
  strict_core_rxn_num <- dim(strict_core_rxn_df)[1]
  core_rxn_num <- dim(core_rxn_df)[1] - strict_core_rxn_num
  shell_rxn_num <- dim(shell_rxn_df)[1] - core_rxn_num - strict_core_rxn_num
  cloud_rxn_num <- dim(rxn2mod_dt)[1] - shell_rxn_num - core_rxn_num - strict_core_rxn_num 
  cat(paste("\nThe total # of rxn is:", dim(rxn2mod_dt)[1], "\n"))
  cat(paste("The # of strict core rxn (all mod) is:", strict_core_rxn_num, "\n"))
  cat(paste("The # of core rxn ( >=", core.th*100, "% mod) is:", core_rxn_num, "\n"))
  cat(paste("The # of shell rxn ( >=", shell.th*100, "% mod) is:", shell_rxn_num, "\n"))
  cat(paste("The # of cloud rxn ( <", shell.th*100, "% mod) is:", cloud_rxn_num, "\n\n"))

  # write statistics on pan-reactome 
  stat <- c(dim(rxn2mod_dt)[2]-1, dim(rxn2mod_dt)[1], strict_core_rxn_num, core_rxn_num, shell_rxn_num, cloud_rxn_num)
  stat_dt <- data.table(t(stat))
  colnames(stat_dt) <- c("nGenome", "total_RXN", "strict_core (rxn freq. 100%)", "core (rxn freq. >95%)", "shell (rxn freq. 5-95%)", "cloud (rxn freq. <5%)")
  fwrite(stat_dt, file = file.path(output.dir,"pan-reactome_stat.tsv"), sep = "\t")

  ### Reconstruct pan-Draft
  # Extract RXN from model, extract for all rxn the info from the first model having that rxn 
  info_all_rxns_mods <- list()
  for (rxn_id in rxn2mod_dt$rxn) {
    first_modID_with_rxn <- colnames(rxn2mod_dt)[rxn2mod_dt[rxn_id, ]==1][1] # first model ID having the reaction
    first_mod_with_rxn <- attr(mod_id2mod_dict, first_modID_with_rxn) # first model having the reaction
    rxn <- getReactionPD(first_mod_with_rxn, rxn_id) # extract the reaction from a model based on rxn_ID
    info_all_rxns_mods <- c(info_all_rxns_mods, list(rxn)) # save rxn in a list
  }

  # find associations between MET_ID and MET_NAME/RXN_ID --- dictionary
  met_id2met_name_dict <- list()
  met_id2rxn_id_dict <- list()
  for (rxn in info_all_rxns_mods) {
    rxn_id <- rxn$react_id
    for (idx in seq_along(rxn$met_id)) { # identify the metabolites associated to each rxn
      met_id <- rxn$met_id[idx] 
      met_name <- rxn$met_name[idx] 

      # MET: metabolite ID to metabolite name dictionary
      if (!(met_id %in% names(met_id2met_name_dict))) {
        met_id2met_name_dict[[met_id]] <- list(met_name)
      } else if (!(met_name %in% met_id2met_name_dict[[met_id]])) { # save all the met_name associated to a specific met_id 
        met_id2met_name_dict[[met_id]] <- c(met_id2met_name_dict[[met_id]], met_name)
      }
      # RXN: metabolite ID to rxn ID dictionary
      if (!(met_id %in% names(met_id2rxn_id_dict))) {
        met_id2rxn_id_dict[[met_id]] <- list(rxn_id)
      } else if (!(rxn_id %in% met_id2rxn_id_dict[[met_id]])) {
        met_id2rxn_id_dict[[met_id]] <- c(met_id2rxn_id_dict[[met_id]], rxn_id) # save all the rxn_id associated to a specific met_id 
      }
    }
  }

  met_id2rxn_id_df <- pad_dict_to_dataframe(met_id2rxn_id_dict)
  met_id2met_name_padded_df <- pad_dict_to_dataframe(met_id2met_name_dict)
  rm(list = c("met_id2rxn_id_dict", "met_id2met_name_dict")) # Remove some variables

  cat(paste("Let's standanrdize the name of the duplicated compounds:\n"))
  dupl.th <- dim(met_id2met_name_padded_df)[2]-2 # ids have duplicated names if at least 2 different names are associated with the same id
  met_id2duplicated_met_name_df <- met_id2met_name_padded_df[!(rowSums(is.na(met_id2met_name_padded_df)) > dupl.th),] # find metabolites that have a duplicated name 
  info_all_rxns_mods <- standardize_duplicated_met_name(met_id2duplicated_met_name_df, met_id2rxn_id_df, info_all_rxns_mods) # Standardize the duplicated MET_NAME 

  # subset rxn2mod_dt based on min.rxn.freq.in.mods to reconstruct the pan-draft
  rxn_PresAbs_dt <- copy(rxn2mod_dt)
  cat(paste("\nth: ", min.rxn.freq.in.mods))
  col_n <- paste0("pan.mod_", min.rxn.freq.in.mods)
  th_rxn_df <- subset(rxn2mod_dt, rowSums(rxn2mod_dt[, ..mod_ids]) >= round(num.mod*min.rxn.freq.in.mods)) # Reactions present/absence in pan-Draft based on frequency th
  rxn_PresAbs_dt[, (col_n) := as.integer(rxn %in% th_rxn_df$rxn)] # add column on dataset 
  row_sub <- rxn_PresAbs_dt[,  ..col_n]==1   
  subSet_rxn_dt <- rxn_PresAbs_dt[row_sub[, 1], .SD, .SDcols = "rxn"] # select id of present reactions
  pan.mod <- build_panDraft(subSet_rxn_dt, info_all_rxns_mods, first_mod_desc) # build pan-Draft
  saveRDS(pan.mod, file.path(output.dir,"panModel-draft.RDS"))

  # Write SBML
  if(!opt$sbml.no.output){
    source(paste0(script.dir,"/sbml_write.R"))
    write_gapseq_sbml(pan.mod, file.path(output.dir,"panModel-draft"))
  }
}
