--- a/create_heatmaps.R	Fri Jan 12 14:14:19 2018 -0500
+++ b/create_heatmaps.R	Fri Jan 12 14:14:38 2018 -0500
@@ -5,7 +5,8 @@
 option_list <- list(
     make_option(c("--input_dir"), action="store", dest="input_dir", help="IDEAS para files directory"),
     make_option(c("--output_dir"), action="store", dest="output_dir", help="PDF output directory"),
-    make_option(c("--script_dir"), action="store", dest="script_dir", help="R script source directory")
+    make_option(c("--script_dir"), action="store", dest="script_dir", help="R script source directory"),
+    make_option(c("--in_training_mode"), action="store_true", dest="in_training_mode", default=FALSE, help="Flag for training mode")
 )
 
 parser <- OptionParser(usage="%prog [options] file", option_list=option_list);
@@ -15,11 +16,17 @@
 heatmap_path <- paste(opt$script_dir, "create_heatmap.R", sep="/");
 source(heatmap_path);
 
-para_files <- list.files(path=opt$input_dir, pattern="\\.para$", full.names=TRUE);
+if (opt$in_training_mode) {
+    ext <- ".para";
+    para_files <- list.files(path=opt$input_dir, pattern="\\.para$", full.names=TRUE);
+} else {
+    ext <- ".para0";
+    para_files <- list.files(path=opt$input_dir, pattern="\\.para0$", full.names=TRUE);
+}
 for (i in 1:length(para_files)) {
     para_file <- para_files[i];
     para_file_base_name <- strsplit(para_file, split="/")[[1]][2];
-    output_file_base_name <- gsub(".para", "", para_file_base_name);
+    output_file_base_name <- gsub(ext, "", para_file_base_name);
     output_file_name <- paste(output_file_base_name, "state", i, "pdf", sep=".");
     output_file_path <- paste(opt$output_dir, output_file_name, sep="/");
     data_frame <- read.table(para_file, comment="!", header=T);