# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
# Copyright 2023 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
load
"""
from typing import Union, Optional, Dict, Sequence, Mapping
from datasets import load_dataset as hf_load
from datasets import Dataset, IterableDataset, Split, Features, \
    DownloadConfig, DownloadMode, VerificationMode, Version
from mindspore.dataset import GeneratorDataset


class TransferIterableDataset():
    """TransferDataset for Huggingface Dataset."""
    def __init__(self, arrow_ds, column_names):
        r"""
        Initializes a new instance of TransferIterableDataset.
        
        Args:
            self (TransferIterableDataset): The instance of the TransferIterableDataset class.
            arrow_ds (ArrowDataset): The Arrow dataset to be used for transferring data.
            column_names (list): A list of column names to be used for data transfer.
        
        Returns:
            None. The method initializes the TransferIterableDataset instance with the provided parameters.
        
        Raises:
            No specific exceptions are raised within this method.
        """
        self.ds = arrow_ds
        self.column_names = column_names

    def __iter__(self):
        r"""
        Method '__iter__' in the class 'TransferIterableDataset'.
        
        Args:
            self: TransferIterableDataset
                The instance of TransferIterableDataset.
                This parameter is required to iterate over the dataset.
                
        Returns:
            None
                This method does not return any explicit value.
        
        Raises:
            None
        """
        for data in self.ds:
            yield tuple(data[name] for name in self.column_names)

class TransferDataset():
    """TransferDataset for Huggingface Dataset."""
    def __init__(self, arrow_ds, column_names):
        r"""
        __init__ method in the TransferDataset class.
        
        Args:
            self (object): The instance of the TransferDataset class.
            arrow_ds (object): An Arrow Dataset representing the data to be transferred.
            column_names (list): A list of column names to be included in the transfer.
        
        Returns:
            None. This method does not return any value.
        
        Raises:
            No specific exceptions are raised by this method.
        """
        self.ds = arrow_ds
        self.column_names = column_names

    def __getitem__(self, index):
        r"""
        Retrieve the item at the specified index.
        
        Args:
            self (TransferDataset): The TransferDataset instance from which to retrieve the item.
            index (int): The index of the item to retrieve. It should be a non-negative integer within the valid range of indices for the dataset.
        
        Returns:
            tuple: A tuple containing the values of the item at the specified index for each column in the dataset.
        
        Raises:
            IndexError: If the index is out of range or not a valid integer.
            KeyError: If the column name is not found in the dataset.
        """
        return tuple(self.ds[int(index)][name] for name in self.column_names)

    def __len__(self):
        r"""
        This method '__len__' is implemented in the class 'TransferDataset' to return the number of rows in the dataset.
        
        Args:
            self: Represents the instance of the class. It is used to access the dataset stored in the class. No restrictions.
        
        Returns:
            None. This method does not return any value directly to the caller but rather accesses the 'num_rows' attribute of the dataset stored in the class.
        
        Raises:
            No exceptions are explicitly raised by this method. However, if the 'num_rows' attribute is not properly initialized or if there are any issues with accessing the dataset, unexpected errors might
occur.
        """
        return self.ds.num_rows


def load_dataset(
    path: str,
    name: Optional[str] = None,
    data_dir: Optional[str] = None,
    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
    split: Optional[Union[str, Split]] = None,
    cache_dir: Optional[str] = None,
    features: Optional[Features] = None,
    download_config: Optional[DownloadConfig] = None,
    download_mode: Optional[Union[DownloadMode, str]] = None,
    verification_mode: Optional[Union[VerificationMode, str]] = None,
    keep_in_memory: Optional[bool] = None,
    save_infos: bool = False,
    revision: Optional[Union[str, Version]] = None,
    streaming: bool = False,
    num_proc: Optional[int] = None,
    storage_options: Optional[Dict] = None,
    **config_kwargs,
):
    """Load a dataset from the Hugging Face Hub, or a local dataset.

    You can find the list of datasets on the [Hub](https://hf-mirror.com/datasets) or with [`huggingface_hub.list_datasets`].

    A dataset is a directory that contains:

    - some data files in generic formats (JSON, CSV, Parquet, text, etc.).
    - and optionally a dataset script, if it requires some code to read the data files. This is used to load any kind of formats or structures.

    Note that dataset scripts can also download and read data files from anywhere - in case your data files already exist online.

    This function does the following under the hood:

        1. Download and import in the library the dataset script from `path` if it's not already cached inside the library.

            If the dataset has no dataset script, then a generic dataset script is imported instead (JSON, CSV, Parquet, text, etc.)

            Dataset scripts are small python scripts that define dataset builders. They define the citation, info and format of the dataset,
            contain the path or URL to the original data files and the code to load examples from the original data files.

            You can find the complete list of datasets in the Datasets [Hub](https://hf-mirror.com/datasets).

        2. Run the dataset script which will:

            * Download the dataset file from the original URL (see the script) if it's not already available locally or cached.
            * Process and cache the dataset in typed Arrow tables for caching.

                Arrow table are arbitrarily long, typed tables which can store nested objects and be mapped to numpy/pandas/python generic types.
                They can be directly accessed from disk, loaded in RAM or even streamed over the web.

        3. Return a dataset built from the requested splits in `split` (default: all).

    It also allows to load a dataset from a local directory or a dataset repository on the Hugging Face Hub without dataset script.
    In this case, it automatically loads all the data files from the directory or the dataset repository.

    Args:

        path (`str`):
            Path or name of the dataset.
            Depending on `path`, the dataset builder that is used comes from a generic dataset script (JSON, CSV, Parquet, text etc.) or
            from the dataset script (a python file) inside the dataset directory.

            For local datasets:

            - if `path` is a local directory (containing data files only)
              -> load a generic dataset builder (csv, json, text etc.) based on the content of the directory
              e.g. `'./path/to/directory/with/my/csv/data'`.
            - if `path` is a local dataset script or a directory containing a local dataset script (if the script has the same name as the directory)
              -> load the dataset builder from the dataset script
              e.g. `'./dataset/squad'` or `'./dataset/squad/squad.py'`.

            For datasets on the Hugging Face Hub (list all available datasets with [`huggingface_hub.list_datasets`])

            - if `path` is a dataset repository on the HF hub (containing data files only)
              -> load a generic dataset builder (csv, text etc.) based on the content of the repository
              e.g. `'username/dataset_name'`, a dataset repository on the HF hub containing your data files.
            - if `path` is a dataset repository on the HF hub with a dataset script (if the script has the same name as the directory)
              -> load the dataset builder from the dataset script in the dataset repository
              e.g. `glue`, `squad`, `'username/dataset_name'`, a dataset repository on the HF hub containing a dataset script `'dataset_name.py'`.

        name (`str`, *optional*):
            Defining the name of the dataset configuration.
        data_dir (`str`, *optional*):
            Defining the `data_dir` of the dataset configuration. If specified for the generic builders (csv, text etc.)
            or the Hub datasets and `data_files` is `None`,
            the behavior is equal to passing `os.path.join(data_dir, **)` as `data_files` to reference all the files in a directory.
        data_files (`str` or `Sequence` or `Mapping`, *optional*):
            Path(s) to source data file(s).
        split (`Split` or `str`):
            Which split of the data to load.
            If `None`, will return a `dict` with all splits (typically `datasets.Split.TRAIN` and `datasets.Split.TEST`).
            If given, will return a single Dataset.
            Splits can be combined and specified like in tensorflow-datasets.
        cache_dir (`str`, *optional*):
            Directory to read/write data. Defaults to `"~/.cache/huggingface/datasets"`.
        features (`Features`, *optional*):
            Set the features type to use for this dataset.
        download_config ([`DownloadConfig`], *optional*):
            Specific download configuration parameters.
        download_mode ([`DownloadMode`] or `str`, defaults to `REUSE_DATASET_IF_EXISTS`):
            Download/generate mode.
        verification_mode ([`VerificationMode`] or `str`, defaults to `BASIC_CHECKS`):
            Verification mode determining the checks to run on the downloaded/processed dataset information (checksums/size/splits/...).
        keep_in_memory (`bool`, defaults to `None`):
            Whether to copy the dataset in-memory. If `None`, the dataset
            will not be copied in-memory unless explicitly enabled by setting `datasets.config.IN_MEMORY_MAX_SIZE` to
            nonzero. See more details in the [improve performance](../cache#improve-performance) section.
        save_infos (`bool`, defaults to `False`):
            Save the dataset information (checksums/size/splits/...).
        revision ([`Version`] or `str`, *optional*):
            Version of the dataset script to load.
            As datasets have their own git repository on the Datasets Hub, the default version "main" corresponds to their "main" branch.
            You can specify a different version than the default "main" by using a commit SHA or a git tag of the dataset repository.
        token (`str` or `bool`, *optional*):
            Optional string or boolean to use as Bearer token for remote files on the Datasets Hub.
            If `True`, or not specified, will get token from `"~/.huggingface"`.
        streaming (`bool`, defaults to `False`):
            If set to `True`, don't download the data files. Instead, it streams the data progressively while
            iterating on the dataset. An [`IterableDataset`] or [`IterableDatasetDict`] is returned instead in this case.

            Note that streaming works for datasets that use data formats that support being iterated over like txt, csv, jsonl for example.
            Json files may be downloaded completely. Also streaming from remote zip or gzip files is supported but other compressed formats
            like rar and xz are not yet supported. The tgz format doesn't allow streaming.
        num_proc (`int`, *optional*, defaults to `None`):
            Number of processes when downloading and generating the dataset locally.
            Multiprocessing is disabled by default.

        storage_options (`dict`, *optional*, defaults to `None`):
            **Experimental**. Key/value pairs to be passed on to the dataset file-system backend, if any.

        **config_kwargs (additional keyword arguments):
            Keyword arguments to be passed to the `BuilderConfig`
            and used in the [`DatasetBuilder`].

    Returns:
        [`GeneratorDataset`] or [`Dict`]:
        - if `split` is not `None`: the dataset requested,
        - if `split` is `None`, a [`Dict`] with each split.

    Example:

    Load a dataset from the Hugging Face Hub:

    ```py
    >>> from datasets import load_dataset
    >>> ds = load_dataset('rotten_tomatoes', split='train')

    # Map data files to splits
    >>> data_files = {'train': 'train.csv', 'test': 'test.csv'}
    >>> ds = load_dataset('namespace/your_dataset_name', data_files=data_files)
    ```

    Load a local dataset:

    ```py
    # Load a CSV file
    >>> from datasets import load_dataset
    >>> ds = load_dataset('csv', data_files='path/to/local/my_dataset.csv')

    # Load a JSON file
    >>> from datasets import load_dataset
    >>> ds = load_dataset('json', data_files='path/to/local/my_dataset.json')

    # Load from a local loading script
    >>> from datasets import load_dataset
    >>> ds = load_dataset('path/to/local/loading_script/loading_script.py', split='train')
    ```

    ```py
    >>> from datasets import load_dataset
    >>> ds = load_dataset('rotten_tomatoes', split='train', streaming=True)
    ```

    Load an image dataset with the `ImageFolder` dataset builder:

    ```py
    >>> from datasets import load_dataset
    >>> ds = load_dataset('imagefolder', data_dir='/path/to/images', split='train')
    ```
    """
    shuffle = config_kwargs.get('shuffle', False)

    ds_ret = hf_load(path,
                     name=name,
                     data_dir=data_dir,
                     data_files=data_files,
                     split=split,
                     cache_dir=cache_dir,
                     features=features,
                     download_config=download_config,
                     download_mode=download_mode,
                     verification_mode=verification_mode,
                     keep_in_memory=keep_in_memory,
                     save_infos=save_infos,
                     revision=revision,
                     streaming=streaming,
                     num_proc=num_proc,
                     storage_options=storage_options,
                     )
    if isinstance(ds_ret, (list, tuple)):
        ds_dict = dict(zip(split, ds_ret))
    elif isinstance(ds_ret, (Dataset, IterableDataset)):
        assert isinstance(split, str)
        ds_dict = {split: ds_ret}
    else:
        ds_dict = ds_ret

    datasets_dict = {}

    for key, raw_ds in ds_dict.items():
        column_names = list(raw_ds.features.keys())
        source = TransferDataset(raw_ds, column_names) if isinstance(raw_ds, Dataset) \
            else TransferIterableDataset(raw_ds, column_names)

        ms_ds = GeneratorDataset(source=source,
        column_names=column_names,
        shuffle=shuffle,
        num_parallel_workers=num_proc if num_proc else 1)
        datasets_dict[key] = ms_ds

    if len(datasets_dict) == 1:
        return datasets_dict.popitem()[1]
    return datasets_dict
