import shlex, subprocess, re
import os,re
from time import sleep
import pickle,argparse
import math
import time

dir_path = os.path.dirname(os.path.realpath(__file__))
def run_test(num_bundles, bundle_size, num_instances=1, num_cores=12):
	start_time = time.time()
	print("\nRunning test:")
	print("  Bundle count:\t{}".format(num_bundles))
	print("  Bundle size:\t{}".format(bundle_size))
	print("  Endpoints:\t{}".format(num_instances))
	print("  Cores:\t{}".format(num_cores))
	print()
	counters = []
	for i in range(num_instances):
		corelist = "0-{}".format(num_cores-1)
		command = "taskset -c {} bpcounter ipn:3.{} {}".format(corelist,i+1,num_bundles)
		args = shlex.split(command)
		p = subprocess.Popen(args,stdout=subprocess.PIPE, cwd=dir_path+"/3.bench.tcp")
		counters.append(p)
	sleep(1)
	for i in range(num_instances):
		command = "taskset -c {} bpdriver {} ipn:2.{} ipn:3.{} -{}".format(corelist,num_bundles,i+1,i+1,bundle_size)
		args = shlex.split(command)
		p = subprocess.Popen(args,stdout=subprocess.PIPE, cwd=dir_path+"/2.bench.tcp")

	s = 0
	count = 0
	for i,p in enumerate(counters):
		while True:
			l = p.stdout.readline().decode("utf-8")
			if re.search('received', l):
				pass

			res = re.search('(?<=Throughput \(Mbps\)\: )[0-9]*\.[0-9]*',l)
			if res:
				split = l.split(' ')
				num = float(split[2].strip())
				s+=num
				count+=1
				# print("endpoint {}: {} Mbps".format(i+1,num))
				break
			if l == "":
				break
	elapsed_time = time.time() - start_time
	print("time: {} s".format(round(elapsed_time,2)))
	print("Throughput: {} Mbps".format(s))
	if count == 0:
		return 0
	else:
		return s


parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('--graph', metavar='data_file',type=str, nargs=1)
args = parser.parse_args()

if args.graph:
	filename = args.graph[0]
	data = pickle.load(open(filename,'rb'))
	print(data)

	quit()

def restart_ion(num_cores):
	print("running killm...\r",end="")
	killm = subprocess.Popen(["killm"],stdout=subprocess.PIPE)
	killm.wait()
	sleep(1)
	print("running killm... done")
	print("starting ion nodes...\r",end="")

	corelist = "0-{}".format(num_cores-1)
	command = "taskset -c {} ./ionstart".format(corelist)
	args = shlex.split(command)
	p1 = subprocess.Popen(args,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,cwd=dir_path+"/2.bench.tcp")
	p2 = subprocess.Popen(args,stdout=subprocess.PIPE,stderr=subprocess.STDOUT,cwd=dir_path+"/3.bench.tcp")
	p1.wait()
	p2.wait()
	print("starting ion nodes... done")
	sleep(3)

# args = shlex.split(command)
# print("running killm...\r",end="")
# killm = subprocess.Popen(["killm"],stdout=subprocess.PIPE)
# killm.wait()
# sleep(1)
# print("running killm... done")
# print("starting ion nodes...\r",end="")
# p1 = subprocess.Popen(["./ionstart"],stdout=subprocess.PIPE,stderr=subprocess.STDOUT,cwd=dir_path+"/2.bench.tcp")
# p2 = subprocess.Popen(["./ionstart"],stdout=subprocess.PIPE,stderr=subprocess.STDOUT,cwd=dir_path+"/3.bench.tcp")
# p1.wait()
# p2.wait()

# print("starting ion nodes... done")
# sleep(3)
# throughput = run_test(50000, 10000, 10)
# print("Throughput: {} Mbps".format(throughput))

tests = {}

tests = pickle.load(open("tests.p","rb"))

bundle_sizes = [4000,8000,16000,32000,64000,128000,256000,512000]
bundle_count = 50000
runs = 1
instances = [1,2,3,4,5,6,7,8,9,10,11,12]
print(instances)
for cores in range(1,13):
	restart_ion(cores)
	core_instances = {}
	if cores in tests:
		core_instances = tests[cores]
	for num_instances in instances:
		bundle_tests = {}
		if num_instances in core_instances:
			bundle_tests = core_instances[num_instances]

		scaled_bundle_count = int(bundle_count/num_instances)
		for bundle_size in bundle_sizes:
			if bundle_size in bundle_tests and bundle_tests[bundle_size] != 0:
				continue
			s = 0
			for i in range(runs):
				s += run_test(scaled_bundle_count,bundle_size,num_instances,cores)

			bundle_tests[bundle_size] = s/runs
		core_instances[num_instances] = bundle_tests
		tests[cores] = core_instances
		pickle.dump(tests,open('tests.p','wb'))
	tests[cores] = core_instances