import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patheffects as PathEffects

N_processes = 5
N_threads = 5

# Load measured timings
times = np.empty((N_processes, N_threads))
for fname in os.listdir('timings'):
    values = open(f'timings/{fname}').read().split()
    n_processes = int(values[0])
    n_threads = int(values[1])
    dt = float(values[2])
    times[n_processes-1][n_threads-1] = dt
print(times)

""" Plot measured time"""
fig_time, axs_time = plt.subplots()
im = axs_time.imshow(times.T, origin='lower')
axs_time.set_title('Computation time')
fig_time.colorbar(im, ax=axs_time, label='Measured computation time (s)')

""" Plot speedup """
workers = np.arange(N_processes + 1)[:, None] * np.arange(N_threads + 1)
speedup = times[0, 0] / times

fig_speedup, axs_speedup = plt.subplots()
im = axs_speedup.imshow(speedup.T, origin='lower')
axs_speedup.set_title('Computation speed-up')
fig_speedup.colorbar(im, ax=axs_speedup, label='Speed-up')

# Set same style for both plots
for axs, data in zip([axs_time, axs_speedup], [times, speedup]):
    axs.set_xlabel('# processes')
    axs.set_ylabel('# threads')
    axs.set_xticks(np.arange(N_processes))
    axs.set_xticklabels(np.arange(N_processes)+1)
    axs.set_yticks(np.arange(N_threads))
    axs.set_yticklabels(np.arange(N_threads)+1)

    for i in range(N_processes):
        for j in range(N_threads):
            txt = axs.text(i, j, f'{data[i, j]:.2f}', fontsize=10, color='w',
                        ha='center', va='center', fontweight='bold')
            txt.set_path_effects([PathEffects.withStroke(linewidth=0.5, foreground='k')])
    axs.spines[['right', 'top']].set_visible(False)

# Save plots
fig_time.savefig('time.png', dpi=300)
fig_speedup.savefig('speedup.png', dpi=300)