# ObtainDataEmbedding.py
# imports
import openai
import pandas as pd

import tiktoken
from openai.embeddings_utils import get_embedding
import config
# set your API key
openai.api_key = "your openai api key"

# embedding model parameters
embedding_model = "text-embedding-ada-002"
embedding_encoding = "cl100k_base"  # this the encoding for text-embedding-ada-002
max_tokens = 8000  # the maximum for text-embedding-ada-002 is 8191

# load & inspect dataset
input_datapath = "data/attachmentchatdata-formated.csv"
df = pd.read_csv(input_datapath, index_col=0)
df = df[["userid", "chathistory", "avoide", "avoida", "avoidb", "avoidc", "avoidd", "anxietye", "anxietya", "anxietyb", "anxietyc", "anxietyd"]]
df = df.dropna()
df.head(2)

# Filter out chat transcripts that are too long to embed, estimate for the maximum number of words would be around 1638 words (8191 tokens / 5).
encoding = tiktoken.get_encoding(embedding_encoding)

df["n_tokens"] = df.chathistory.apply(lambda x: len(encoding.encode(x)))
df = df[df.n_tokens <= max_tokens]
len(df)

# Ensure you have your API key set in your environment per the README: https://github.com/openai/openai-python#usage

# This may take a few minutes
df["embedding"] = df.chathistory.apply(lambda x: get_embedding(x, engine=embedding_model))
df.to_csv("data/chat_transcripts_with_embeddings_and_scores.csv")


# Please replace "data/chat_transcripts.csv" with the path to your actual data file. Also, replace 'ChatTranscript', 'Attachment', 'Avoidance' with the actual column names of your chat transcripts and attachment scores in your data file.

# Also, remember to set the API key for OpenAI in your environment before running the get_embedding function.