import openai
import os
from task import Task


def read_api_key(v=3.5):
    if v == 3.5:
        with open("api_key.txt", "r") as f:
            return f.read()
    elif v == 4:
        with open("api_key_4.txt", "r") as f:
            return f.read()


def run_ARC_subset(encoding="word", delimiter="|", prompt_type="direct_grid_in_context_few_shot_cot", gpt_v="3.5"):

    with open("dataset/subset/simplest.txt") as f:
        lines = f.readlines()

    for file_num, file_name in enumerate(lines):

        print(f"Running task {file_num + 1}: {file_name} out of {len(lines)}")
        file_name = file_name.strip()

        task = Task("dataset/training/" + file_name)
        task.run_GPT(mode="grid", encoding=encoding, delimiter=delimiter, json_str="", gpt=gpt_v, task_type="2d",
                         prompt_type=prompt_type, temp=0)


def run_subset(data_path="dataset/generated/1d_move_1p/", abs="nbccg", gpt_v="3.5"):
    subset_name = data_path.strip("/").split("/")[-1]
    for ix, filename in enumerate(os.listdir(data_path)):
        print(f"Running task {ix + 1}: {filename} out of {len(os.listdir(data_path))}")
        task = Task(data_path + filename)
        task.abstraction = abs
        task.run_GPT(mode="image", encoding="object_json", delimiter="", json_str="", gpt=gpt_v, task_type="2d",
                             prompt_type="object_based_in_context_few_shot_cot_1d", temp=0, subset_name=subset_name)


if __name__ == "__main__":

    api_key = read_api_key()
    openai.api_key = api_key

    # runs ARC-subset
    # prompt types:
    #     object_based: object_based_few_shot, object_based_in_context_few_shot_cot
    #     direct_grid: direct_grid_few_shot, direct_grid_in_context_few_shot_cot
    #
    # encodings:
    #     object_based: "object_json", "object_descriptor", "object_json_w_edge", "object_descriptor_w_edge"
    #     direct_grid: "number", "word"
    #
    # delimiter:
    #     direct_grid: "", "|", ","

    run_ARC_subset(encoding="word", delimiter="|", prompt_type="direct_grid_in_context_few_shot_cot", gpt_v="3.5")

    # runs generated-subset
    # data_path is the path to the folder containing the generated tasks
    # abs is the ARGA abstraction to be used
    run_subset(data_path="dataset/generated/1d_move_1p/", abs="nbccg", gpt_v="3.5")
