# coding=utf-8
# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.
#
# 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.

# Lint as: python3
"""Some python utils function and classes.

"""

import contextlib
import functools
import itertools
import os
import pickle
import sys
import types
from io import BytesIO as StringIO
from multiprocessing import Pool, RLock
from shutil import disk_usage
from types import CodeType, FunctionType
from typing import Callable, ClassVar, Generic, Optional, Tuple, Union

import dill
import numpy as np
from tqdm.auto import tqdm

from .. import utils
from . import logging


try:  # pragma: no branch
    import typing_extensions as _typing_extensions
    from typing_extensions import Final, Literal
except ImportError:
    _typing_extensions = Literal = Final = None


logger = logging.get_logger(__name__)


# NOTE: When used on an instance method, the cache is shared across all
# instances and IS NOT per-instance.
# See
# https://stackoverflow.com/questions/14946264/python-lru-cache-decorator-per-instance
# For @property methods, use @memoized_property below.
memoize = functools.lru_cache


def size_str(size_in_bytes):
    """Returns a human readable size string.

    If size_in_bytes is None, then returns "Unknown size".

    For example `size_str(1.5 * datasets.units.GiB) == "1.50 GiB"`.

    Args:
        size_in_bytes: `int` or `None`, the size, in bytes, that we want to
            format as a human-readable size string.
    """
    if not size_in_bytes:
        return "Unknown size"

    _NAME_LIST = [("PiB", 2 ** 50), ("TiB", 2 ** 40), ("GiB", 2 ** 30), ("MiB", 2 ** 20), ("KiB", 2 ** 10)]

    size_in_bytes = float(size_in_bytes)
    for (name, size_bytes) in _NAME_LIST:
        value = size_in_bytes / size_bytes
        if value >= 1.0:
            return "{:.2f} {}".format(value, name)
    return "{} {}".format(int(size_in_bytes), "bytes")


@contextlib.contextmanager
def temporary_assignment(obj, attr, value):
    """Temporarily assign obj.attr to value."""
    original = getattr(obj, attr, None)
    setattr(obj, attr, value)
    try:
        yield
    finally:
        setattr(obj, attr, original)


def zip_dict(*dicts):
    """Iterate over items of dictionaries grouped by their keys."""
    for key in set(itertools.chain(*dicts)):  # set merge all keys
        # Will raise KeyError if the dict don't have the same keys
        yield key, tuple(d[key] for d in dicts)


class NonMutableDict(dict):
    """Dict where keys can only be added but not modified.

    Will raise an error if the user try to overwrite one key. The error message
    can be customized during construction. It will be formatted using {key} for
    the overwritten key.
    """

    def __init__(self, *args, **kwargs):
        self._error_msg = kwargs.pop(
            "error_msg",
            "Try to overwrite existing key: {key}",
        )
        if kwargs:
            raise ValueError("NonMutableDict cannot be initialized with kwargs.")
        super(NonMutableDict, self).__init__(*args, **kwargs)

    def __setitem__(self, key, value):
        if key in self:
            raise ValueError(self._error_msg.format(key=key))
        return super(NonMutableDict, self).__setitem__(key, value)

    def update(self, other):
        if any(k in self for k in other):
            raise ValueError(self._error_msg.format(key=set(self) & set(other)))
        return super(NonMutableDict, self).update(other)


class classproperty(property):  # pylint: disable=invalid-name
    """Descriptor to be used as decorator for @classmethods."""

    def __get__(self, obj, objtype=None):
        return self.fget.__get__(None, objtype)()


def _single_map_nested(args):
    """Apply a function recursively to each element of a nested data struct."""
    function, data_struct, types, rank, disable_tqdm = args

    # Singleton first to spare some computation
    if not isinstance(data_struct, dict) and not isinstance(data_struct, types):
        return function(data_struct)

    # Reduce logging to keep things readable in multiprocessing with tqdm
    if rank is not None and logging.get_verbosity() < logging.WARNING:
        logging.set_verbosity_warning()
    # Print at least one thing to fix tqdm in notebooks in multiprocessing
    # see https://github.com/tqdm/tqdm/issues/485#issuecomment-473338308
    if rank is not None and not disable_tqdm and "notebook" in tqdm.__name__:
        print(" ", end="", flush=True)

    # Loop over single examples or batches and write to buffer/file if examples are to be updated
    pbar_iterable = data_struct.items() if isinstance(data_struct, dict) else data_struct
    pbar_desc = "#" + str(rank) if rank is not None else None
    pbar = utils.tqdm(pbar_iterable, disable=disable_tqdm, position=rank, unit="obj", desc=pbar_desc)

    if isinstance(data_struct, dict):
        return {k: _single_map_nested((function, v, types, None, True)) for k, v in pbar}
    else:
        mapped = [_single_map_nested((function, v, types, None, True)) for v in pbar]
        if isinstance(data_struct, list):
            return mapped
        elif isinstance(data_struct, tuple):
            return tuple(mapped)
        else:
            return np.array(mapped)


def map_nested(
    function,
    data_struct,
    dict_only: bool = False,
    map_list: bool = True,
    map_tuple: bool = False,
    map_numpy: bool = False,
    num_proc: Optional[int] = None,
    types=None,
):
    """Apply a function recursively to each element of a nested data struct.
    If num_proc > 1 and the length of data_struct is longer than num_proc: use multi-processing
    """
    if types is None:
        types = []
        if not dict_only:
            if map_list:
                types.append(list)
            if map_tuple:
                types.append(tuple)
            if map_numpy:
                types.append(np.ndarray)
        types = tuple(types)

    # Singleton
    if not isinstance(data_struct, dict) and not isinstance(data_struct, types):
        return function(data_struct)

    disable_tqdm = bool(logger.getEffectiveLevel() > logging.INFO) or not utils.is_progress_bar_enabled()
    iterable = list(data_struct.values()) if isinstance(data_struct, dict) else data_struct

    if num_proc is None:
        num_proc = 1
    if num_proc <= 1 or len(iterable) <= num_proc:
        mapped = [
            _single_map_nested((function, obj, types, None, True))
            for obj in utils.tqdm(iterable, disable=disable_tqdm)
        ]
    else:
        split_kwds = []  # We organize the splits ourselve (contiguous splits)
        for index in range(num_proc):
            div = len(iterable) // num_proc
            mod = len(iterable) % num_proc
            start = div * index + min(index, mod)
            end = start + div + (1 if index < mod else 0)
            split_kwds.append((function, iterable[start:end], types, index, disable_tqdm))
        assert len(iterable) == sum(len(i[1]) for i in split_kwds), (
            f"Error dividing inputs iterable among processes. "
            f"Total number of objects {len(iterable)}, "
            f"length: {sum(len(i[1]) for i in split_kwds)}"
        )
        logger.info(
            "Spawning {} processes for {} objects in slices of {}".format(
                num_proc, len(iterable), [len(i[1]) for i in split_kwds]
            )
        )
        initargs, initializer = None, None
        if not disable_tqdm:
            initargs, initializer = (RLock(),), tqdm.set_lock
        with Pool(num_proc, initargs=initargs, initializer=initializer) as pool:
            mapped = pool.map(_single_map_nested, split_kwds)
        logger.info("Finished {} processes".format(num_proc))
        mapped = [obj for proc_res in mapped for obj in proc_res]
        logger.info("Unpacked {} objects".format(len(mapped)))

    if isinstance(data_struct, dict):
        return dict(zip(data_struct.keys(), mapped))
    else:
        if isinstance(data_struct, list):
            return mapped
        elif isinstance(data_struct, tuple):
            return tuple(mapped)
        else:
            return np.array(mapped)


def zip_nested(arg0, *args, **kwargs):
    """Zip data struct together and return a data struct with the same shape."""
    # Python 2 do not support kwargs only arguments
    dict_only = kwargs.pop("dict_only", False)
    assert not kwargs

    # Could add support for more exotic data_struct, like OrderedDict
    if isinstance(arg0, dict):
        return {k: zip_nested(*a, dict_only=dict_only) for k, a in zip_dict(arg0, *args)}
    elif not dict_only:
        if isinstance(arg0, list):
            return [zip_nested(*a, dict_only=dict_only) for a in zip(arg0, *args)]
    # Singleton
    return (arg0,) + args


def flatten_nest_dict(d):
    """Return the dict with all nested keys flattened joined with '/'."""
    # Use NonMutableDict to ensure there is no collision between features keys
    flat_dict = NonMutableDict()
    for k, v in d.items():
        if isinstance(v, dict):
            flat_dict.update({"{}/{}".format(k, k2): v2 for k2, v2 in flatten_nest_dict(v).items()})
        else:
            flat_dict[k] = v
    return flat_dict


class NestedDataStructure:
    def __init__(self, data=None):
        self.data = data if data is not None else []

    def flatten(self, data=None):
        data = data if data is not None else self.data
        if isinstance(data, dict):
            return self.flatten(list(data.values()))
        elif isinstance(data, (list, tuple)):
            return [flattened for item in data for flattened in self.flatten(item)]
        else:
            return [data]


def has_sufficient_disk_space(needed_bytes, directory="."):
    try:
        free_bytes = disk_usage(os.path.abspath(directory)).free
    except OSError:
        return True
    return needed_bytes < free_bytes


class Pickler(dill.Pickler):
    """Same Pickler as the one from dill, but improved for notebooks and shells"""

    dispatch = dill._dill.MetaCatchingDict(dill.Pickler.dispatch.copy())

    def save_global(self, obj, name=None):
        if sys.version_info[:2] < (3, 7) and _CloudPickleTypeHintFix._is_parametrized_type_hint(
            obj
        ):  # noqa  # pragma: no branch
            # Parametrized typing constructs in Python < 3.7 are not compatible
            # with type checks and ``isinstance`` semantics. For this reason,
            # it is easier to detect them using a duck-typing-based check
            # (``_is_parametrized_type_hint``) than to populate the Pickler's
            # dispatch with type-specific savers.
            _CloudPickleTypeHintFix._save_parametrized_type_hint(self, obj)
        else:
            dill.Pickler.save_global(self, obj, name=name)


def dump(obj, file):
    """pickle an object to a file"""
    Pickler(file, recurse=True).dump(obj)
    return


@contextlib.contextmanager
def _no_cache_fields(obj):
    try:
        if (
            "PreTrainedTokenizerBase" in [base_class.__name__ for base_class in type(obj).__mro__]
            and hasattr(obj, "cache")
            and isinstance(obj.cache, dict)
        ):
            with temporary_assignment(obj, "cache", {}):
                yield
        else:
            yield

    except ImportError:
        yield


def dumps(obj):
    """pickle an object to a string"""
    file = StringIO()
    with _no_cache_fields(obj):
        dump(obj, file)
    return file.getvalue()


def pklregister(t):
    def proxy(func):
        Pickler.dispatch[t] = func
        return func

    return proxy


class _CloudPickleTypeHintFix:
    """
    Type hints can't be properly pickled in python < 3.7
    CloudPickle provided a way to make it work in older versions.
    This class provide utilities to fix pickling of type hints in older versions.
    from https://github.com/cloudpipe/cloudpickle/pull/318/files
    """

    def _is_parametrized_type_hint(obj):
        # This is very cheap but might generate false positives.
        origin = getattr(obj, "__origin__", None)  # typing Constructs
        values = getattr(obj, "__values__", None)  # typing_extensions.Literal
        type_ = getattr(obj, "__type__", None)  # typing_extensions.Final
        return origin is not None or values is not None or type_ is not None

    def _create_parametrized_type_hint(origin, args):
        return origin[args]

    def _save_parametrized_type_hint(pickler, obj):
        # The distorted type check sematic for typing construct becomes:
        # ``type(obj) is type(TypeHint)``, which means "obj is a
        # parametrized TypeHint"
        if type(obj) is type(Literal):  # pragma: no branch
            initargs = (Literal, obj.__values__)
        elif type(obj) is type(Final):  # pragma: no branch
            initargs = (Final, obj.__type__)
        elif type(obj) is type(ClassVar):
            initargs = (ClassVar, obj.__type__)
        elif type(obj) in [type(Union), type(Tuple), type(Generic)]:
            initargs = (obj.__origin__, obj.__args__)
        elif type(obj) is type(Callable):
            args = obj.__args__
            if args[0] is Ellipsis:
                initargs = (obj.__origin__, args)
            else:
                initargs = (obj.__origin__, (list(args[:-1]), args[-1]))
        else:  # pragma: no cover
            raise pickle.PicklingError("Datasets pickle Error: Unknown type {}".format(type(obj)))
        pickler.save_reduce(_CloudPickleTypeHintFix._create_parametrized_type_hint, initargs, obj=obj)


@pklregister(CodeType)
def _save_code(pickler, obj):
    """
    From dill._dill.save_code
    This is a modified version that removes the origin (filename + line no.)
    of functions created in notebooks or shells for example.
    """
    dill._dill.log.info("Co: %s" % obj)
    # The filename of a function is the .py file where it is defined.
    # Filenames of functions created in notebooks or shells start with '<'
    # ex: <ipython-input-13-9ed2afe61d25> for ipython, and <stdin> for shell
    # Moreover lambda functions have a special name: '<lambda>'
    # ex: (lambda x: x).__code__.co_name == "<lambda>"  # True
    # For the hashing mechanism we ignore where the function has been defined
    # More specifically:
    # - we ignore the filename of special functions (filename starts with '<')
    # - we always ignore the line number
    # Only those two lines are different from the original implementation:
    co_filename = "" if obj.co_filename.startswith("<") or obj.co_name == "<lambda>" else obj.co_filename
    co_firstlineno = 1
    # The rest is the same as in the original dill implementation
    if dill._dill.PY3:
        if hasattr(obj, "co_posonlyargcount"):
            args = (
                obj.co_argcount,
                obj.co_posonlyargcount,
                obj.co_kwonlyargcount,
                obj.co_nlocals,
                obj.co_stacksize,
                obj.co_flags,
                obj.co_code,
                obj.co_consts,
                obj.co_names,
                obj.co_varnames,
                co_filename,
                obj.co_name,
                co_firstlineno,
                obj.co_lnotab,
                obj.co_freevars,
                obj.co_cellvars,
            )
        else:
            args = (
                obj.co_argcount,
                obj.co_kwonlyargcount,
                obj.co_nlocals,
                obj.co_stacksize,
                obj.co_flags,
                obj.co_code,
                obj.co_consts,
                obj.co_names,
                obj.co_varnames,
                co_filename,
                obj.co_name,
                co_firstlineno,
                obj.co_lnotab,
                obj.co_freevars,
                obj.co_cellvars,
            )
    else:
        args = (
            obj.co_argcount,
            obj.co_nlocals,
            obj.co_stacksize,
            obj.co_flags,
            obj.co_code,
            obj.co_consts,
            obj.co_names,
            obj.co_varnames,
            co_filename,
            obj.co_name,
            co_firstlineno,
            obj.co_lnotab,
            obj.co_freevars,
            obj.co_cellvars,
        )
    pickler.save_reduce(CodeType, args, obj=obj)
    dill._dill.log.info("# Co")
    return


@pklregister(FunctionType)
def save_function(pickler, obj):
    """
    From dill._dill.save_function
    This is a modified version that make globs deterministic since the order of
    the keys in the output dictionary of globalvars can change.
    """
    if not dill._dill._locate_function(obj):
        dill._dill.log.info("F1: %s" % obj)
        if getattr(pickler, "_recurse", False):
            # recurse to get all globals referred to by obj
            globalvars = dill.detect.globalvars
            globs = globalvars(obj, recurse=True, builtin=True)
            if id(obj) in dill._dill.stack:
                globs = obj.__globals__ if dill._dill.PY3 else obj.func_globals
        else:
            globs = obj.__globals__ if dill._dill.PY3 else obj.func_globals
        # globs is a dictionary with keys = var names (str) and values = python objects
        # however the dictionary is not always loaded in the same order
        # therefore we have to sort the keys to make deterministic.
        # This is important to make `dump` deterministic.
        # Only this line is different from the original implementation:
        globs = {k: globs[k] for k in sorted(globs.keys())}
        # The rest is the same as in the original dill implementation
        _byref = getattr(pickler, "_byref", None)
        _recurse = getattr(pickler, "_recurse", None)
        _memo = (id(obj) in dill._dill.stack) and (_recurse is not None)
        dill._dill.stack[id(obj)] = len(dill._dill.stack), obj
        if dill._dill.PY3:
            _super = ("super" in getattr(obj.__code__, "co_names", ())) and (_byref is not None)
            if _super:
                pickler._byref = True
            if _memo:
                pickler._recurse = False
            fkwdefaults = getattr(obj, "__kwdefaults__", None)
            pickler.save_reduce(
                dill._dill._create_function,
                (obj.__code__, globs, obj.__name__, obj.__defaults__, obj.__closure__, obj.__dict__, fkwdefaults),
                obj=obj,
            )
        else:
            _super = (
                ("super" in getattr(obj.func_code, "co_names", ()))
                and (_byref is not None)
                and getattr(pickler, "_recurse", False)
            )
            if _super:
                pickler._byref = True
            if _memo:
                pickler._recurse = False
            pickler.save_reduce(
                dill._dill._create_function,
                (obj.func_code, globs, obj.func_name, obj.func_defaults, obj.func_closure, obj.__dict__),
                obj=obj,
            )
        if _super:
            pickler._byref = _byref
        if _memo:
            pickler._recurse = _recurse
        if (
            dill._dill.OLDER
            and not _byref
            and (_super or (not _super and _memo) or (not _super and not _memo and _recurse))
        ):
            pickler.clear_memo()
        dill._dill.log.info("# F1")
    else:
        dill._dill.log.info("F2: %s" % obj)
        name = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
        dill._dill.StockPickler.save_global(pickler, obj, name=name)
        dill._dill.log.info("# F2")
    return


def copyfunc(func):
    result = types.FunctionType(func.__code__, func.__globals__, func.__name__, func.__defaults__, func.__closure__)
    result.__kwdefaults__ = func.__kwdefaults__
    return result


try:
    import regex

    @pklregister(type(regex.Regex("", 0)))
    def _save_regex(pickler, obj):
        dill._dill.log.info("Re: %s" % obj)
        args = (
            obj.pattern,
            obj.flags,
        )
        pickler.save_reduce(regex.compile, args, obj=obj)
        dill._dill.log.info("# Re")
        return


except ImportError:
    pass
