from Generator import *

# 2D ARC, horizontal to vertical
class Gen_h_to_v(Generator):

    def __init__(self, save_folder, file_prefix, min_len=12, max_len=33):
        super().__init__(save_folder, file_prefix, min_len, max_len)

    def gen_fix_hv(self,origin_task_path,fix_idx=[]):
        task_1d = json.load(open(origin_task_path))

        # rotate io pairs
        cnt=0
        for demo in task_1d['train'] + task_1d['test']:
            if cnt in fix_idx:
                demo = self.helper_to_vertical(demo)
            cnt+=1

        print(task_1d)
        filepath = os.path.basename(origin_task_path).replace(".json","_h.json")
        print(filepath)
        #with open("../dataset/arc_hv_fixed/"+filepath, 'w') as f:
        #        json.dump(task_1d, f)

    def gen_h_to_v(self, h_task_folder, v_task_folder,degree=3):
        for entry in os.scandir(h_task_folder):
            task_1d = json.load(open(entry.path))

            # rotate
            for demo in task_1d['train'] + task_1d['test']:
                for i in range(degree):
                    demo = self.helper_to_vertical(demo)

                print(demo['output'])

            print(task_1d)
            filepath = v_task_folder + os.path.basename(entry.path).replace("_h","_v")
            print(filepath)
            with open(filepath, 'w') as f:
                json.dump(task_1d, f)

gen1=Gen_h_to_v(save_folder="../dataset/hv_tasks/test",file_prefix="hv_test")

# Fix exisiting ARC task - align all orientation
#gen1.gen_fix_hv(origin_task_path="../dataset/training/05f2a901.json",fix_idx=[0,2])
#gen1.gen_fix_hv(origin_task_path="../dataset/training/1e0a9b12.json",fix_idx=[0,1,2,3])
#gen1.gen_fix_hv(origin_task_path="../dataset/training/25ff71a9.json",fix_idx=[0,1,2,3,4])
#gen1.gen_fix_hv(origin_task_path="../dataset/training/d037b0a7.json",fix_idx=[0,1,2,3])
#gen1.gen_fix_hv(origin_task_path="../dataset/training/3906de3d.json",fix_idx=[0,1,2,3])
#gen1.gen_fix_hv(origin_task_path="../dataset/training/6855a6e4.json",fix_idx=[0,3])
#gen1.gen_fix_hv(origin_task_path="../dataset/training/f8a8fe49.json",fix_idx=[0,3])


# Rotate the generated datasets from horizontal to vertical
#gen1.gen_h_to_v(h_task_folder="../dataset/arc_hv_tasks/arc_extend_h/",v_task_folder="../dataset/arc_hv_tasks/arc_extend_v/")
#gen1.gen_h_to_v(h_task_folder="../dataset/arc_hv_tasks/arc_pile_h/",v_task_folder="../dataset/arc_hv_tasks/arc_pile_v/")
#gen1.gen_h_to_v(h_task_folder="../dataset/arc_hv_tasks/arc_extend_L_h/",v_task_folder="../dataset/arc_hv_tasks/arc_extend_L_v/")
#gen1.gen_h_to_v(h_task_folder="../dataset/arc_hv_tasks/arc_downone_h/",v_task_folder="../dataset/arc_hv_tasks/arc_downone_v/", degree=1)
#gen1.gen_h_to_v(h_task_folder="../dataset/arc_hv_tasks/arc_3906de3d_h/",v_task_folder="../dataset/arc_hv_tasks/arc_3906de3d_v/", degree=1)
