import os
import json

auto_messages = [
    {
      "role": "system",
      "content": "You are a helpful assistant that can write SMT-LIB language that solves mathematical reasoning questions similarly to the examples that you will be provided."
    }]

inauto_messages = [
    {
      "role": "system",
      "content": "You are a helpful assistant that can write SMT-LIB language that solves mathematical reasoning questions similarly to the examples that you will be provided."
    }
]
sol_messages = [
    {
      "role": "system",
      "content": "Follow the given examples and answer the question."
    },
]

json_files = []
for root, dirs, files in os.walk("../data/train"):
    for file in files:
        if file.endswith('.json'):
            file_path = os.path.join(root, file)
            json_files.append(file_path)

for json_file in json_files:
    with open(json_file, 'r') as f:
        data = json.load(f)
        auto_messages += [{"role": "user", "content": f"Translate the natural language problem into SMT-LIB language: \"{data['problem']}\""}, 
                            {"role": "assistant", "content": data['formal problem']}]
        inauto_messages += [{"role": "user","content": f"Translate the math problem formulated with SMT-LIB back to a natural language problem. \"{data['formal problem']}\""}, 
                            {"role": "assistant","content": data['problem']}]
        sol_messages += [{"role": "user","content": data['problem']}, 
                            {"role": "assistant","content": data['solution'].replace("####", "The answer is")}]

print("auto messages length ==== ", (len(auto_messages)-1)/2)
print("inauto messages length ==== ", (len(inauto_messages)-1)/2)
print("sol messages length ==== ", (len(sol_messages)-1)/2)

# save
with open("./data/auto_examples.json", "w") as f:
    json.dump({"messages": auto_messages}, f, indent=2)
with open("./data/inauto_examples.json", "w") as f:
    json.dump({"messages": inauto_messages}, f, indent=2)
with open("./data/sol_examples.json", "w") as f:
    json.dump({"messages": sol_messages}, f, indent=2)