###########################################
## Required Packages (install these first):
###########################################

#install.packages("ggplot2", dependencies=TRUE) # install if necessary
#install.packages("gridExtra") # install if necessary

setwd("Z:/eval")

############################
## Import Required Packages:
############################

require(ggplot2) # the base ggplot library
require(grid) # required for advanced grid layouts
require(gridExtra) # required for advanced grid layouts
require(reshape2) # required for melting/transforming data

##########################
## User-Defined Constants:
##########################

USER.DEFINED.COLORS = c("#AA0000", "#00AA00", "#0055DD", "#AA00AA", "#00AAAA", "#AAAA00", "#AAAAAA")
USER.DEFINED.SHAPES = 21:25 # these shapes can be given a separate fill and border color

######################
## Data For Bar Plots:
######################

## Data files are originally in "wide format"; each column contains the data points for a single treatment.
## We use melt() to convert data from this format into the "long format" used by ggplot.
## Don't forget to set your working directory to include these files!

dataMR = melt(cbind(read.csv("evalMR_rewards.dat")))
dataMRAck = melt(cbind(read.csv("evalAckMR_rewards.dat")))
dataAm = melt(cbind(read.csv("evalAm_rewards.dat")))
dataSea = melt(cbind(read.csv("evalSea_rewards.dat")))
dataGrav = melt(cbind(read.csv("evalGrav_rewards.dat")))

#################################
## The AWESOME Bar Plot Function:
#################################

## Creates a box plot for each treatment in the given data set.
## Data should be be organized in "long" format (typically by using melt()).
plotBox = function(dataToUse, showLegend=FALSE, xLabel="", yLabel="", plotSignificance=TRUE, colorsToUse=rainbow(6)) {
  
  ## Internal function to draw significance bars/text on the given plot.
  ## Lines with the given sigGeometry will be drawn on the given plot.
  ## The annotation is drawn from strings in sigList; if sigList contains NA values, nothing will be drawn for those values.
  ## After all annotations have been added, the updated plot will be returned.
  plotBoxSignificance = function(gPlot, sigGeometry=NULL, sigList=NULL) {
    if (!is.null(sigGeometry) && !is.null(sigList)) {
      for (i in 1:length(sigGeometry)) {
        if (!is.na(sigList[[i]])) {
          gPlot = gPlot + geom_line(data=sigGeometry[[i]], aes(x=between, y=value)) +
            annotate("text",
                     x=mean(sigGeometry[[i]]$between),      # text annotation is in the middle of the drawn bar
                     y=max(sigGeometry[[i]]$value) + 0.02,  # text annotation is slightly above the top of the drawn bar
                     label=sigList[[i]], size=8)            # text annotation taken from sigList
        }
      }
    }
    return(gPlot)
  }
  
  ## Internal function to compute two-tailed Mann-Whitney U significance.
  ## Returns a string indicating the statistical significance between two data sets, or NA if there is no significant difference.
  ## To use another statistical test, simply modify the line responsible for performing the test.
  ## Mann-Whitney U is used here only because the example data is NOT normally distributed.
  computeSignificance = function(data1, data2) {
    if (identical(sort(data1), sort(data2))) {
      return(NA) # Mann-Whitney U cannot compute significance for identical sets,
                 #  but identical sets couldn't be significantly different anyway
    }
    pValue = wilcox.test(data1, data2, exact=FALSE)$p.value
    if (pValue < 0.0001) {
      return("****") # most stars = most significant
    } else if (pValue < 0.001) {
      return("***")
    } else if (pValue < 0.01) {
      return("**")
    } else if (pValue < 0.05) {
      return("*") # least stars = least significant
    }
    return(NA) # no significance
  }
  
  ## Draw the base box plot.
  ## If you don't need any of the significance frills, this is the ONLY code required to generate awesome box plots!
  ## You can modify most visual attributes of the plot here (except for significance bars/text, below).
  ## Because significance bars/text take up vertical space, the y-axis limits will be defined later.
  basePlot = ggplot() +
              stat_boxplot(data=dataToUse, aes(x=variable, y=value), geom="errorbar", lwd=1.2, width=0.5) + # plots the "ends" of the whiskers
              geom_boxplot(data=dataToUse, aes(x=variable, y=value, fill=variable), linetype="solid", lwd=1.2) + # the main box plot
              geom_jitter(data=dataToUse, aes(x=variable, y=value, fill=variable), color="#000000", pch=21, stroke=2, size=3, alpha=0.5,
                          position=position_jitter(height=0)) +
              #geom_dotplot(data=dataToUse, aes(x=variable, y=value, fill=variable), binaxis="y", stackdir="center", dotsize=0.25, alpha=0.5) +
              xlab(xLabel) + # x-axis label
              ylab(yLabel) + # y-axis label
              guides(fill=showLegend) + # whether a legend should be displayed
              scale_y_continuous(labels=scales::comma, breaks=pretty(dataToUse$value, n=5), limits=range(pretty(dataToUse$value))) +
              theme(axis.ticks.x=element_blank(),
                    axis.text.x=element_blank(),
                    axis.text.y=element_text(size=24),
                    axis.title.x=element_text(size=32, face="bold"),
                    axis.title.y=element_text(size=40, face="bold")
              ) +
              scale_fill_manual(values=colorsToUse) + # any "fill" parameter (e.g. box fill) will use these colors
              scale_color_manual(values=colorsToUse) # any "color" parameter (e.g. line color) will use these colors
  
  ## This is where things get a bit more complicated, but only because we want significance bars over the existing plot.
  ## The number of significance comparisons between treatments is defined as N choose 2.
  ## We need to (1) compute statistical significance and (2) manually create the geometry to display significance text/stars.
  ## Although ggplot doesn't do this by default, we can thankfully accomplish both tasks in the same loop (below).
  ## Of course, if we only graphing a single treatment, we don't need to do this since there's nothing to compare it against.
  numSig = choose(nlevels(dataToUse$variable), 2) # the number of pairwise comparisons is [numTreatments CHOOSE 2]
  if (plotSignificance && numSig >= 1) {
    # basePlot = basePlot + scale_y_continuous(breaks=c(0.2, 0.4, 0.6, 0.8, 1.0), # force y-axis text to these values
    #                                          limits=c(0, 1.47), # more area than needed (make room for significance bars/text)
    #                                          labels=scales::percent) # interpret decimals visually as a percentage
    sigGeom = list()
    sig = list()
    currEntry = 1
    currMax = max(dataToUse$value) + 0.02 # start drawing bars/text slightly above the maximum data value
    # Iterate through each pair of treatments in the data set.
    # At each iteration, we compute the significance between that treatment and another, then draw the geometry (stars, etc.).
    # Each subsequent bar/text is drawn slightly above the previous one to prevent collisions.
    for (i in 1:(nlevels(dataToUse$variable))) { # the treatment we are currently examining
      for (j in (i + 1):(nlevels(dataToUse$variable))) { # the treatment we are using as a comparison against i
        if (currEntry <= numSig) {
          maxI = currMax
          maxIPlus = maxI + 0.02 # vertical space to avoid collisions
          currMax = maxIPlus + 0.06 # vertical space to avoid collisions
          # Sets the geometry to be drawn.
          # This will be a square bracket drawn from the first treatment to the second treatment, with room for significance text/stars.
          sigGeom[[currEntry]] = data.frame(between=c(i, i:j, j), value=c(maxI, rep(maxIPlus, times=(j - i + 1)), maxI))
          # Compute the significance for the treatments being compared.
          # The function gives us a string that will be drawn above each bracket, but only if there is a statistical difference.
          sig[[currEntry]] = computeSignificance(dataToUse$value[dataToUse$variable == levels(dataToUse$variable)[i]],
                                                 dataToUse$value[dataToUse$variable == levels(dataToUse$variable)[j]])
          currEntry = currEntry + 1
        }
      }
    }
    return(plotBoxSignificance(basePlot, sigGeom, sig)) # actually plots the significance we obtained above, then returns the result
    
  ## If no significance test is desired, the box plot is ready to be returned.
  ## Only the y-axis limits need to be configured.
  } else {
    # basePlot = basePlot + scale_y_continuous(breaks=c(0.2, 0.4, 0.6, 0.8, 1.0), # force y-axis text to these values
    #                                          limits=c(0, 1.0), # a hard limit allows multiple plots to be compared side-by-side
    #                                          labels=scales::percent) # interpret decimals visually as a percentage
    return(basePlot)
  }
}

###################################
## Generate some AWESOME bar plots!
###################################

plot1 = plotBox(dataMR, xLabel="MR (A2C)", yLabel="Game Score (100 trials)",
                colorsToUse=USER.DEFINED.COLORS[1])
plot2 = plotBox(dataMRAck, xLabel="MR (ACKTR)", yLabel="",
                colorsToUse=USER.DEFINED.COLORS[2])
plot3 = plotBox(dataAm, xLabel="Amidar", yLabel="",
                colorsToUse=USER.DEFINED.COLORS[3])
plot4 = plotBox(dataSea, xLabel="Seaquest", yLabel="",
                colorsToUse=USER.DEFINED.COLORS[4])
plot5 = plotBox(dataGrav, xLabel="Gravitar", yLabel="",
                colorsToUse=USER.DEFINED.COLORS[5])

########################################
## Prepare plots for multi-graph layout.
########################################

## Combine all plots into a single list.
## This step is necessary for arranging the plots into an arbitrary layout.
## If you only have one plot, you can simply display or save it; this step is not necessary.
plots = list(plot1, plot2, plot3, plot4, plot5)

############################
## Multi-Graph: Side-by-Side
############################

## Note: if you only have one plot, skip this section and go to the next.
## The idea of grid.arrange() is to take a list of plots and then place them in a grid defined by the layout_matrix.
## This allows you to plot four plots as 2x2, 4x1, 1x4, etc., which would not be possible in regular ggplot.
## Each number in the layout matrix refers to the graph that "goes" there; 1 is the first graph in the list, 2 is the second, etc.
## Passing NA to the layout matrix creates a blank cell instead.

## Arrange the plots in a single row (i.e. side-by-side).
arrangedPlots1 = grid.arrange(grobs=plots, layout_matrix = rbind(c(1, 2, 3, 4, 5)),
                              top=textGrob("Best-Policy Evaluations", # single, centered graph title
                                           gp=gpar(fontsize=48, fontface="bold"))) # title formatting

#######################################
## Multi-Graph: Two Rows, Three Columns
#######################################

## Note: if you only have one plot, skip this section and go to the next.
## The idea of grid.arrange() is to take a list of plots and then place them in a grid defined by the layout_matrix.
## This allows you to plot four plots as 2x2, 4x1, 1x4, etc., which would not be possible in regular ggplot.
## Each number in the layout matrix refers to the graph that "goes" there; 1 is the first graph in the list, 2 is the second, etc.
## Passing NA to the layout matrix creates a blank cell instead.

## Arrange the plots in a grid 3 wide and 2 tall.
# arrangedPlots2 = grid.arrange(grobs=plots, layout_matrix = rbind(c(1, 2, 3),
#                                                                  c(4, 5, 6)),
#                               top=textGrob("Best Individual Fitness: All Mazes & Aggregate\n", # single, centered graph title
#                                            gp=gpar(fontsize=32, fontface="bold"))) # title formatting


###############################################
## Multi-Graph: 4x4 Grid With One Enlarged Plot
###############################################

## Note: if you only have one plot, skip this section and go to the next.
## The idea of grid.arrange() is to take a list of plots and then place them in a grid defined by the layout_matrix.
## This allows you to plot four plots as 2x2, 4x1, 1x4, etc., which would not be possible in regular ggplot.
## Each number in the layout matrix refers to the graph that "goes" there; 1 is the first graph in the list, 2 is the second, etc.
## Passing NA to the layout matrix creates a blank cell instead.

## Arrange the plots in a 4x4 grid.
## If you supply the same number to layout_matrix multiple times, the plot will be stretched to fill that space.
## In this particular example, the first plot will be enlarged relative to all other plots.
# arrangedPlots3 = grid.arrange(grobs=plots, layout_matrix = rbind(c(1, 1, 2),
#                                                                  c(1, 1, 3),
#                                                                  c(4, 5, NA),
#                                                                  c(NA, NA, 6)),
#                               top=textGrob("Best Individual Fitness: All Mazes & Aggregate\n", # single, centered graph title
#                                            gp=gpar(fontsize=32, fontface="bold"))) # title formatting


###############################
## Saving the AWESOME Bar Plots
###############################

## Saves each plot as a PNG image.
## To change the file type, just change the extension from .png to .jpg or whatever you want.
## Vector-based graphics will be used for PDFs, meaning you can zoom in without losing image quality.
## Keep in mind that some formats (e.g. JPG) will not preserve transparency even if your plot uses it.
## All dimensions are given in inches.
ggsave(filename="barPlot1.pdf", plot=arrangedPlots1, width=20, height=10)
# ggsave(filename="barPlot2.png", plot=arrangedPlots2, width=10, height=20)
# ggsave(filename="barPlot3.pdf", plot=arrangedPlots3, width=20, height=20)
