import pickle
import numpy as np
from timeit import default_timer as timer
import matplotlib.pyplot as plt

from acceleration import *
import analysis.utils as utils


timescale = 10 # ms, approximate time per attempt
massive = True # whether all particles are massive

functions = [bruteForce, bruteForceNumba, bruteForceNumbaOptimized, bruteForceCPP, barnesHutCPP]
labels = ['Brute force', 'Brute force Numba', 'Brute force Numba opt.', 'Brute force [C++]', 'Barnes-Hut [C++]']
styles = ['solid', 'dashed', 'solid', 'dashed', 'solid']
colours = ['cornflowerblue', 'cornflowerblue', 'darkorange', 'darkorange', 'black']
if not massive: #do not evaluate barnes hut
	functions, labels, styles, colours = functions[:-1], labels[:-1], styles[:-1], colours[:-1]


timess = []

for i in range(len(functions)):
	n = 10 #First data point
	# Run attempts with multiple iterations each
	iterations = 1000
	attempts = 6
	timess.append([])
	while(True):
		# Evaluate the performance
		localTimes = []
		for _ in range(attempts):
			# Initialize random data
			r_vec = np.random.rand(n, 3).astype(np.float64)
			if massive:
				m_vec = np.random.rand(n).astype(np.float64)
			else:
				m_vec = np.zeros(n).astype(np.float64)
				m_vec[0] = 1
			start = timer()
			for _ in range(iterations):
				functions[i](r_vec, m_vec)
			stop = timer()
			localTimes.append((stop - start)*1000/iterations) #in ms
		# Calculate mean and error
		# Discard high times since these are likely due to interferences
		localTimes = np.sort(localTimes)
		localTimes = localTimes[:attempts-2]
		print(localTimes, iterations)
		error = np.std(localTimes)/np.sqrt(attempts-1)
		timess[i].append([n, np.mean(localTimes), error])

		# Recalculate how many iterations would have been sensible
		iterations = int(timescale/np.mean(localTimes))
		if iterations > 1000: iterations = 1000
		if iterations < 1: break;

		# Increase the number of points
		n *= 2

# Plot the results
f, ax = plt.subplots(1, 1,  figsize=(6, 4), sharey=True)

for times, label, style, colour in zip(timess, labels, styles, colours):
	times = np.array(times)
	ax.errorbar(times[:,0], times[:,1], yerr=times[:,2], 
		label=label, marker='+', linestyle=style, c=colour)

ax.loglog()
if massive:
	utils.setAxes(ax, x='Number of massive particles', y='Time / ms', 
	xcoords=(.7, -0.1), ycoords=(-0.1,.7))
else:
	utils.setAxes(ax, x='Number of massless particles', y='Time / ms', 
	xcoords=(.7, -0.1), ycoords=(-0.1,.7))
utils.stylizePlot([ax])
ax.legend()
ax.set_xlim((10, None))
plt.show()


file = open('data/performance.pickle', "wb")
pickle.dump(timess, file)