def full_to_interiornet():
    scene_file_path = "/nvme/datasets/diml_depth/scenes.txt"
    base_path = "/nvme/datasets/diml_depth/"
    out_path = "/nvme/datasets/diml_depth/HD7/"
    cam0_render = "/nvme/datasets/interiornet/3FO4IDEI1LAV_Bedroom/cam0.render"
    num_frames = 20
    shape = (672, 378)
    np.random.seed(123)

    with open(scene_file_path, 'r') as f:
        scene_lines = f.readlines()

    scene_lines = [sn.split('\n')[0] for sn in scene_lines]
    scene_paths = [os.path.join(base_path, sn.split('-')[0]) for sn in scene_lines]
    scene_ranges = [sn.split('-')[1] for sn in scene_lines]
    scene_ranges = [(int(rn[1:-1].split(':')[0]), int(rn[1:-1].split(':')[1])) for rn in scene_ranges]

    for i, scene_path in enumerate(scene_paths):
        file_list = []
        for j in range(scene_ranges[i][0], scene_ranges[i][1]+1):
            scene_path_col = os.path.join(scene_path, "{}/col".format(j))
            if os.path.exists(scene_path_col):
                file_list += [os.path.join(scene_path_col, dn) for dn in os.listdir(scene_path_col)]

        scene_count = len(os.listdir(out_path))
        scene_out_path = "{:02d}DIML_{}".format(scene_count + 1, scene_path.split('/')[-2].split(' ')[1])
        scene_out_path = os.path.join(out_path, scene_out_path)

        if os.path.exists(scene_out_path):
            shutil.rmtree(scene_out_path)
        os.mkdir(scene_out_path)
        os.mkdir(os.path.join(scene_out_path, "cam0"))
        os.mkdir(os.path.join(scene_out_path, "depth0"))
        os.mkdir(os.path.join(scene_out_path, "label0"))
        os.mkdir(os.path.join(scene_out_path, "cam0", "data"))
        os.mkdir(os.path.join(scene_out_path, "depth0", "data"))
        os.mkdir(os.path.join(scene_out_path, "label0", "data"))
        shutil.copyfile(cam0_render, os.path.join(scene_out_path, "cam0.render"))
        print(scene_out_path)

        frame_paths = np.random.choice(file_list, num_frames, False)
        for j, frame_path in enumerate(frame_paths):
            img = cv2.imread(frame_path, cv2.IMREAD_UNCHANGED)
            depth_path = frame_path.replace('/col/', '/up_png/')
            depth_path = depth_path.replace('_c.png', '_ud.png')
            depth_img = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)

            if depth_img is None:
                print(depth_path)
                exit()
            if img is None:
                print(frame_path)
                exit()

            img = cv2.resize(img, dsize=shape, interpolation=cv2.INTER_LINEAR)
            depth_img = cv2.resize(depth_img, dsize=shape,
                                   interpolation=cv2.INTER_LINEAR)
            label_img = depth_img.copy()
            label_img[:, :] = 3

            cv2.imwrite(os.path.join(scene_out_path, "cam0", "data", "{}.png".format(j)), img)
            cv2.imwrite(os.path.join(scene_out_path, "depth0", "data", "{}.png".format(j)), depth_img)
            cv2.imwrite(os.path.join(scene_out_path, "label0", "data", "{}_instance.png".format(j)), label_img)
            cv2.imwrite(os.path.join(scene_out_path, "label0", "data", "{}_nyu.png".format(j)), label_img)


if __name__ == '__main__':
    full_to_interiornet()