import torch.utils.data as data
import numpy as np
from imageio import imread
from path import Path
import pdb

def crawl_folders(folders_list):
        imgs = []
        depth = []
        for folder in folders_list:
            current_imgs = sorted(folder.files('*.jpg'))
            current_depth = []
            for img in current_imgs:
                d = img.dirname()/(img.name[:-4] + '.npy')
                assert(d.isfile()), "depth file {} not found".format(str(d))
                depth.append(d)
            imgs.extend(current_imgs)
            depth.extend(current_depth)
        return imgs, depth


def load_as_float(path):
    return imread(path).astype(np.float32)


class ValidationSet(data.Dataset):
    """A sequence data loader where the files are arranged in this way:
        root/scene_1/0000000.jpg
        root/scene_1/0000000.npy
        root/scene_1/0000001.jpg
        root/scene_1/0000001.npy
        ..
        root/scene_2/0000000.jpg
        root/scene_2/0000000.npy
        .

        transform functions must take in a list a images and a numpy array which can be None
    """

    def __init__(self, root, transform=None):
        self.root = Path(root)
        scene_list_path = self.root/'val.txt'
        self.scenes = [self.root/folder[:-1] for folder in open(scene_list_path)]
        self.imgs, self.depth = crawl_folders(self.scenes)
        self.transform = transform

    def __getitem__(self, index):
        img = load_as_float(self.imgs[index])
        depth = np.load(self.depth[index]).astype(np.float32) #;pdb.set_trace()
        if self.transform is not None: 
            img, _, _ = self.transform([img], depth, None); #this depth is just used to fill the compose transform that is shared(no need for the result)
            img = img[0]
        return img, depth

    def __len__(self):
        return len(self.imgs)

[((41, 20, 41, 30), 'path.Path', 'Path', ({(41, 25, 41, 29): 'root'}, {}), '(root)', False, 'from path import Path\n'), ((23, 11, 23, 23), 'imageio.imread', 'imread', ({(23, 18, 23, 22): 'path'}, {}), '(path)', False, 'from imageio import imread\n'), ((49, 16, 49, 42), 'numpy.load', 'np.load', ({(49, 24, 49, 41): 'self.depth[index]'}, {}), '(self.depth[index])', True, 'import numpy as np\n')]