import openai
import os
from dotenv import load_dotenv

load_dotenv()

openai.api_key = os.getenv("API_KEY")


#Takes in a sentence and returns a list of dicts consisiting of key-value pairs of masked words and lists of the possible replacements
def predict_masked_words(sentence, n_suggestions=5):

    prompt = (
          f"Given a sentence with masked words, masked word can be one or more than one, indicated by [MASK], generate {n_suggestions} possible words to fill each mask. "
          "Return the results as a list of dictionaries, where each dictionary key is a masked word and its value is a list of 5 potential words to fill that mask.\n\n"
          "Example input: \"The [MASK] fox [MASK] over the [MASK] dog.\"\n\n"
          "Example output:\n"
          "[\n"
          "  {\n"
          "    \"[MASK]1\": [\"quick\", \"sly\", \"red\", \"clever\", \"sneaky\"]\n"
          "  },\n"
          "  {\n"
          "    \"[MASK]2\": [\"jumped\", \"leaped\", \"hopped\", \"sprang\", \"bounded\"]\n"
          "  },\n"
          "  {\n"
          "    \"[MASK]3\": [\"lazy\", \"sleeping\", \"brown\", \"tired\", \"old\"]\n"
          "  }\n"
          "]\n\n"
          "Example input: \"The [MASK] [MASK] ran swiftly across the [MASK] field.\"\n\n"
          "Example output:\n"
          "[\n"
          "  {\n"
          "    \"[MASK]1\": [\"tall\", \"fierce\", \"young\", \"old\", \"beautiful\"]\n"
          "  },\n"
          "  {\n"
          "    \"[MASK]2\": [\"lion\", \"tiger\", \"horse\", \"cheetah\", \"deer\"]\n"
          "  },\n"
          "  {\n"
          "    \"[MASK]3\": [\"green\", \"wide\", \"sunny\", \"open\", \"empty\"]\n"
          "  }\n"
          "]\n\n"
          "Example input: \"It was a [MASK] day when the train arrived at the station.\"\n\n"
          "Example output:\n"
          "[\n"
          "  {\n"
          "    \"[MASK]1\": [\"sunny\", \"rainy\", \"cloudy\", \"foggy\", \"stormy\"]\n"
          "  },\n"
          "]\n\n"
          "Now, please process the following sentence:\n"
          f"{sentence}"
      )


    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=100,
        n=1,
        stop=None,
        temperature=0.7
    )

    print(response['choices'][0]['message']['content'])


# sentence = "Evacuations and storm [MASK] began on Sunday night as forecasters projected that Hurricane Dorian would hit into Florida’s west coast on Wednesday as a major hurricane packing life-threatening winds and storm surge."
# predict_masked_words(sentence, n_suggestions=5)