### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. 
### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
import os.path
from data.base_dataset import BaseDataset, get_params, get_transform, normalize
from data.image_folder import make_dataset
from PIL import Image
import time as time

class RecursiveDataset(BaseDataset):
    def initialize(self, opt):
        self.opt = opt
        self.root = opt.dataroot    

        ### input A (label maps)
        self.dataset_size = 1000000
    
    def __getitem__(self, index):
        ### input A (label maps)
        A_path = os.path.join(self.opt.render_dir, "frame_{:05d}.png".format(index))
        if not os.path.exists(A_path):
            # print("{} doesn't exist, waiting for it".format(A_path))
            while not os.path.exists(A_path):
                # print('sleeping for {}'.format(self.opt.poll_delay))
                time.sleep(self.opt.poll_delay)
        # print("got {}".format(A_path))
        A = Image.open(A_path)
        params = get_params(self.opt, A.size)
        transform_A = get_transform(self.opt, params)
        A_tensor = transform_A(A.convert('RGB'))

        B_tensor = inst_tensor = feat_tensor = 0

        input_dict = {'label': A_tensor, 'inst': inst_tensor, 'image': B_tensor, 
                      'feat': feat_tensor, 'path': A_path}

        return input_dict

    def __len__(self):
        return self.dataset_size

    def name(self):
        return 'RecursiveDataset'