# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

'''
Classes:
    DatasetLoader - Generate OCW Dataset objects from a variety of sources.
'''

import ocw.data_source.local as local
import ocw.data_source.esgf as esgf
import ocw.data_source.rcmed as rcmed
import ocw.data_source.dap as dap


class DatasetLoader:
    '''Generate a list of OCW Dataset objects from a variety of sources.'''

    def __init__(self, *loader_opts):
        '''Generate a list of OCW Dataset objects from a variety of sources.

        Each keyword argument can be information for a dataset in dictionary
        form. For example:
        ``
        >>> loader_opt1 = {'loader_name': 'rcmed', 'name': 'cru',
                           'dataset_id': 10, 'parameter_id': 34}
        >>> loader_opt2 = {'path': './data/TRMM_v7_3B43_1980-2010.nc,
                           'variable': 'pcp'}
        >>> loader = DatasetLoader(loader_opt1, loader_opt2)
        ``

        Or more conveniently if the loader configuration is defined in a
        yaml file named config_file (see RCMES examples):
        ``
        >>> import yaml
        >>> config = yaml.load(open(config_file))
        >>> obs_loader_config = config['datasets']['reference']
        >>> loader = DatasetLoader(*obs_loader_config)
        ``

        As shown in the first example, the dictionary for each argument should
        contain a loader name and parameters specific to the particular loader.
        Once the configuration is entered, the datasets may be loaded using:
        ``
        >>> loader.load_datasets()
        >>> obs_datasets = loader.datasets
        ``

        Additionally, each dataset must have a ``loader_name`` keyword. This may
        be one of the following:
        * ``'local'`` - One or multiple dataset files in a local directory
        * ``'local_split'`` - A single dataset split accross multiple files in a
                              local directory
        * ``'esgf'`` - Download the dataset from the Earth System Grid
                       Federation
        * ``'rcmed'`` - Download the dataset from the Regional Climate Model
                        Evaluation System Database
        * ``'dap'`` - Download the dataset from an OPeNDAP URL

        Users who wish to download datasets from loaders not described above
        may define their own custom dataset loader function and incorporate it
        as follows:
        >>> loader.add_source_loader('my_loader_name', my_loader_func)

        :param loader_opts: Dictionaries containing the each dataset loader
                            configuration, representing the keyword arguments of
                            the loader function specified by an additional key
                            called 'loader_name'. If not specified by the user,
                            this defaults to local.
        :type loader_opts: :class:`dict`

        :raises KeyError: If an invalid argument is passed to a data source
        loader function.
        '''
        # Reference dataset config
        self.set_loader_opts(*loader_opts)

        # Default loaders
        self._source_loaders = {
            'local': local.load_multiple_files,
            'local_split': local.load_dataset_from_multiple_netcdf_files,
            'esgf': esgf.load_dataset,
            'rcmed': rcmed.parameter_dataset,
            'dap': dap.load
        }

    def add_source_loader(self, loader_name, loader_func):
        '''
        Add a custom source loader.

        :param source_name: The name of the data source.
        :type source_name: :mod:`string`

        :param loader_func: Reference to a custom defined function. This should
        return an OCW Dataset object.
        :type loader_func: :class:`callable`
        '''
        self._source_loaders[loader_name] = loader_func

    def add_loader_opts(self, *loader_opts):
        '''
        A convenient means of adding loader options for each dataset to the
        loader. If 'loader_name' is not entered as a keyword argument, then
        'local' is used by default.

        :param loader_opts: Dictionaries containing the each dataset loader
                            configuration, representing the keyword arguments of
                            the loader function specified by an additional key
                            called 'loader_name'. If not specified by the user,
                            this defaults to local.
        :type loader_opts: :mod:`dict`
        '''
        for opt in loader_opts:
            if 'loader_name' not in opt:
                opt['loader_name'] = 'local'
        self._config.extend(loader_opts)

    def set_loader_opts(self, *loader_opts):
        '''
        Reset the dataset loader config.

        :param loader_opts: Dictionaries containing the each dataset loader
                            configuration, representing the keyword arguments of
                            the loader function specified by an additional key
                            called 'loader_name'. If not specified by the user,
                            this defaults to local.
        :type loader_opts: :mod:`dict`
        '''
        self._config = []
        self.add_loader_opts(*loader_opts)

    def load_datasets(self):
        '''
        Loads the datasets from the given loader configurations.
        '''
        # Ensure output is clear if loading is performed more than once to
        # prevent duplicates.
        self.datasets = []

        # Load the target datasets
        for loader_opt in self._config:
            output = self._load(**loader_opt)

            # Need to account for the fact that some loaders return lists
            # of OCW Dataset objects instead of just one
            if isinstance(output, list):
                self.datasets.extend(output)
            else:
                self.datasets.append(output)

    def _load(self, **kwargs):
        '''
        Generic dataset loading method.
        '''
        # Extract the loader name
        loader_name = kwargs.pop('loader_name')

        # Find the correct loader function for the given data source
        loader_func = self._source_loaders[loader_name]

        # The remaining kwargs should be specific to the loader
        output = loader_func(**kwargs)

        # Preserve loader_name info for later use
        kwargs['loader_name'] = loader_name
        return output