from api import GPT, Example, set_openai_key
set_openai_key('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')

# gpt = GPT(input_prefix="", input_suffix="", output_prefix="", output_suffix="")
gpt = GPT()

def r(prompt):
    return gpt.get_top_reply(prompt)

def a(input, output):
    gpt.add_example(Example(input, output))


with open("spatial_examples.txt", 'r') as f:
    se = f.read()
s = list(map(lambda x: x + '\n', se.split('\n\n')))
for ss in s[:-1]:
    world, answer = ss.split("Answer: ")
    world += "Answer: "
    a(world, answer)

from generate_examples import SpatialExampleGenerator
from tqdm import tqdm
import time

seg = SpatialExampleGenerator(min_size=2, max_size=10)
pos, neg = 0, 0
valid = 0
# for _ in tqdm(range(100)):
#     try:
#         test = seg.generate()
#         world, asnwer = test.split("Answer: ")
#         world += "Answer: "
#         # print(world)
#         pred = r(world)
#         # print(pred)
#         # print("gold:", answer)
#         if "left" in pred or "right" in pred:
#             valid += 1
#         if answer.lower() in pred:
#             pos += 1
#         else:
#             neg += 1
#         time.sleep(1)
#     except Exception as e:
#         print(e)
#         pass

print(valid)
print(pos + neg)
# print(pos / (pos + neg))