import numpy as np
from loguru import logger
import time

total_vectors_num = 3_000
query_vectors_num = 1_000

def generate_random_vectors(num_vectors:int)-> np.array:
    logger.info(f"Generating {num_vectors} vectors...")
    rng = np.random.default_rng() 
    vectors = rng.random((num_vectors, 768))
    return vectors

def get_dot_products_vectorized(vectors_file:np.array, query_vectors:np.array):
    dot_products = vectors_file @ query_vectors.T
    return dot_products.flatten() # collapse into single dim

# Generate initial vectors and query vectors and write to disk
ram_vectors = generate_random_vectors(total_vectors_num)
query_vectors = generate_random_vectors(query_vectors_num)
np.save('vectors.npy', ram_vectors)

# Load vectors from disk 
logger.info("Loading file from disk...")
vectors_file = np.load('vectors.npy')

start_time = time.time()
logger.info("Getting dot products...")
results =  get_dot_products_vectorized(vectors_file, query_vectors)
end_time = time.time()

logger.info(f"Execution time: {end_time - start_time:.4f} seconds")
logger.info(f"Number of dot products computed: {len(results)}")