# -*- coding: utf-8 -*-
"""
Tools to deal with HDF5 cache files
HDF5 cache files are used to cache Energy Database files, and Line Database
files, and yield a much faster access time.

Routine Listing
---------------


- :func:`~radis.misc.cache_file.check_cache_file`
- :func:`~radis.misc.cache_file.check_not_deprecated`
- :func:`~radis.misc.cache_file.save_to_hdf`

See Also
--------

:py:func:`~radis.io.hitran.hit2df`,
:py:func:`~radis.io.cdsd.cdsd2df`

-------------------------------------------------------------------------------


"""
# TODO: start using .feather format. Faster than HDF5 for pandas Dataframe,
# and can be multithreaded.
# see:
# https://gist.github.com/gansanay/4514ec731da1a40d8811a2b3c313f836
# and pd.read_feather(file, nthreads=3)


# Note: don't import unicode_literals because it breaks the df.to_hdf of
# save_to_hdf because of a stupid unicode/str error in Python 2.7
import os
from os.path import exists, splitext
from warnings import warn

import h5py
import pandas as pd
from packaging.version import parse

import radis
from radis import OLDEST_COMPATIBLE_VERSION
from radis.misc.basics import compare_dict, is_float
from radis.misc.printer import printm, printr


class DeprecatedFileError(DeprecationWarning):
    pass


"""str: forces to regenerate cache files that were created in a previous version"""

# Just make sure LAST_BACKWARD_COMPATIBLE_VERSION is valid
assert radis.__version__ >= OLDEST_COMPATIBLE_VERSION

# Utils


def load_h5_cache_file(
    cachefile,
    use_cached,
    metadata,
    current_version,
    last_compatible_version=OLDEST_COMPATIBLE_VERSION,
    verbose=True,
):
    """Function to load a h5 cache file

    Parameters
    ----------

    cachefile: str
        cache file path

    use_cached: str
        use cache file if value is not ``False``:

        - if ``True``, use (and generate if doesnt exist) cache file.
        - if ``'regen'``, delete cache file (if exists) so it is regenerated
        - if ``'force'``, use cache file and raises an error if it doesnt exist

        if using the cache file, check if the file is deprecated. If it
        is deprecated, regenerate the file unless ``'force'`` was used
        (in that case, raise an error)

    metadata: dict
        values are compared to cache file attributes. If they dont match,
        the file is considered deprecated. See ``use_cached`` to know
        how to handle deprecated files

    current_version: str
        version is compared to cache file version (part of attributes).
        If current version is superior, a simple warning is triggered.

    last_compatible_version: str
        if file version is inferior to this, file is considered deprecated.
        See ``use_cached`` to know how to handle deprecated files.
        Default :data:`~radis.OLDEST_COMPATIBLE_VERSION`.

    Returns
    -------

    df: pandas DataFrame, or None
        None if no cache file was found, or if it was deleted


    See Also
    --------

    :data:`~radis.OLDEST_COMPATIBLE_VERSION`

    """

    # 1. know if we have to load the file
    if not use_cached:
        return None
    elif use_cached == "regen" and exists(cachefile):
        os.remove(cachefile)
        if verbose:
            printm("Deleted h5 cache file : {0}".format(cachefile))
        return None

    # 2. check the file is here
    if not exists(cachefile):
        if use_cached == "force":
            raise ValueError("Cache file {0} doesnt exist".format(cachefile))
        else:
            return None  # File doesn't exist. It's okay.

    # 3. read file attributes to know if it's deprecated
    try:
        check_not_deprecated(
            cachefile,
            metadata,
            current_version=current_version,
            last_compatible_version=last_compatible_version,
        )
    # ... if deprecated, raise an error only if 'force'
    except DeprecatedFileError as err:
        if use_cached == "force":
            raise err
        else:
            if verbose:
                printr(
                    "File {0} deprecated:\n{1}\nDeleting it!".format(
                        cachefile, str(err)
                    )
                )
            os.remove(cachefile)
            return None
    # 4. File is not not deprecated: read the content.
    else:
        df = None
        if verbose >= 2:
            printm("Reading cache file ({0})".format(cachefile))
        try:
            df = pd.read_hdf(cachefile, "df")
        except KeyError as err:  # An error happened during file reading.
            # Fail safe by deleting cache file (unless we explicitely wanted it
            # with 'force')
            if use_cached == "force":
                raise
            else:
                if verbose:
                    printr(
                        "An error happened during cache file reading "
                        + "{0}:\n{1}\n".format(cachefile, str(err))
                        + "Deleting cache file to regenerate it"
                    )
                os.remove(cachefile)
                df = None

    return df


def get_cache_file(fcache, verbose=True):
    """Load HDF5 cache file

    Parameters
    ----------

    fcache: str
        file name

    Other Parameters
    ----------------

    verbose: bool
        If >=2, also warns if non numeric values are present (it would make
        calculations slower)

    Notes
    -----

    we could start using FEATHER format instead. See notes in cache_files.py
    """

    # Load file

    df = pd.read_hdf(fcache, "df")

    # Check file

    # ... 'object' columns slow everything down (not fixed format strings!)
    if verbose >= 2:
        _warn_if_object_columns(df, fcache)

    return df


def check_cache_file(fcache, use_cached=True, metadata={}, verbose=True):
    """Quick function that check status of cache file generated by RADIS:

    Parameters
    ----------

    fcache: str
        cache file name

    use_cached: ``True``, ``False``, ``'force'``, ``'regen'``
        see notes below. Default ``True``.

    metadata: dict
        attributes to check

    verbose: boolean
        print stuff

    Notes
    -----

    The function first checks the existence of ``fcache``. What is does depends
    on the value of ``use_cached``:

    - if ``True``, check it exists.
    - if ``'regen'``, delete cache file to regenerate it later.
    - if ``'force'``, raise an error if file doesnt exist.

    Then look if it is deprecated (we just look at the attributes, the file
    is never fully read). Deprecation is done by :py:func:`~radis.misc.cache_files.check_not_deprecated`
    comparing the ``metadata=`` content.

    - if deprecated, deletes it to regenerate later unless 'force' was used

    See Also
    --------

    :py:func:`~radis.misc.cache_files.check_not_deprecated`
    """

    # Test existence of file:
    if not use_cached:
        return  # we dont want a cache file, no need to test it
    elif use_cached == "regen":
        if exists(fcache):
            os.remove(fcache)
            if verbose:
                print(("Deleted h5 cache file : {0}".format(fcache)))
    elif use_cached == "force":
        if not exists(fcache):
            raise ValueError("Cache file {0} doesnt exist".format(fcache))
    else:  # use_cached == True
        pass  # just use the file as is

    # If file is still here, test if it is deprecated:
    # (we just read the attributes, the file is never fully read)
    if exists(fcache):
        if verbose:
            print(("Using cache file: {0}".format(fcache)))
        try:
            check_not_deprecated(
                fcache,
                metadata=metadata,
                current_version=radis.__version__,
                last_compatible_version=OLDEST_COMPATIBLE_VERSION,
            )
        except DeprecatedFileError as err:
            if use_cached == "force":
                raise
            else:  # delete file to regenerate it in the end of the script
                if verbose:
                    printr(
                        "File {0} deprecated:\n{1}\nDeleting it!".format(
                            fcache, str(err)
                        )
                    )
                os.remove(fcache)

    return


def check_not_deprecated(
    file,
    metadata,
    current_version=None,
    last_compatible_version=OLDEST_COMPATIBLE_VERSION,
):
    """Make sure cache file is not deprecated: checks that ``metadata`` is the same,
    and that the version under which the file was generated is valid.

    Parameters
    ----------

    file: str
        a `` .h5`` cache file for Energy Levels

    metadata: dict
        list of variables used to create the file. If the values dont match,
        an error is raised.

    current_version: str, or ``None``
        current version number. If the file was generated in a previous version
        a warning is raised. If ``None``, current version is read from
        :data:`radis.__version__`.

    last_backward_compatible_version: str
        If the file was generated in a non-compatible version, an error is raised.
        Default :py:data:`~radis.OLDEST_COMPATIBLE_VERSION`

    """

    # Get attributes (metadata+version)
    hf = h5py.File(file, "r")
    try:
        attrs = dict(hf.attrs)
    except OSError:
        attrs = {}
    finally:
        hf.close()

    # Raise an error if version is not found
    try:
        file_version = attrs.pop("version")
    except KeyError:
        raise DeprecatedFileError(
            "File {0} has been generated in a deprecated ".format(file)
            + "version. Delete it to regenerate it on next run"
        )

    # Get current version
    if current_version is None:
        current_version = radis.__version__

    # If file version is anterior to a major change
    # ... Update here versions afterwhich Deprecated cache file is not safe
    # ... (example: a key name was changed)
    if parse(file_version) < parse(last_compatible_version):
        raise DeprecatedFileError(
            "File {0} has been generated in a deprecated ".format(file)
            + "version ({0}). Oldest compatible version is {1}. ".format(
                file_version, last_compatible_version
            )
            + "Delete the file to regenerate it on next run"
        )

    # If file version is outdated: Warning, but no error
    if parse(current_version) > parse(file_version):
        warn(
            DeprecationWarning(
                "File {0} has been generated in ".format(file)
                + "a deprecated version ({0}) compared to current ({1})".format(
                    file_version, current_version
                )
                + ". Delete it to regenerate it on next run"
            )
        )
        out = False
    elif parse(current_version) == parse(file_version):
        out = True
    else:
        raise ValueError(
            "Cache file ({0}) generated with a future version ({1} > {2})? ".format(
                file, file_version, current_version
            )
            + "Do you own a DeLorean? Delete the file manually if you understand what happened"
        )

    # Compare metadata
    metadata = _h5_compatible(metadata)
    out, compare_string = compare_dict(
        metadata, attrs, verbose=False, return_string=True
    )
    if out != 1:
        raise DeprecatedFileError(
            "Metadata in file {0} dont match ".format(file)
            + "expected values. See comparison below:"
            + "\n\tExpected\tFile\n{0}".format(compare_string)
        )

    return out


def _h5_compatible(a_dict):
    """ Make dictionary ``a_dict`` h5 compatible """
    out = {}
    for k, v in a_dict.items():
        if v is None:
            continue  # dont store None
        elif is_float(v):
            out[k] = v
        else:
            out[k] = str(v)  # convert to str
    return out


def _warn_if_object_columns(df, fname):
    """  'object' columns slow everything down (not fixed format strings!) """
    objects = [k for k, v in df.dtypes.items() if v == object]
    if len(objects) > 0:
        warn(
            "Dataframe in {0} contains `object` format columns: {1}. ".format(
                fname, objects
            )
            + "Operations will be slower. Try to convert them to numeric."
        )


def save_to_hdf(
    df, fname, metadata, version=None, key="df", overwrite=True, verbose=True
):
    """Save energy levels to HDF5 file. Add metadata and version

    Parameters
    ----------

    df: a pandas DataFrame
        data will be stored in the key ``'df'``

    fname: str
        ``.h5`` file where to store.

    metadata: dict
         dictionary of values that were used to generate the DataFrame. Metadata
         will be asked again on file load to ensure it hasnt changed. ``None``
         values will not be stored.

    version: str, or ``None``
        file version. If ``None``, the current :data:`radis.__version__` is used.
        On file loading, a warning will be raised if the current version is
        posterior, or an error if the file version is set to be uncompatible.
        See :py:data:`~radis.OLDEST_COMPATIBLE_VERSION`

    key: str
        dataset name. Default ``'df'``

    overwrite: boolean
        if ``True``, overwrites file. Else, raise an error if it exists.

    Other Parameters
    ----------------

    verbose: bool
        If >=2, also warns if non numeric values are present (it would make
        calculations slower)

    Notes
    -----

    ``None`` values are not stored

    """

    # Check file
    assert fname.endswith(".h5")
    assert "version" not in metadata
    # ... 'object' columns slow everything down (not fixed format strings!)
    if verbose >= 2:
        _warn_if_object_columns(df, fname)

    # Update metadata format
    metadata = _h5_compatible(metadata)

    # Overwrite file
    if exists(fname) and not overwrite:
        raise ValueError("File exist: {0}".format(fname))
    hf = h5py.File(fname, "w")

    try:

        # Start by adding version
        if version is None:
            version = radis.__version__
        hf.attrs["version"] = version

        # Add metadata
        for k, v in metadata.items():
            hf.attrs[k] = v

        if verbose >= 3:
            print("... saved {0} with metadata: {1}".format(fname, metadata))

    except:
        raise

    finally:
        hf.close()

    # now export dataframe
    df.to_hdf(fname, key, format="fixed", mode="a", complevel=1, complib="blosc")


#    df.to_hdf(fname, 'df', format='fixed', mode='a')  # would be 10-20% faster, but take 2x more space


def filter_metadata(arguments, discard_variables=["self", "verbose"]):
    """Filter arguments (created with  ``locals()`` at the beginning
    of the script) to extract metadata. Metadata is stored as attributes
    in the cached file:

    - remove variables in ``discard_variables``
    - remove variables that start with ``'_'``
    - remove varibles whose value is ``None``

    Parameters
    ----------

    arguments: dict
        list of local variables. For instance::

            arguments = locals()

    discard_variables: list of str
        variable names to discard

    Returns
    -------

    metadata: dict
        a (new) dictionary built from arguments by removing ``discard_variables``
        and variables starting with ``'_'``

    Examples
    --------

    How to get only function argument::

        def some_function(*args):
            metadata = locals()     # stores only function arguments because it's the first line

            ...

            metadata = filter_metadata(metadata)
            save_to_hdf(df, fname, metadata=metadata)

            ...

    """

    metadata = {k: v for (k, v) in arguments.items() if not k.startswith("_")}
    metadata = {k: v for (k, v) in metadata.items() if k not in discard_variables}
    metadata = {k: v for (k, v) in metadata.items() if v is not None}

    return metadata


def cache_file_name(fname):
    """ Returns the corresponding cache file name for fname """
    return splitext(fname)[0] + ".h5"
