import pickle
# from openai import OpenAI
import openai
from tqdm import tqdm
import json
import numpy as np

# client = OpenAI("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")

openai.api_key = "sk-proj-fVs1blmKr8rJn1hciiEkT3BlbkFJIUrppQa46tDswSa8tOuh"

def get_embedding(text, model="text-embedding-3-small"):
   text = text.replace("\n", " ")
   return openai.Embedding.create(input = [text], model=model).data[0].embedding


data_path = "/home/czr/Agent-Driver/data/memory/database.pkl"
data_sample_path = "/home/czr/Agent-Driver/data/finetune/data_samples_train.json"

data = pickle.load(open(data_path, 'rb'))
with open(data_sample_path, 'r') as file:
    data_samples = json.load(file)#[:20000]

data_sample_dict = {}
data_sample_val_dict = {}
for data_sample in data_samples:
    data_sample_dict[data_sample["token"]] = data_sample

embeddings_database = []

for idx, token in tqdm(enumerate(data), desc="Embedding original database with OpenAI ADA model"):
    if idx >= 10000:
        break
    # print("data[token]", data_sample_dict[token])
    # input()
    # try:
    working_memory = {}
    working_memory["ego_prompts"] = data_sample_dict[token]["ego"]
    perception = data_sample_dict[token]["perception"]
    working_memory["perception"] = perception
    
    text = working_memory["ego_prompts"] + " " + working_memory["perception"]
    
    try:
        while True:
            embedding = get_embedding(text)
            break
    except:
        continue

    embeddings_database.append(embedding)
    
    print("embeddings_database", np.shape(embeddings_database))
    with open("data/memory/AgentDriver_database_embeddings_5000.pkl", "wb") as f:
        pickle.dump(embeddings_database, f)

    # except:
    #     continue