import openai
from task_fewshot import Task
import time
from random import randrange


def read_api_key(gpt_version="3.5"):
    version_map = {"3.5":"api_key_project.txt","4":"api_key.txt"}
    with open(version_map[gpt_version], "r") as f:
        return f.read()


if __name__ == "__main__":

    #gpt_version = "3.5"
    gpt_version = "4"
    
    #task_type="2d"
    task_type="1d"

    api_key = read_api_key(gpt_version=gpt_version)
    print(api_key)
    openai.api_key = api_key.rstrip()

    custom_runs = [21]
    for i in range(4,50):
    #for i in range(50):
    #for i in custom_runs:
        print(str(i)+".json")

        # 1D basic moves
        # 1 pixel right
        #task = Task(filepath="dataset/1d_tasks/basic_move_1/bmove_1p_"+str(i)+".json", db_path="fewshot_results/1d_bmove_1_solutions.csv")
        #task = Task(filepath="dataset/1d_tasks/basic_move_2/bmove_2p_"+str(i)+".json", db_path="fewshot_results/1d_bmove_2_solutions.csv")
        # +3 pixel right
        #task = Task(filepath="dataset/1d_tasks/basic_move/bmove_"+str(i)+".json", db_path="fewshot_results/1d_bmove_solutions.csv")
        # dynamic pixels
        #task = Task(filepath="dataset/1d_tasks/basic_move_dp/bmove_dp_"+str(i)+".json", db_path="fewshot_results/1d_bmove_dp_solutions.csv")
        #task = Task(filepath="dataset/1d_tasks/basic_move_2p_p/bmove_2p_p_"+str(i)+".json", db_path="fewshot_results/1d_bmove_2_p_solutions.csv")

        # 1D basic fill
        #task = Task(filepath="dataset/1d_tasks/basic_fill/bfill_"+str(i)+".json", db_path="fewshot_results/1d_bfill_solutions.csv")

        # Expands x3
        task = Task(filepath="dataset/1d_tasks/basic_fill_3/bfill_3_"+str(i)+".json", db_path="fewshot_results/1d_bfill_3_solutions.csv")

        # Reverse versions
        #task = Task(filepath="dataset/1d_tasks/basic_fill_reverse/bfill_r_"+str(i)+".json", db_path="fewshot_results/1d_bfill_r_solutions.csv")


        # 1D flip
        #task = Task(filepath="dataset/1d_tasks/1d_flip_single/1dflip_s_"+str(i)+".json", db_path="fewshot_results/1d_flip_single_solutions.csv")

        # 1D mirror
        #task = Task(filepath="dataset/1d_tasks/1d_mirror_single/1dmirror_s_"+str(i)+".json", db_path="fewshot_results/1d_mirror_single_solutions.csv")

        # From here, all are 32 pixel grids

        # 1D denoising
        #task = Task(filepath="dataset/1d_tasks/1d_1c_denoising_single/1d1cdenoi_s_"+str(i)+".json", db_path="fewshot_results/1d_1c_denoising_single_solutions.csv")
        #task = Task(filepath="dataset/1d_tasks/1d_mc_denoising_single/1dmcdenoi_s_"+str(i)+".json", db_path="fewshot_results/1d_mc_denoising_single_solutions.csv")

        # Pattern copy
        #task = Task(filepath="dataset/1d_tasks/1d_pcopy_1c/1dpcopy_1c_"+str(i)+".json", db_path="fewshot_results/1d_pcopy_1c_solutions.csv")
        #task = Task(filepath="dataset/1d_tasks/1d_pcopy_mc/1dpcopy_mc_"+str(i)+".json", db_path="fewshot_results/1d_pcopy_mc_solutions.csv")

        # Recolors
        #task = Task(filepath="dataset/1d_tasks/1d_recolor_oe/1drec_oe_"+str(i)+".json", db_path="fewshot_results/1d_recolor_oe_solutions.csv")
        #task = Task(filepath="dataset/1d_tasks/1d_recolor_cnt/1drec_cnt_"+str(i)+".json", db_path="fewshot_results/1d_recolor_cnt_solutions.csv")
        #task = Task(filepath="dataset/1d_tasks/1d_recolor_cmp/1drec_cmp_"+str(i)+".json", db_path="fewshot_results/1d_recolor_cmp_solutions.csv")

        # 1D scale
        #task = Task(filepath="dataset/1d_tasks/1d_scale_dp/1dscale_dp_"+str(i)+".json", db_path="fewshot_results/1d_scale_dp_solutions.csv")

        # LEVEL-2s
        #task = Task(filepath="dataset/1d_tasks/fill_n_move/fmove_"+str(i)+".json", db_path="fewshot_results/1d_fmove_solutions.csv")
        #task = Task(filepath="dataset/1d_tasks/1d_move_1p_recolor/1dm1p_rec_"+str(i)+".json", db_path="fewshot_results/1d_move_1p_recolor_solutions.csv")


        # Best performance on ARC-simple-50 is 5/50, word+pipe
        task.run_GPT(mode="grid", encoding="word", delimiter="|", json_str="", gpt=gpt_version, task_type=task_type)
        #time.sleep(randrange(2,4))

        #break

    print(gpt_version,task_type,task.db_path)

