"""Implements loading of Hdf5 files into the dantro data tree"""

import logging

import numpy as np
import h5py as h5

from ..base import BaseDataGroup, BaseDataContainer
from ..containers import NumpyDataContainer
from ..groups import OrderedDataGroup
from ..proxy import Hdf5DataProxy
from ..tools import fill_line, decode_bytestrings
from ._tools import add_loader

# Local constants
log = logging.getLogger(__name__)


# -----------------------------------------------------------------------------

class Hdf5LoaderMixin:
    """Supplies functionality to load hdf5 files into the data manager.
    
    It resolves the hdf5 groups into corresponding data groups and the datasets
    into NumpyDataContainers.

    If ``enable_mapping`` is set, the class variables ``_HDF5_DSET_MAP`` and
    ``_HDF5_GROUP_MAP`` are used to map from a string to a container type. The
    class variable ``_HDF5_MAP_FROM_ATTR`` determines the default value of the
    attribute to read and use as input string for the mapping. 
    """

    # Define the container classes to use when loading data with this mixin
    # The default class to use for containers
    _HDF5_DSET_DEFAULT_CLS = NumpyDataContainer

    # The mapping of types to data groups
    _HDF5_GROUP_MAP = None

    # The mapping of types to data containers
    _HDF5_DSET_MAP = None

    # The name of the attribute to read for mapping
    _HDF5_MAP_FROM_ATTR = None

    # Whether to decode strings stored as byte arrays to regular python strings
    _HDF5_DECODE_ATTR_BYTESTRINGS = True


    @add_loader(TargetCls=OrderedDataGroup, omit_self=False)
    def _load_hdf5(self, filepath: str, *,
                   TargetCls: type,
                   load_as_proxy: bool=False,
                   proxy_kwargs: dict=None,
                   lower_case_keys: bool=False,
                   enable_mapping: bool=False, 
                   map_from_attr: str=None,
                   print_params: dict=None
                   ) -> OrderedDataGroup:
        """Loads the specified hdf5 file into DataGroup- and DataContainer-like
        objects; this completely recreates the hierarchic structure of the hdf5
        file. The data can be loaded into memory completely, or be loaded as
        a proxy object.
        
        The h5py File and Group objects will be converted to the specified
        DataGroup-derived objects; the Dataset objects to the specified
        DataContainer-derived object.
        
        All HDF5 group or dataset attributes are carried over and are
        accessible under the ``attrs`` attribute of the respective dantro
        objects in the tree.
        
        Args:
            filepath (str): The path to the HDF5 file that is to be loaded
            TargetCls (type): The group type this is loaded into
            load_as_proxy (bool, optional): if True, the leaf datasets are
                loaded as :py:class:`dantro.proxy.hdf5.Hdf5DataProxy`
                objects. That way, the data is only loaded into memory when
                their ``.data`` property is accessed the first time, either
                directly or indirectly.
            proxy_kwargs (dict, optional): When loading as proxy, these
                parameters are unpacked in the ``__init__`` call. For available
                argument see :py:class:`dantro.proxy.hdf5.Hdf5DataProxy`.
            lower_case_keys (bool, optional): whether to use only lower-case
                versions of the paths encountered in the HDF5 file.
            enable_mapping (bool, optional): If true, will use the class
                variables _HDF5_GROUP_MAP and _HDF5_DSET_MAP to map groups or
                datasets to a custom container class during loading. Which
                attribute to read is determined by the `map_from_attr` argument
            map_from_attr (str, optional): From which attribute to read the
                key that is used in the mapping. If nothing is given, the
                class variable _HDF5_MAP_FROM_ATTR is used.
            print_params (dict, optional): parameters for the status report.
                Available keys:

                level (int):
                    how verbose to print loading info; possible values are:
                    ``0``: None, ``1``: on file level, ``2``: on dataset level
                fstr1:
                    format string level 1, receives keys ``name`` and ``file``,
                    which is the file path.
                fstr2:
                    format string level 2, receives keys ``name``, ``file`` and
                    ``obj``, which is an ``h5py.Dataset``.
        
        Returns:
            OrderedDataGroup: The populated root-level group, corresponding to
                the base group of the file
        
        Raises:
            ValueError: If `enable_mapping`, but no map attribute can be
                determined from the given argument or the class variable
                _HDF5_MAP_FROM_ATTR
        """
        
        def recursively_load_hdf5(src: h5.Group, target: BaseDataGroup, *,
                                  load_as_proxy: bool,
                                  proxy_kwargs: dict, 
                                  lower_case_keys: bool,
                                  DsetCls: BaseDataContainer,
                                  enable_mapping: bool, GroupMap: dict,
                                  DsetMap: dict,map_attr: str):
            """Recursively loads the data from the source hdf5 file into the 
            target DataGroup object.
            If given, each group or dataset is checked whether an attribute
            `container_type` or `dset_type` exists, which is then used to
            apply a mapping from that attribute to a certain type of DataGroup
            or DataContainer, respectively.
            
            Args:
                src (h5.Group): The source group to iterate over
                target (BaseDataGroup): The target group where the content from
                    the source group is loaded into
                load_as_proxy (bool): Whether to load as Hdf5DataProxy
                proxy_kwargs (dict): Unpacked in Hdf5DataProxy.__init__
                lower_case_keys (bool): Whether to make keys lower-case
                DsetCls (BaseDataContainer): The type that is used to create
                    the dataset-equivalents in ``target``
                enable_mapping (bool): Whether type mapping should be used
                GroupMap (dict): Map of names to BaseDataGroup-derived types
                DsetMap (dict): Map of names to BaseDataContainer-derived types
                map_attr (str): The HDF5 attribute to inspect in order to
                    determine the name of the mapping
            
            Raises:
                NotImplementedError: When encountering objects other than
                    groups or datasets in the HDF5 file
            """

            def get_map_attr_val(attrs) -> str:
                """Make sure the map attribute isn't a 1-sized array!"""
                attr_val = attrs[map_attr] # map_attr and check in outer scope

                if isinstance(attr_val, np.ndarray):
                    # Need be single item and already decoded
                    attr_val = attr_val.item()

                return attr_val

            def decode_attr_val(attr_val) -> str:
                """Wrapper around decode_bytestrings"""
                # If feature not activated, return without doing anything
                if not self._HDF5_DECODE_ATTR_BYTESTRINGS:
                    return attr_val

                return decode_bytestrings(attr_val)

            # Go through the elements of the source object
            for key, obj in src.items():
                if lower_case_keys and isinstance(key, str):
                    key = key.lower()

                if isinstance(obj, h5.Group):
                    # Need to continue recursion
                    # Extract attributes manually
                    attrs = {k: decode_attr_val(v)
                             for k, v in obj.attrs.items()}

                    # Determine the class to use for this group
                    if enable_mapping and GroupMap and attrs.get(map_attr):
                        # Try to resolve the mapping
                        try:
                            _GroupCls = GroupMap[get_map_attr_val(attrs)]

                        except KeyError:
                            # Fall back to default
                            log.warning("Could not find a mapping from map "
                                        "attribute %s='%s' (originally %s) "
                                        "to a DataGroup class. Available "
                                        "keys: %s. Falling back to default "
                                        "class ...",
                                        map_attr, get_map_attr_val(attrs),
                                        attrs[map_attr],
                                        ", ".join([k
                                                   for k in GroupMap.keys()]))
                            _GroupCls = None
                    
                    else:
                        # Use the default of the target group
                        _GroupCls = None
                        # None value will lead to new_group determining the
                        # type of this group. Unlike for datasets, this is
                        # always possible, because the type of the target
                        # group is used as fallback type.

                    # Create and add the group, passing the attributes
                    target.new_group(path=key, Cls=_GroupCls,
                                     attrs=attrs)

                    # Continue recursion
                    recursively_load_hdf5(obj, target[key],
                                          load_as_proxy=load_as_proxy,
                                          proxy_kwargs=proxy_kwargs,
                                          lower_case_keys=lower_case_keys,
                                          DsetCls=DsetCls,
                                          enable_mapping=enable_mapping,
                                          GroupMap=GroupMap, DsetMap=DsetMap,
                                          map_attr=map_attr)

                elif isinstance(obj, h5.Dataset):
                    # Reached a leaf -> Import the data and attributes into a
                    # BaseDataContainer-derived object. This assumes that the
                    # given DsetCls supports np.ndarray-like data

                    # Progress information on loading this dataset
                    if plvl >= 2:
                        line = fstr2.format(name=target.name, key=key, obj=obj)
                        print(fill_line(line), end="\r")

                    if load_as_proxy:
                        # Instantiate a proxy object
                        data = Hdf5DataProxy(obj,
                                             **(proxy_kwargs if proxy_kwargs
                                                else {}))
                    else:
                        # Import the data completely
                        data = np.array(obj)

                    # Extract attributes manually
                    attrs = {k: decode_attr_val(v)
                             for k, v in obj.attrs.items()}

                    # Determine the class to use for this dataset
                    if enable_mapping and DsetMap and attrs.get(map_attr):
                        # Try to resolve the mapping
                        try:
                            _DsetCls = DsetMap[get_map_attr_val(attrs)]

                        except KeyError:
                            # Fall back to default
                            log.warning("Could not find a mapping from map "
                                        "attribute %s='%s' (originally %s) to "
                                        "a DataContainer class. Available "
                                        "keys: %s. "
                                        "Falling back to default class %s...",
                                        map_attr, get_map_attr_val(attrs),
                                        attrs[map_attr],
                                        ", ".join([k for k in DsetMap.keys()]),
                                        DsetCls.__name__)
                            _DsetCls = DsetCls
                    
                    else:
                        # If the target group supplies a default container
                        # type, use that one; otherwise fall back to default.
                        if target._NEW_CONTAINER_CLS is not None:
                            _DsetCls = None
                            # Leads to new_container taking care of the type

                        else:
                            _DsetCls = DsetCls

                    # Now create and add the dataset, passing data and attrs
                    target.new_container(path=key, Cls=_DsetCls,
                                         data=data, attrs=attrs)

                else:
                    raise NotImplementedError("Object {} is neither a dataset "
                                              "nor a group, but of type {}. "
                                              "Cannot load this!"
                                              "".format(key, type(obj)))

        # Prepare print format strings
        print_params = print_params if print_params else {}
        plvl = print_params.get('level', 0)
        fstr1 = print_params.get('fstr1', "  Loading {name:} ... ")
        fstr2 = print_params.get('fstr2', "  Loading {name:} - {key:} ...")

        # Initialize the root group
        log.debug("Loading hdf5 file %s into %s ...",
                  filepath, TargetCls.__name__)
        root = TargetCls()

        # Get the classes to use for groups and/or containers
        DsetCls = self._HDF5_DSET_DEFAULT_CLS

        # Determine from which attribute to read the mapping
        if not map_from_attr:
            # No custom value was given; use the class variable, if available
            if self._HDF5_MAP_FROM_ATTR:
                map_from_attr = self._HDF5_MAP_FROM_ATTR

            elif enable_mapping:
                # Mapping was enabled but it is unclear from which attribute
                # the map should be read. Need to raise an exception
                raise ValueError("Could not determine from which attribute "
                                 "to read the mapping. Either set the loader "
                                 "argument `map_from_attr`, the class "
                                 "variable _HDF5_MAP_FROM_ATTR, or disable "
                                 "mapping altogether via the `enable_mapping` "
                                 "argument.")

        # Now recursively go through the hdf5 file and add them to the roo
        with h5.File(filepath, 'r') as h5file:
            if plvl >= 1:
                # Print information on the level of this file
                line = fstr1.format(name=root.name, file=filepath)
                print(fill_line(line), end="\r")

            # Load the file level attributes, manually re-creating the dict
            root.attrs = {k:v for k, v in h5file.attrs.items()}

            # Now recursively load the data into the root group
            recursively_load_hdf5(h5file, root,
                                  load_as_proxy=load_as_proxy,
                                  proxy_kwargs=proxy_kwargs,
                                  lower_case_keys=lower_case_keys,
                                  DsetCls=DsetCls,
                                  enable_mapping=enable_mapping,
                                  GroupMap=self._HDF5_GROUP_MAP,
                                  DsetMap=self._HDF5_DSET_MAP,
                                  map_attr=map_from_attr)

        return root

    @add_loader(TargetCls=OrderedDataGroup, omit_self=False)
    def _load_hdf5_proxy(self, *args, **kwargs) -> OrderedDataGroup:
        """This is a shorthand for
        :py:meth:`~dantro.data_loaders.load_hdf5.Hdf5LoaderMixin._load_hdf5`
        with the ``load_as_proxy`` flag set.
        """
        return self._load_hdf5(*args, load_as_proxy=True, **kwargs)

    @add_loader(TargetCls=OrderedDataGroup, omit_self=False)
    def _load_hdf5_as_dask(self, *args, **kwargs) -> OrderedDataGroup:
        """This is a shorthand for
        :py:meth:`~dantro.data_loaders.load_hdf5.Hdf5LoaderMixin._load_hdf5`
        with the ``load_as_proxy`` flag set and ``resolve_as_dask`` passed as
        additional arguments to the proxy via ``proxy_kwargs``.
        """
        return self._load_hdf5(*args,
                               load_as_proxy=True,
                               proxy_kwargs=dict(resolve_as_dask=True),
                               **kwargs)
