import numpy as np
from loguru import logger
import time
import os

# start with 3_000 vectors to keep things small 
total_vectors_num = 3_000_000_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(vectors_file:np.array, query_vectors:np.array) -> list[np.array]:
    total_products_computed = 0

    dot_products = []
    for v in vectors_file:
        for qv in query_vectors:
            dot_product = v @ qv
            dot_products.append(dot_product)
            total_products_computed += 1
            if total_products_computed % 100000 == 0:
                 logger.info(f"Total vectors processed:{total_products_computed}")

    return dot_products


# Generate initial vectors and query vectors and write to disk
doc_vectors = generate_random_vectors(total_vectors_num)
query_vectors = generate_random_vectors(query_vectors_num)
np.save('vectors.npy', doc_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(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)}")