import numpy as np
import matplotlib.pyplot as plt
import timeit

nthreads = [1, 2, 4, 8, 16]
sizes = [100, 400, 1600, 6400, 12800, 25600, 51200]

thread_timings = [] 
for t in nthreads:
    print(f"Measure timing for {t} thread(s)")
    size_timings = []
    for s in sizes:
        print(f"--- Measure timing for size {s}")
        command = f"subprocess.run(\"mpirun --use-hwthread-cpus -n {t} python3 exercise_04.py {s}\", shell = True)"
        size_timings.append(timeit.timeit(command, setup = "import subprocess", number = 1) * 1000)
    thread_timings.append(size_timings)

plt.plot(sizes, thread_timings[0], "green", label=f"{nthreads[0]}")
plt.plot(sizes, thread_timings[1], "blue", label=f"{nthreads[1]}")
plt.plot(sizes, thread_timings[2], "purple", label=f"{nthreads[2]}")
plt.plot(sizes, thread_timings[3], "red", label=f"{nthreads[3]}")
plt.plot(sizes, thread_timings[4], "orange", label=f"{nthreads[4]}")

plt.title("Matrix-Vector-Multiplication time/threads comparison")
plt.xlabel("Sizes")
plt.ylabel("Time (ms)")
plt.loglog()
plt.legend(title = "# of Threads")

plt.show()