import chromadb
import os
import pandas as pd
from chromadb.utils import embedding_functions

client = chromadb.PersistentClient()
em = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2", device="cuda")
collection = client.get_or_create_collection("eligibility_criteria_collection", metadata={"hnsw:space": "cosine"}, embedding_function=em) # cosine is the default

# Create an empty list to store filenames without extension
folder_path='../data/preprocessed_data/clinical_trials'
files = os.listdir(folder_path)
filenames_without_extension = []
for file in files:
    filename, file_extension = os.path.splitext(file)
    filenames_without_extension.append(filename)

def add_documents_collection(filenames, collection, folder_path):
    for f in filenames:
        print(f)
        df = pd.read_csv(f"{folder_path}/{f}.csv")
        df = df.dropna()
        for index, row in df.iterrows():
            text = row["sentence"]
            nct_id = row["id"]
            metadatas = {
                "criteria": row["criteria"],  # Replace "criteria" with the actual column name
                "sub-criteria": row["sub_criteria"]  # Replace "sub-criteria" with the actual column name
            }
            # Add the document to the collection
            collection.add(
                documents=[text],
                ids=["{}-{}".format(nct_id, index + 1)],
                metadatas=metadatas
            )
add_documents_collection(filenames=filenames_without_extension, collection=collection, folder_path=folder_path);
print(collection.count())