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

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

F_SUFFICES = ["CS", "NoCS"]
PLOT_NAMES = ["DeepCS", "A2C"]

STEP_DIVISOR = 1000000.0
MAJOR_LINE_WIDTH = 6
MINOR_LINE_WIDTH = 4
PLOT_EVERY = 125
MARK_EVERY = 30
SIGNIFICANCE_SKIP = 8
NUM_BOOTSTRAPS = 1000

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

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

# Set figure parameters
params = {
	"axes.labelsize": 44, #18
	"text.fontsize": 8,
	"legend.fontsize": 32, #18
	"xtick.labelsize": 28, #16
	"ytick.labelsize": 28, #16
	"text.usetex": False,
	"pdf.fonttype": 42,
	"figure.figsize": [25, 10],
	"xtick.major.pad" : 16,
	"xtick.minor.pad" : 16,
	"ytick.major.pad" : 16,
	"ytick.minor.pad" : 16,
	"lines.markersize": 24,
	"legend.borderaxespad": 0.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(suffixToLoad):
	dataToReturn = []
	xRangesToReturn = []
	for plotSuffix in xrange(0, len(F_SUFFICES)):
		# Load all files
		fileList = []
		for i in xrange(1, numFilesToLoad + 1):
			fileList.append(fPrefix + F_SUFFICES[plotSuffix] + "/" + str(i) + suffixToLoad + ".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, colorStart, suffix, xRangeMax,
		median0, lower0, upper0, xRange0, yRange0, 
		median1, lower1, upper1, xRange1, yRange1,
		minimum0=None, maximum0=None, minimum1=None, maximum1=None, inset=False):
	# Set quartile fill
	fill_between(xRange0, lower0, upper0, alpha=0.25, linewidth=0, color=colors[colorStart])
	fill_between(xRange1, lower1, upper1, alpha=0.25, linewidth=0, color=colors[colorStart + 1])

	# Plot the median
	plot(xRange0, median0, linewidth=MAJOR_LINE_WIDTH, marker="v", markevery=(MARK_EVERY), markersize=25.0, color=colors[colorStart])
	plot(xRange1, median1, linewidth=MAJOR_LINE_WIDTH, marker="s", markevery=(MARK_EVERY), markersize=25.0, color=colors[colorStart + 1])	

	# Plot the minimum
	if minimum0 is not None:
		plot(xRange0, minimum0, linewidth=MINOR_LINE_WIDTH, color=colors[colorStart], linestyle="--")
	if minimum1 is not None:
		plot(xRange1, minimum1, linewidth=MINOR_LINE_WIDTH, color=colors[colorStart + 1], linestyle="--")

	# Plot the maximum
	if maximum0 is not None:
		plot(xRange0, maximum0, linewidth=MINOR_LINE_WIDTH, color=colors[colorStart], linestyle="--")
	if maximum1 is not None:
		plot(xRange1, maximum1, linewidth=MINOR_LINE_WIDTH, color=colors[colorStart + 1], linestyle="--")

	# Set axis bounds
	axes.set_xlim(0, xRangeMax)
	absoluteYMax = max(np.max(upper0), np.max(upper1)) # default to upper quartiles
	axes.set_ylim(0, absoluteYMax * 1.05)

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

	# Legend
	if inset:
		l = legend(PLOT_NAMES, loc=9, ncol=2)
		axes.axhline(y=absoluteYMax * 1.05, linewidth=4, color="black")
	else:
		l = legend(PLOT_NAMES, loc=0, ncol=2)
	l.get_frame().set_facecolor("0.9")
	l.get_frame().set_edgecolor("0.9")

# Plots an inset plot.
def plotInset(axes, colorStart, suffix, xRangeMax,
		median0, lower0, upper0, xRange0, yRange0, 
		median1, lower1, upper1, xRange1, yRange1,
		minimum0=None, maximum0=None, minimum1=None, maximum1=None):
	#fill_between(xRange0, lower0, upper0, alpha=0.25, linewidth=0, color=colors[colorStart])
	#fill_between(xRange1, lower1, upper1, alpha=0.25, linewidth=0, color=colors[colorStart + 1])
	#plot(xRange0, median0, linewidth=MAJOR_LINE_WIDTH, marker="v", markevery=(10), markersize=25.0, color=colors[colorStart])
	#plot(xRange1, median1, linewidth=MAJOR_LINE_WIDTH, marker="s", markevery=(10), markersize=25.0, color=colors[colorStart + 1])
	if maximum0 is not None:
		plot(xRange0, maximum0, linewidth=MINOR_LINE_WIDTH, color=colors[colorStart], linestyle="--")
	if maximum1 is not None:
		plot(xRange1, maximum1, linewidth=MINOR_LINE_WIDTH, color=colors[colorStart + 1], linestyle="--")
	axes.set_xlim(0, xRangeMax)
	plt.xticks([])
	axes.set_ylim(40000, 80000)
	plt.yticks([40000, 60000, 80000])
	axes.axhline(y=40000, linewidth=4, color="black")
	
	
# Plots the significance scores for a given data set
def plotSignificance(axes, colorToUse, markerToUse, data0, data1, absoluteMaxXRange):
	significanceValues = computeSignificance(data0, data1)
	for i in xrange(0, len(significanceValues), SIGNIFICANCE_SKIP):
		if (significanceValues[i] < significanceThreshold):
			plot(i, 0, marker=markerToUse, markersize=10, color=colorToUse)
	
	# Set axis bounds
	axes.set_xlim(0, absoluteMaxXRange)
	axes.set_ylim(-10, 10)

	# Disable tick labels
	axes.set_xticklabels([])
	axes.set_yticklabels([])
	plt.xticks([])
	plt.yticks([])

	# Indicate significance threshold
	plt.xlabel("p < " + significanceThresholdString, size=32)

# Obtain rewards data	
dataItems, xRanges = loadFiles(SUFFIX_SCORES)
data0rew = dataItems[0]
data1rew = dataItems[1]
xRange0rew = xRanges[0]
xRange1rew = xRanges[1]

# Obtain intrinsic data
dataItems, xRanges = loadFiles(SUFFIX_INTRINSIC)
data0int = dataItems[0]
data1int = dataItems[1]
xRange0int = xRanges[0]
xRange1int = xRanges[1]

# Data distribution
if PLOT_CI:
	median0rew, lower0rew, upper0rew = computeCIs(data0rew)
	median1rew, lower1rew, upper1rew = computeCIs(data1rew)
	median0int, lower0int, upper0int = computeCIs(data0int)
	median1int, lower1int, upper1int = computeCIs(data1int)
else:
	median0rew, lower0rew, upper0rew = computeQuartiles(data0rew)
	median1rew, lower1rew, upper1rew = computeQuartiles(data1rew)
	median0int, lower0int, upper0int = computeQuartiles(data0int)
	median1int, lower1int, upper1int = computeQuartiles(data1int)

# Minimum values
minimum0rew = computeMinimum(data0rew)
minimum1rew = computeMinimum(data1rew)
minimum0int = computeMinimum(data0int)
minimum1int = computeMinimum(data1int)

# Maximum values
maximum0rew, yRange0rew = computeMaximum(data0rew)
maximum1rew, yRange1rew = computeMaximum(data1rew)
maximum0int, yRange0int = computeMaximum(data0int)
maximum1int, yRange1int = computeMaximum(data1int)

# Maximum x-range
xRangeMax0rew = np.max(xRange0rew) - 1
xRangeMax1rew = np.max(xRange1rew) - 1
xRangeMax0int = np.max(xRange0int) - 1
xRangeMax1int = np.max(xRange1int) - 1

# Create the figure
fig = figure()

# Create the grid layout
gs = gridspec.GridSpec(24, 2)

# First plot (scores)
ax1 = fig.add_subplot(gs[5:19, 0])
plotData(ax1, 0, SUFFIX_SCORES, max(xRangeMax0rew, xRangeMax1rew),
	median0rew, lower0rew, upper0rew, xRange0rew, yRange0rew,
	median1rew, lower1rew, upper1rew, xRange1rew, yRange1rew,
	minimum0=minimum0rew, maximum0=maximum0rew, inset=True)

ax3 = fig.add_subplot(gs[0:5, 0])
plotInset(ax3, 0, SUFFIX_SCORES, max(xRangeMax0rew, xRangeMax1rew),
	median0rew, lower0rew, upper0rew, xRange0rew, yRange0rew,
	median1rew, lower1rew, upper1rew, xRange1rew, yRange1rew,
	minimum0=minimum0rew, maximum0=maximum0rew)

# Second plot (intrinsic scores)
ax2 = fig.add_subplot(gs[:19, 1])
plotData(ax2, 0, SUFFIX_INTRINSIC, max(xRangeMax0int, xRangeMax1int),
	median0int, lower0int, upper0int, xRange0int, yRange0int,
	median1int, lower1int, upper1int, xRange1int, yRange1int,
	minimum0=minimum0int, maximum0=maximum0int)

# Third plot (scores significance)
ax3 = fig.add_subplot(gs[23, 0])
plotSignificance(ax3, colors[1], "s", data0rew, data1rew, max(xRangeMax0rew, xRangeMax1rew))

# Fourth plot (intrinsic scores significance)
ax4 = fig.add_subplot(gs[23, 1])
plotSignificance(ax4, colors[1], "s", data0int, data1int, max(xRangeMax0int, xRangeMax1int))

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

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