from cnvrg.modules.cnvrg_files import CnvrgFiles
import cnvrg.helpers.param_build_helper as param_build_helper
import cnvrg.helpers.config_helper as config_helper
from cnvrg.modules.errors import CnvrgError
from cnvrg.modules.storage import storage_factory
import os.path

class Dataset(CnvrgFiles):
    def __init__(self, dataset=None, dataset_url=None, working_dir=None):
        if dataset_url:
            owner_slug, project_slug = Dataset.get_owner_and_project_from_url(dataset_url)
        else:
            owner_slug, project_slug = param_build_helper.parse_params(dataset, param_build_helper.DATASET, working_dir=working_dir)
        self.__owner = owner_slug
        self.__dataset = project_slug

        in_dir = config_helper.is_in_dir(config_helper.CONFIG_TYPE_DATASET, project_slug, working_dir)
        super(Dataset, self).__init__()
        self._set_working_dir(working_dir)
        if in_dir:
            self._set_working_dir(config_helper.find_config_dir(path=working_dir))
            self.in_dir = True
        elif not self.__dataset or not self.__owner:
            raise CnvrgError("Cant init project without params and outside project directory")

    def get_base_url(self):
        return "users/{owner}/datasets/{dataset}".format(owner=self.__owner, dataset=self.__dataset)

    def fetch_all_files(self):
        return self.get_commit_files()

    def pytorch_dataset(self, loader=None):
        from torch.utils.data.dataset import Dataset as TorchDataset
        class CnvrPytorchDataset(TorchDataset):
            'Characterizes a dataset for PyTorch'
            def __init__(self, dataset: CnvrgFiles=None):
                'Initialization'
                self.dataset = dataset
                self.files = dict(map(lambda x: [x.get("fullpath"), x], dataset.get_commit_files()))
                self.keys = list(self.files.keys())
                self.storage = storage_factory(dataset)

            def __len__(self):
                'Denotes the total number of samples'
                return len(self.files.keys())

            def __getitem__(self, index):
                'Generates one sample of data'
                # Select sample
                file_def = self.files[self.keys[index]]
                file_path = os.path.join(self.dataset.get_working_dir(), file_def.get("fullpath"))
                if not os.path.exists(file_path):
                    self.storage.download_single_file(file_def)
                return file_path if not loader else loader(file_path)


        return CnvrPytorchDataset(dataset=self)
