class ExampleDataset:
    def __init__(self, root, policy_container=None, is_train=True):
        self.root = root
        self.data_type = "train" if is_train else "val"
        self.policy_container = policy_container
        # main
        # self.imgs = list(sorted(glob.glob(f'{root}/images/{self.data_type}/*.jpg')))
        # self.boxes = list(sorted(glob.glob(f'{root}/labels/{self.data_type}/*.txt')))
        # test
        self.imgs = list(sorted(glob.glob(f'{root}/test_data/*.jpg')))
        self.boxes = list(sorted(glob.glob(f'{root}/test_data/*txt')))
        self.out_dir = f'{root}/test_data_out/'
        if not os.path.exists(self.out_dir):
            os.mkdir(self.out_dir)

    def __len__(self):
        return len(self.imgs)
        
    # def __getitem__(self, idx):
    #     img = np.array(Image.open(self.imgs[idx]))
    #     boxes_path = self.boxes[idx]
    #     height, width, _ = img.shape
    #     # For convenience I’ve hard coded the label and co-ordinates as label, x_min, y_min, x_max, y_max
    #     # for each bounding box in the image. For your own model you will need to load
    #     # in the coordinates and do the appropriate transformations.
    #     boxes = []
    #     labels = []
    #     with open(boxes_path, 'r') as in_box:
    #         for line in in_box:
    #             if line:
    #                 line = line.split()
    #                 xywh = list(map(int, map(float, line[1:])))
    #                 xyxy = self.convert_xyxy(xywh, width=width, height=height)
    #                 boxes.append(xyxy)
    #                 labels.append(int(line[0]))
        
    #     if self.policy_container:

    #         # Select a random sub-policy from the policy list
    #         random_policy = self.policy_container.select_random_policy()
    #         print(random_policy)

    #         # Apply this augmentation to the image, returns the augmented image and bounding boxes
    #         # The boxes must be at a pixel level. e.g. x_min, y_min, x_max, y_max with pixel values
    #         img_aug, bbs_aug = self.policy_container.apply_augmentation(
    #             random_policy,
    #             img,
    #             boxes,
    #             labels,
    #         )
    #         labels = np.array(labels)
    #         boxes = np.hstack((np.vstack(labels), np.array(boxes))) # Add the labels to the boxes
    #         bbs_aug= np.array(bbs_aug)
            
    #         # Only return the augmented image and bounded boxes if there are
    #         # boxes present after the image augmentation
    #         if bbs_aug.size > 0:
    #             return img, boxes, img_aug, bbs_aug
    #         else:
    #             return img, boxes, [], np.array([])
    #     return img, boxes

    def run(self, num_random=10):
        for idx in range(len(self.imgs)):
            img = np.array(Image.open(self.imgs[idx]))
            boxes_path = self.boxes[idx]
            height, width, _ = img.shape
            # For convenience I’ve hard coded the label and co-ordinates as label, x_min, y_min, x_max, y_max
            # for each bounding box in the image. For your own model you will need to load
            # in the coordinates and do the appropriate transformations.
            boxes = []
            labels = []
            with open(boxes_path, 'r') as in_box:
                for line in in_box:
                    if line:
                        line = line.split()
                        xywh = list(map(int, map(float, line[1:])))
                        xyxy = self.convert_xyxy(xywh, width=width, height=height)
                        boxes.append(xyxy)
                        labels.append(int(line[0]))
            
            if self.policy_container:
                # run $num_random times
                print("Processing: " + self.imgs[idx])
                i = 0
                for i in range(num_random):

                    # Select a random sub-policy from the policy list
                    random_policy = self.policy_container.select_random_policy()

                    # Apply this augmentation to the image, returns the augmented image and bounding boxes
                    # The boxes must be at a pixel level. e.g. x_min, y_min, x_max, y_max with pixel values
                    img_aug, bbs_aug = self.policy_container.apply_augmentation(
                        random_policy,
                        img,
                        boxes,
                        labels
                    )
                    labels = np.array(labels)
                    boxes = np.hstack((np.vstack(labels), np.array(boxes))) # Add the labels to the boxes
                    bbs_aug= np.array(bbs_aug)
                    
                    # Only return the augmented image and bounded boxes if there are
                    # boxes present after the image augmentation
                    if bbs_aug.size > 0:
                        print("Step: " + str(i))
                        print(random_policy)
                        # img, boxes, img_aug, bbs_aug
                        # to write
                        cv2.imwrite(str(i)+"_bbaug_"+self.imgs[idx], img_aug)
                        with open(str(i)+"_bbaug_"+self.boxes[idx], "w") as fw:
                            fw.writelines(bbs_aug)
                        i += 1
    
    def convert_xyxy(self, xywh, width, height):
        x, w = xywh[0] * width, xywh[2] * width
        y, h = xywh[1] * height, xywh[3] * height
        x1 = x - w / 2
        x2 = x + w / 2
        y1 = y - h / 2
        y2 = y + h / 2

        return list(map(int, [x1, x2, y1, y2]))