--- a/ks_distribution.R	Thu Jun 08 10:32:44 2017 -0400
+++ b/ks_distribution.R	Fri Jun 23 14:12:59 2017 -0400
@@ -5,22 +5,15 @@
 option_list <- list(
     make_option(c("-c", "--components_input"), action="store", dest="components_input", help="Ks significant components input dataset"),
     make_option(c("-k", "--kaks_input"), action="store", dest="kaks_input", help="KaKs analysis input dataset"),
-    make_option(c("-o", "--output"), action="store", dest="output", help="Output dataset")
+    make_option(c("-n", "--num_comp"), action="store", dest="num_comp", type="integer", help="Number of significant components in the Ks distribution"),
+    make_option(c("-o", "--output"), action="store", dest="output", help="Output dataset"),
+    make_option(c("-r", "--colors"), action="store", default=NA, help="List of component colors"),
 )
 
 parser <- OptionParser(usage="%prog [options] file", option_list=option_list)
 args <- parse_args(parser, positional_arguments=TRUE)
 opt <- args$options
 
-
-get_num_components = function(components_data)
-{
-    # Get the max of the number_comp column.
-    number_comp = components_data[, 3]
-    num_components <- max(number_comp, na.rm=TRUE)
-    return(num_components)
-}
-
 get_pi_mu_var = function(components_data, num_components)
 {
     # FixMe: enhance this to generically handle any integer value for num_components.
@@ -64,7 +57,7 @@
     return(results)
 }
 
-plot_ks<-function(kaks_input, output, pi, mu, var)
+plot_ks<-function(kaks_input, num_components, colors, output, pi, mu, var)
 {
     # Start PDF device driver to save charts to output.
     pdf(file=output, bg="white")
@@ -89,7 +82,51 @@
     barplot(nc, space=0.25, offset=0, width=0.04, xlim=c(0, max_ks), ylim=c(0, ymax), col="lightpink1", border="lightpink3")
     # Add x-axis.
     axis(1)
-    color <- c('red', 'yellow','green','black','blue', 'darkorange' )
+    if(is.na(colors)
+    {
+        color <- c('red', 'yellow', 'green', 'black', 'blue', 'darkorange')
+    }
+    else
+    {
+        # Handle specified colors for components.
+        cStr <- unlist(colors)
+        color <- c()
+        items <- strsplit(cStr, ",")
+        for (item in items) {
+            color <- c(color, item)
+        }
+        num_colors_specified = length(color)
+        if num_colors_specified < num_components:
+        {
+            for (i in num_colors_specified:num_components)
+            {
+                if !(any(color=='red')
+                {
+                    color <- c(color, 'red')
+                }
+                else if !(any(color=='yellow')
+                {
+                    color <- c(color, 'yellow')
+                }
+                else if !(any(color=='green')
+                {
+                    color <- c(color, 'green')
+                }
+                else if !(any(color=='black')
+                {
+                    color <- c(color, 'black')
+                }
+                else if !(any(color=='blue')
+                {
+                    color <- c(color, 'blue')
+                }
+                else
+                {
+                    color <- c(color, 'darkorange')
+                }
+            }
+        }
+    }
     for (i in 1:length(mu))
     {
        lines(vx, g[,i] * h, lwd=2, col=color[i])
@@ -117,8 +154,7 @@
 
 # Read in the components data.
 components_data <- read.delim(opt$components_input, header=TRUE)
-# Get the number of components.
-num_components <- get_num_components(components_data)
+num_components <- opt$num_comp
 
 # Set pi, mu, var.
 items <- get_pi_mu_var(components_data, num_components)
@@ -160,4 +196,4 @@
 }
 
 # Plot the output.
-plot_ks(opt$kaks_input, opt$output, pi, mu, var)
+plot_ks(opt$kaks_input, num_components, opt$colors, opt$output, pi, mu, var)