import matplotlib
matplotlib.use('macosx')
import matplotlib.pyplot as plt
import numpy as np
import requests
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--max_workers', type=int, required=True)
parser.add_argument('--server_addr', type=str, required=True)
parser.add_argument('--repo', type=str, required=True)
args = parser.parse_args()

SERVER_ADDR = args.server_addr
REPO = args.repo
MAX_WORKERS = args.max_workers
workers = [2 ** i for i in range(int(np.log2(MAX_WORKERS) + 1))]

times = {}

for data_parallelisation_strategy in ['files', 'commits']:
  for work_distribution_strategy in ['work_stealing', 'work_pushing']:
    times[(data_parallelisation_strategy, work_distribution_strategy)] = []
    for num_workers in workers:
        r = requests.get('{0}/?git_url={1}&max_workers={2}&data_parallelisation_strategy={3}&work_distribution_strategy={4}'.format(SERVER_ADDR, REPO, num_workers, data_parallelisation_strategy, work_distribution_strategy))
        res = r.json()

        times[(data_parallelisation_strategy, work_distribution_strategy)].append(res['time_taken'])

max_time = max(map(max, times.values()))
min_time = min(map(min, times.values()))

print times

plt.plot(workers, times[('files', 'work_stealing')], 'r-', label='files | work stealing')
plt.plot(workers, times[('files', 'work_pushing')], 'g-', label='files | work pushing')
plt.plot(workers, times[('commits', 'work_stealing')], 'b-', label='commits | work stealing')
plt.plot(workers, times[('commits', 'work_pushing')], 'm-', label='commits | work pushing')
plt.axis([0, MAX_WORKERS + 1, 0, 1.1 * max_time])
plt.xlabel('Num workers')
plt.ylabel('Time (seconds)')
plt.legend(loc=1, title='Data parallelisation | Work distribution')
plt.title('Argon repo execution times - Macbook Air - Num cores = 4')
# plt.show()
plt.savefig('images/argon_mac.png')

plt.clf()
plt.cla()
plt.close()

plt.plot(workers, map(lambda t: times[('files', 'work_stealing')][0] / t, times[('files', 'work_stealing')]), 'r-', label='files | work stealing')
plt.plot(workers, map(lambda t: times[('files', 'work_pushing')][0] / t, times[('files', 'work_pushing')]), 'g-', label='files | work pushing')
plt.plot(workers, map(lambda t: times[('commits', 'work_stealing')][0] / t, times[('commits', 'work_stealing')]), 'b-', label='commits | work stealing')
plt.plot(workers, map(lambda t: times[('commits', 'work_pushing')][0] / t, times[('commits', 'work_pushing')]), 'm-', label='commits | work pushing')
plt.axis([0, MAX_WORKERS + 1, 0, MAX_WORKERS])
plt.xlabel('Num workers')
plt.ylabel('Parallel Speedup')
plt.legend(loc=1, title='Data parallelisation | Work distribution')
plt.title('Argon repo parallel speedup - Macbook Air - Num cores = 4')
# plt.show()
plt.savefig('images/argon_mac_speedup.png')

