import os
import timeit
import numpy as np
from datetime import datetime
import time

# Timestamp that will be put in the file name
timestamp = datetime.now().strftime("%H%M%S%f")

# Get the environment variable for threads
threads = os.getenv('OMP_NUM_THREADS')

# A relatively large matrix to work on
n = 5_000
x = np.random.random(size=(n, n))

print(f"We are executed with OMP_NUM_THREADS={threads} for {n=}")

# Measure the time required for matrix multiplication
start_time = time.time()
y = x @ x  # The heavy compute
elapsed_time = time.time() - start_time

print(f'Time used for matrix multiplication: {elapsed_time:.2f} s')

# Check if timings folder exists
if not os.path.isdir('timings/'): os.mkdir('timings')

# IO: Save the timing to a unique txt file
with open(f'timings/{threads}_threads_t{timestamp}.txt', 'w') as file:
    file.write(f'{threads},{elapsed_time:.6f}')