import numpy as np
import time
import sys
import os

verbose=False
nthreads=os.environ['MKL_NUM_THREADS']
N = int(sys.argv[1])
# Create two large random matrices
a = np.random.randn(N, N)
b = np.random.randn(N)
Mem = (a.nbytes + b.nbytes)


t1 = time.time()
# That's the expensive call:
np.linalg.solve(a, b)
dt1  = time.time() - t1


# matrix vector mult
nrep=100
t2 = time.time()
for i in range(nrep):
    b= a@b

dt2 = (time.time()-t2)/1000

if verbose == True:
    print(f"Number of threads used by MKL: {nthreads}")
    print(f"Solving A x = b, with A a ({N},{N}) matrix and b a {N} vector")
    print(f"Memory used {(a.nbytes + b.nbytes)*1e-6}MB")
    print(f"time to solve A x = b {dt1} seconds")
    print(f"Time to 1 matrix vector multiplication {dt2} seconds")

#print("N,nthreads,Mem,linalg time, matrix vector time")
print(f"{N},{nthreads},{Mem},{dt1},{dt2}")