import time

import numpy as np
import cupy as cp

from matplotlib import pyplot as plt

def warmup_cupy():
	a = cp.random.random((10,10))
	b = cp.random.random((10,10))
	_ = cp.dot(a, b)

def calc_cpu(a: np.array, b: np.array):
	return np.dot(a, b)

def calc_gpu(a: np.array, b: np.array):
	return cp.dot(a, b)

def perf_test_2():
	### Numpy and CPU
	print("Creation")
	s = time.time()
	x_cpu = np.ones((1000, 1000, 1000))
	e = time.time()
	time_cpu = e - s
	s = time.time()
	x_gpu = cp.ones((1000, 1000, 1000))
	cp.cuda.Stream.null.synchronize()
	e = time.time()
	time_gpu = e - s
	print(f"Speedup: {time_cpu/time_gpu}")

	### Numpy and CPU
	print("Multiply with constant")
	s = time.time()
	x_cpu *= 5
	e = time.time()
	time_cpu = e - s
	s = time.time()
	x_gpu *= 5
	cp.cuda.Stream.null.synchronize()
	e = time.time()
	time_gpu = e - s
	print(f"Speedup: {time_cpu/time_gpu}")

	### Numpy and CPU
	print("Multiply with constant, Square, Square")
	s = time.time()
	x_cpu *= 5
	x_cpu *= x_cpu
	x_cpu += x_cpu
	e = time.time()
	time_cpu = e - s
	s = time.time()
	x_gpu *= 5
	x_gpu *= x_gpu
	x_gpu += x_gpu
	cp.cuda.Stream.null.synchronize()
	e = time.time()
	time_gpu = e - s
	print(f"Speedup: {time_cpu/time_gpu}")

def perf_test_1():
	n_values = []
	np_times = []
	cp_times = []
	ratio_times = []

	for N in range(100, 5100, 100):
		a_np = np.random.rand(1024, N)
		b_np = np.random.rand(N, 1024)

		a_cp = cp.asarray(a_np)
		b_cp = cp.asarray(b_np)

		start_time = time.time()
		c_np = calc_cpu(a_np, b_np)
		end_time = time.time()
		numpy_time = end_time - start_time

		start_time = time.time()
		c_cp = calc_gpu(a_cp, b_cp)
		c_np = cp.asnumpy(c_cp)
		end_time = time.time()
		cupy_time = end_time - start_time

		n_values.append(N)
		np_times.append(numpy_time)
		cp_times.append(cupy_time)
		ratio_times.append(numpy_time/cupy_time)

	plt.plot(n_values, np_times, label='numpy time')
	plt.plot(n_values, cp_times, label='cupy time')
	plt.plot(n_values, ratio_times, label='numpy/cupy time')
	plt.xlabel("Matrix size [N]")
	plt.ylabel("Time [s]")
	plt.title("Performance numpy vs cupy")
	plt.grid()
	plt.legend()
	plt.show()

def main():
	warmup_cupy()
	perf_test_2()

if __name__ == "__main__":
	main()