import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
import numpy as np
import os
import sys

def wykres(threadsy,capacitiesy,wyniki,iteracje,atrybut):
	kolory=['red','green','blue']
	alfy=[0.1,0.1,0.8]
	fig = plt.figure()
	ax = fig.add_subplot(111, projection='3d')
	ax.view_init(30,-150)
	for i,iteracja in enumerate(iteracje):
		X, Y = np.meshgrid(threadsy, capacitiesy) #X,Y are arrays of same size, 2D, with separated indexes
		Z=np.empty(X.shape)
		Z.fill(np.nan) #instead of nan, one could set -10 to clearly see on the plot which values are missing
		for wynik in wyniki:
			threads,capacity,slownik=wynik
			if slownik['iter']==iteracja:
				print threads,capacity,slownik[atrybut] if atrybut in slownik else np.nan
				threadsindex=threadsy.index(threads)
				capacitiesindex=capacitiesy.index(capacity)
				Z[capacitiesindex,threadsindex]=slownik[atrybut] if atrybut in slownik else np.nan
		print Z
		ktore=len(kolory)-len(iteracje)+i
		ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color=kolory[ktore], linewidth=1, alpha=alfy[ktore],edgecolor=(0,0,0,alfy[ktore]))
	ax.set_xlabel('threads')
	ax.set_ylabel('capacity')
	ax.set_zlabel(atrybut)
	ax.set_xticks(threadsy)
	ax.set_yticks(capacitiesy)
	ax.set_zlim((0,ax.get_zlim()[1])) #scale from (forced) zero to auto
	plt.savefig("wykres_%d_%s.png" % (iteracja,atrybut))
	#plt.show() #if you wanted to rotate the chart interactively 
	plt.close()

	



####################################################### main #######################################################

threadsy=set() #used to store values of 'threads' that are in use. This is helpful because later a single multi-dimensional numpy array is used to keep all results
capacitiesy=set() #used to store values of 'capacity' that are in use
wyniki=[] #list of all results

os.chdir(sys.argv[1]) #files with plots will also be created in that directory
for plik in os.listdir('.'):
	if plik.endswith(".out") and plik.count('_')==2:
		podzielone=plik.replace('.','_').split('_')
		threads=int(podzielone[1])
		capacity=int(podzielone[2])
		threadsy.add(threads)
		capacitiesy.add(capacity)
		plik = open("test_%d_%d.out" % (threads,capacity), "r") #should be the same name as the original file
		slowniki = []
		for linia in plik:
			linia=linia.strip()
			if linia.startswith("[INFO] Script::Message - iter="):
				slownik={}
				for pole in linia.split(' '):
					if '=' in pole:
						pole=pole.strip(',')
						pole=pole.split('=')
						slownik[pole[0]]=float(pole[1])
				print slownik
				slowniki.append(slownik)
				wyniki.append((threads,capacity,slownik))
		plik.close()

threadsy=sorted(list(threadsy))
capacitiesy=sorted(list(capacitiesy))
print threadsy
print capacitiesy

wykres(threadsy,capacitiesy,wyniki,[3.0,6.0,9.0],'simsteps')
wykres(threadsy,capacitiesy,wyniki,[3.0,6.0,9.0],'evals')
wykres(threadsy,capacitiesy,wyniki,[3.0,6.0,9.0],'migrations')
wykres(threadsy,capacitiesy,wyniki,[4.0,9.0],'fit_avg')
wykres(threadsy,capacitiesy,wyniki,[4.0,9.0],'fit_max')

#draws simple 2D plots of progress in consecutive iter's:
#x=[]
#y1=[]
#y2=[]
#for s in slowniki:
#	x.append(s['iter'])
#	y1.append(s['avg_fit'])
#	y2.append(s['max_fit'])
#plt.plot(x,y1)
#plt.plot(x,y2)
#plt.savefig("avg_%d_%d.png" % (threads,capacity))
#plt.close()