import string
import sys
from pylab import *
import scipy.stats
import matplotlib.gridspec as gridspec

SUFFIX_SCORES = "_rewards"
SUFFIX_INTRINSIC = "_intrinsic"
SUFFIX_ROOMS = "_rooms"
SUFFIX_ALL = "_all"

F_SUFFICES = ["CS", "IntrinsicOnly", "GridOnly", "NoCS"]
PLOT_NAMES = ["DeepCS", "No Grid", "No Intrinsic", "A2C"]

STEP_DIVISOR = 1000000.0
MAJOR_LINE_WIDTH = 6
MINOR_LINE_WIDTH = 4
PLOT_EVERY = 125
MARK_EVERY = 15
SIGNIFICANCE_SKIP = 3
PLOT_EXTREMES = False
NUM_BOOTSTRAPS = 1000

# Get command-line arguments
fPrefix = sys.argv[1]
fSuffix = sys.argv[2]
numFilesToLoad = int(sys.argv[3])
significanceThresholdString = sys.argv[4]
significanceThreshold = float(significanceThresholdString)
fileName = sys.argv[5]
plotTitle = sys.argv[6]
PLOT_CI = (sys.argv[7] == 'true')

# Define colors
colors = ["#AA0000", "#0055DD", "#00AA00", "#AA00AA", "#00AAAA", "#AAAA00", "#AAAAAA"]

# Set figure parameters
params = {
	"axes.labelsize": 48, #18
	"text.fontsize": 8,
	"legend.fontsize": 46, #18
	"xtick.labelsize": 36, #16
	"ytick.labelsize": 36, #16
	"text.usetex": False,
	"pdf.fonttype": 42,
	"figure.figsize": [25, 15],
	"xtick.major.pad" : 16,
	"xtick.minor.pad" : 16,
	"ytick.major.pad" : 16,
	"ytick.minor.pad" : 16,
	"lines.markersize": 24,
	"legend.borderaxespad": 0.1,
	"legend.columnspacing": 1
}
rcParams.update(params)

def average(arr, n):
    end =  n * int(len(arr)/n)
    return np.mean(arr[:end].reshape(-1, n), 1)

def median(arr, n):
    end =  n * int(len(arr)/n)
    return np.median(arr[:end].reshape(-1, n), 1)

# Loads the given run files
def loadFiles():
	dataToReturn = []
	xRangesToReturn = []
	for plotSuffix in xrange(0, len(F_SUFFICES)):
		# Load all files
		fileList = []
		for i in xrange(1, numFilesToLoad + 1):
			tempData = np.loadtxt(fPrefix + F_SUFFICES[plotSuffix] + "/" + str(i) + fSuffix + ".dat")
			if np.any(tempData[:, 1] >= 400):
				fileList.append(fPrefix + F_SUFFICES[plotSuffix] + "/" + str(i) + fSuffix + ".dat")
		# If all runs zero, just load all of them		
		if (len(fileList) == 0):
			for i in xrange(1, numFilesToLoad + 1):
				fileList.append(fPrefix + F_SUFFICES[plotSuffix] + "/" + str(i) + fSuffix + ".dat")
		# Get the maximum number of epochs to plots (each run may have +/- some epochs)
		maxEpochs = 99999999
		for f in fileList:
			tempFile = np.loadtxt(f)[:, 1]
			if (tempFile.shape[0] < maxEpochs):
				maxEpochs = tempFile.shape[0]
	
		# Load the data in a format that can be plotted
		i = 0
		data = np.zeros((len(fileList), maxEpochs))
		for f in fileList:
			data[i, :] = [0] * maxEpochs
			dataTemp = np.loadtxt(f)[:, 1]
			for item in xrange(0, maxEpochs):
				data[i, item] = dataTemp[item]
			i += 1

		# Get the x-labels
		xRange = np.zeros(maxEpochs)
		dataTemp = np.loadtxt(fileList[0])[:, 0]
		for item in xrange(0, maxEpochs):
			xRange[item] = dataTemp[item]
		xRange = np.divide(xRange, STEP_DIVISOR)

		# Downsample if necessary
		if PLOT_EVERY > 1:
			#data = [data[i][0::PLOT_EVERY] for i in range(0, len(fileList))]
			data = [median(data[i], PLOT_EVERY) for i in range(0, len(fileList))]
			data = np.array(data)
			xRange = xRange[0::PLOT_EVERY]
			if xRange.shape[0] > data.shape[1]:
				diff = -1 * (xRange.shape[0] - data.shape[1])
				xRange = xRange[:diff]

		dataToReturn.append(data)
		xRangesToReturn.append(xRange)
	return dataToReturn, xRangesToReturn

# Computes the median for a given data set
def computeMedian(data):
	median = np.zeros(data.shape[1])
	for i in xrange(0, len(median)):
		median[i] = np.median(data[:, i])
	return median

# Computes the maximum for a given data set
def computeMaximum(data):
	maximum = np.zeros(data.shape[1])
	absMax = 0
	for i in xrange(0, len(maximum)):
		maximum[i] = np.amax(data[:, i])
		if maximum[i] > absMax:
			absMax = maximum[i]
	return maximum, absMax

# Computes the minimum for a given data set
def computeMinimum(data):
	minimum = np.zeros(data.shape[1])
	for i in xrange(0, len(minimum)):
		minimum[i] = np.amin(data[:, i])
	return minimum

# Computes the median and 25/75 quartiles for a given data set
def computeQuartiles(data):
	median = np.zeros(data.shape[1])
	lowerQuartile = np.zeros(data.shape[1])
	upperQuartile = np.zeros(data.shape[1])
	for i in xrange(0, len(median)):
		median[i] = np.median(data[:, i])
		lowerQuartile[i] = np.percentile(data[:, i], 25)
		upperQuartile[i] = np.percentile(data[:, i], 75)
	return median, lowerQuartile, upperQuartile

# Computes the bootstrapped CIs of the median for a given data set
def computeCIs(data):
	median = np.zeros(data.shape[1])
	lowerCI = np.zeros(data.shape[1])
	upperCI = np.zeros(data.shape[1])
	for i in xrange(0, len(median)):
		median[i] = np.median(data[:, i])
		CI = scipy.stats.mstats.median_cihs(np.random.choice(data[:, i], NUM_BOOTSTRAPS, replace=True), alpha=0.05)
		
		lowerCI[i] = CI[0]
		upperCI[i] = CI[1]
	return median, lowerCI, upperCI

# Computes the Mann-Whitney U significance between two data sets
def computeSignificance(data1, data2):
	if (data1.shape[1] < data2.shape[1]):
		significance = np.zeros(data1.shape[1])
	else:
		significance = np.zeros(data2.shape[1])
	for i in xrange(0, len(significance)):
		setsAreIdentical = True
		setSize = len(data1[:, i])
		for j in xrange(0, setSize):
			if data1[j, i] != data2[j, i]:
				setsAreIdentical =  False
				break
		if setsAreIdentical:
			significance[i] = 1.0
		else:
			z, p = scipy.stats.mannwhitneyu(data1[:, i], data2[:, i])
			significance[i] = p * 2
	return significance

# Plots the given data
def plotData(axes):
	# Add a grid
	#grid(linewidth=2)

	# Set quartile fill
	fill_between(xRange0, lower0, upper0, alpha=0.25, linewidth=0, color=colors[0])
	fill_between(xRange1, lower1, upper1, alpha=0.25, linewidth=0, color=colors[1])
	fill_between(xRange2, lower2, upper2, alpha=0.25, linewidth=0, color=colors[2])
	fill_between(xRange3, lower3, upper3, alpha=0.25, linewidth=0, color=colors[3])

	# Plot the median
	plot(xRange0, median0, linewidth=MAJOR_LINE_WIDTH, marker="v", markevery=(10), markersize=30.0, color=colors[0])
	plot(xRange1, median1, linewidth=MAJOR_LINE_WIDTH, marker="s", markevery=(10), markersize=30.0, color=colors[1])
	plot(xRange2, median2, linewidth=MAJOR_LINE_WIDTH, marker="o", markevery=(10), markersize=30.0, color=colors[2])
	plot(xRange3, median3, linewidth=MAJOR_LINE_WIDTH, marker="*", markevery=(10), markersize=33.0, color=colors[3])

	# Fill in missing data
	if (xRangeMax0 < absoluteMaxXRange):
		fakeDataX = np.arange(xRangeMax0 + 1, absoluteMaxXRange)
		fakeDataY = np.zeros(len(fakeDataX))
		for i in xrange(len(fakeDataX)):
			fakeDataY[i] = median0[len(median0) - 1]
		plot(fakeDataX, fakeDataY, linewidth=MAJOR_LINE_WIDTH, linestyle="--", color=colors[0])
	if (xRangeMax1 < absoluteMaxXRange):
		fakeDataX = np.arange(xRangeMax1 + 1, absoluteMaxXRange)
		fakeDataY = np.zeros(len(fakeDataX))
		for i in xrange(len(fakeDataX)):
			fakeDataY[i] = median1[len(median1) - 1]
		plot(fakeDataX, fakeDataY, linewidth=MAJOR_LINE_WIDTH, linestyle="--", color=colors[1])
	if (xRangeMax2 < absoluteMaxXRange):
		fakeDataX = np.arange(xRangeMax2 + 1, absoluteMaxXRange)
		fakeDataY = np.zeros(len(fakeDataX))
		for i in xrange(len(fakeDataX)):
			fakeDataY[i] = median2[len(median2) - 1]
		plot(fakeDataX, fakeDataY, linewidth=MAJOR_LINE_WIDTH, linestyle="--", color=colors[2])
	if (xRangeMax3 < absoluteMaxXRange):
		fakeDataX = np.arange(xRangeMax3 + 1, absoluteMaxXRange)
		fakeDataY = np.zeros(len(fakeDataX))
		for i in xrange(len(fakeDataX)):
			fakeDataY[i] = median3[len(median3) - 1]
		plot(fakeDataX, fakeDataY, linewidth=MAJOR_LINE_WIDTH, linestyle="--", color=colors[3])

	# Plot the minimum
	#plot(xRange, minimum, linewidth=MINOR_LINE_WIDTH, color=colorToUse, linestyle="--")

	# Plot the maximum
	#plot(xRange, maximum, linewidth=MINOR_LINE_WIDTH, color=colorToUse, linestyle="--")

	# Set axis bounds
	axes.set_xlim(0, absoluteMaxXRange)
	#axes.set_xticks(np.arange(0, xRange.shape[0] + tickEveryX, tickEveryX))
	axes.set_ylim(0, absoluteMaxYRange * 1.2)
	#axes.set_yticks(np.arange(0, yRange + tickEveryY, tickEveryY))

	# Label x-axis
	#plt.xlabel("Training Steps (x" + str(int(STEP_DIVISOR)) + ")", weight="bold")
	plt.xlabel("Training Steps (millions)", weight="bold")
	
	# Label y-axis
	if (fSuffix == SUFFIX_SCORES):
		plt.ylabel("Game Score", weight="bold")
	elif (fSuffix == SUFFIX_INTRINSIC):
		plt.ylabel("Intrinsic Reward", weight="bold")
	elif (fSuffix == SUFFIX_ROOMS):
		plt.ylabel("Rooms Explored", weight="bold")
	elif (fSuffix == SUFFIX_ALL):
		plt.ylabel("Combined Reward", weight="bold")

	# Legend (top-center)
	l = legend(PLOT_NAMES, loc=9, ncol=4)
	l.get_frame().set_facecolor("0.9")
	l.get_frame().set_edgecolor("0.9")

# Plots the significance scores for a given axis size
def plotSignificance(axes):
	significanceValues0v1 = computeSignificance(data0, data1)
	significanceValues0v2 = computeSignificance(data0, data2)
	significanceValues0v3 = computeSignificance(data0, data3)
	
	# Add a grid
	grid(linestyle="-", linewidth=2)

	# Plot significance values (0 vs. 1)
	for i in xrange(0, len(significanceValues0v1), SIGNIFICANCE_SKIP):
		if significanceValues0v1[i] < significanceThreshold:
			plot(i, 5, marker="s", markersize=20.0, color=colors[1])
	# Plot significance values (0 vs. 2)
	for i in xrange(0, len(significanceValues0v2), SIGNIFICANCE_SKIP):
		if significanceValues0v2[i] < significanceThreshold:
			plot(i, -5, marker="o", markersize=20.0, color=colors[2])
	# Plot significance values (0 vs. 3)	
	for i in xrange(0, len(significanceValues0v3), SIGNIFICANCE_SKIP):
		if significanceValues0v3[i] < significanceThreshold:
			plot(i, -15, marker="*", markersize=23.0, color=colors[3])

	# Set axis bounds
	axes.set_xlim(0, absoluteMaxXRange)
	axes.set_ylim(-20, 10)

	# Disable labels
	axes.set_xticklabels([])
	axes.set_yticklabels([])
	plt.xticks([])
	plt.yticks([0, -10])

	# Indicate significance threshold
	plt.xlabel("p < " + significanceThresholdString + " vs. DeepCS", size=50) #14

# Obtain data	
dataItems, xRanges = loadFiles()
data0 = dataItems[0]
data1 = dataItems[1]
data2 = dataItems[2]
data3 = dataItems[3]
xRange0 = xRanges[0]
xRange1 = xRanges[1]
xRange2 = xRanges[2]
xRange3 = xRanges[3]

# Data distribution
if PLOT_CI:
	median0, lower0, upper0 = computeCIs(data0)
	median1, lower1, upper1 = computeCIs(data1)
	median2, lower2, upper2 = computeCIs(data2)
	median3, lower3, upper3 = computeCIs(data3)
else:
	median0, lower0, upper0 = computeQuartiles(data0)
	median1, lower1, upper1 = computeQuartiles(data1)
	median2, lower2, upper2 = computeQuartiles(data2)
	median3, lower3, upper3 = computeQuartiles(data3)

# Minimum values
minimum0 = computeMinimum(data0)
minimum1 = computeMinimum(data0)
minimum2 = computeMinimum(data0)
minimum3 = computeMinimum(data0)

# Maximum values
maximum0, yRange0 = computeMaximum(data0)
maximum1, yRange1 = computeMaximum(data1)
maximum2, yRange2 = computeMaximum(data2)
maximum3, yRange3 = computeMaximum(data3)

# Maximum x-range
xRangeMax0 = np.max(xRange0) - 1
xRangeMax1 = np.max(xRange1) - 1
xRangeMax2 = np.max(xRange2) - 1
xRangeMax3 = np.max(xRange3) - 1
absoluteMaxXRange = max(xRangeMax0, xRangeMax1, xRangeMax2, xRangeMax3) / 2

# Maximum y-range
if not PLOT_EXTREMES:
	absoluteMaxYRange = max(np.max(upper0), np.max(upper1), np.max(upper2), np.max(upper3))
else:
	absoluteMaxYRange = max(yRange0, yRange1, yRange2, yRange3)

# Create the figure
fig = figure()

# Create the grid layout
gs = gridspec.GridSpec(25, 1) # 1 columns by 22 rows

# First plot (scores)
ax1 = fig.add_subplot(gs[:19, 0])
plotData(ax1)

# Second plot (significance scores)
ax2 = fig.add_subplot(gs[22:, 0])
plotSignificance(ax2)

# Create figure title
fig.suptitle(plotTitle, weight="bold", fontsize=60)

# Save figure
#plt.show()
plt.savefig(fileName, bbox_inches="tight")
