from Generator import *

# 2D arc horizontal vs vertical generators
class Gen_arc_3906de3d(Generator):
    # Allow only <= 2 blank rows in each io pair, and must be in training set
    num_blank=0

    def __init__(self, save_folder, file_prefix, min_len=32, max_len=33):
        super().__init__(save_folder, file_prefix, min_len, max_len)
        self.gen_func = self.gen_3906de3d_pair

    def gen_3906de3d_pair(self, seq_len, io_idx, shape1_color, shape2_color):
        """
        Generate 3906de3d, horizontally and vertically, size=according to shapes
        """
        grid_len = 10

        input, output=[],[]
        self.num_blank=0
        bar_start = randrange(1,3)
        bar_end = randrange(6,10)

        for i in range(grid_len):
            input_r,output_r = [self.back_ground]*grid_len,[self.back_ground]*grid_len
            if i>=bar_start and i<=bar_end:
                bar_len = randrange(4,8)
                has_bar2 = randrange(2)
                if has_bar2==1:
                    bar2_start = randrange(1,bar_len)
                    bar2_len = bar_len-bar2_start-1

                    for j in range(bar_len):
                        if j < bar2_start:
                            input_r[j] = shape1_color
                            output_r[j] = shape1_color
                        else:
                            output_r[j] = shape2_color

                    for j in range(grid_len-bar2_len-1,grid_len):
                            input_r[j] = shape2_color
                else:
                    for j in range(bar_len):
                        input_r[j] = shape1_color
                        output_r[j] = shape1_color


            input.append(input_r)
            output.append(output_r)

        return input, output


gen1=Gen_arc_3906de3d(save_folder="../dataset/arc_hv_tasks/arc_3906de3d_h",file_prefix="arc_3906de3d_h")
for i in range(50):
    other_colors = list(range(1,gen1.max_digits+1))
    shape1_color = random.choice(other_colors)
    other_colors.remove(shape1_color)
    shape2_color = random.choice(other_colors)

    gen1.gen_json(dryrun=False, gen_type="2d", shape1_color=shape1_color,shape2_color=shape2_color)
