"""
CASTEP HELPER
Check for errors in input dictionary
{"CELL": {...},
 "PARAM": {...}

"""

from __future__ import absolute_import
from __future__ import print_function
import os
import json
from glob import glob

import logging
from six.moves import zip
logger = logging.getLogger(__name__)

path = os.path.abspath(__file__)
module_path = os.path.dirname(path)

incompatible_keys = [
    ('backup_interval', 'num_backup_iter'),
    ('opt_strategy', 'opt_strategy_bias'),
    ('basis_precision', 'cut_off_energy'),
    ('fine_grid_scale', 'fine_gmax', 'find_cut_off_energy'),
    ('nelectrons', 'charge'),
    ('nelectrons', 'nup'),
    ('nelectrons', 'ndown'),
    ('charge', 'ndown'),
    ('charge', 'nup'),
    ('spin', 'nup'),
    ('spin', 'ndown'),
    ('nextra_bands', 'perc_extra_bands'),
    ('metals_method', 'elec_method'),
]


class HelperCheckError(RuntimeError):
    pass


class CastepHelper(object):
    """
    A class for helping castep inputs
    """
    def __init__(self, version=None):
        # Instance level parameter for by passing any check at all
        self.BY_PASS = False
        self._HELP_DICT = None
        self.version = version
        # Try to load th helper dictionary
        self.load_helper_dict()

    def load_helper_dict(self):
        """
        Attempts to load a dictionary containing the helper information.
        This dictionary should be generated by the .generate module
        We look for .castep_helpinfo.json in the $HOME directory.
        """
        pairs = find_help_info()
        path = None
        for p, v in pairs:
            if v == self.version:
                path = p
                break
        # If we did not found a suitable version
        if path is None:
            if self.version is not None:
                raise RuntimeError(
                    "Cannot found help info for requested version {}".format(
                        self.version))
            try:
                path = pairs[0][
                    0]  # No explicitly requested - use the first one
            except IndexError:
                print("No CASTEP help info detected")
                self.BY_PASS = True
                help_dict = {}
            else:
                help_dict = load_json(path)
        else:
            help_dict = load_json(path)
        # Save the help_dict as a class attribute
        self._HELP_DICT = help_dict

    @property
    def help_dict(self):
        """A dictionary containing information of the helper"""
        return self._HELP_DICT

    @property
    def castep_help_version(self):
        """Version number of CASTEP that the help info is for"""
        return self.help_dict["_CASTEP_VERSION"]

    def save_helper_dict(self, help_dict, file_path):
        """
        Save the helper json file in to sensible location. By default it is here this module is
        """

        try:
            with open(file_path, "w") as fp:
                json.dump(help_dict, fp)
        except OSError:
            try:
                with ("castep_helpinfo.json", "w") as fp:
                    json.dump(help_dict, fp)
                    file_path = os.path.realpath("castep_helpinfo.json")
                    print(("Saving in current path. "
                           "Please move it to {}".format(file_path)))
            except OSError:
                print("Cannot save the retrieved help information")
                return

        print(("\n\nJSON file saved in {}".format(file_path)))

    def _check_dict(self, input_dict):
        """
        Check a dictionary of inputs. Return invalid and wrong keys

        :param dict input_dict: A dictonary with "PARAM" and "CELL" keys.

        :return: A list of [invalid_keys, wrong_keys].
        wrong_keys is a tuple of where the key should have being.
        (key, "PARAM") or (key, "CELL")
        """
        invalid_keys = []
        wrong_keys = []  # a list of tulple (key, "PARAM") or (key, "CELL")

        if self.BY_PASS:
            return invalid_keys, wrong_keys

        for kwtype in input_dict:
            # Check each key
            for key in input_dict[kwtype]:
                # key maybe be both lower or upper case
                info = self.help_dict.get(key.lower(), None)

                if info is None:
                    invalid_keys.append(key)
                    continue

                # Check if the type is correct
                if info["key_type"] != kwtype:
                    wrong_keys.append((key, info["key_type"]))
                    continue

        return invalid_keys, wrong_keys

    def _from_flat_dict(self, input_dict):
        """
        Construct a {"PARAM":{}, "CELL":{}} dictionary from dictionary with top-level keys other
        than PARAM and CELL
        :returns: a list of [out_dict, keys_not_found]
        """

        if self.BY_PASS:
            raise RuntimeError(
                "Cannot construct dictionary - No help info found")

        hinfo = self.help_dict

        # extract copies of CELL and PARAM fields from the input
        cell_dict = dict(input_dict.get("CELL", {}))
        param_dict = dict(input_dict.get("PARAM", {}))
        not_found = []
        for key in input_dict:
            key_entry = hinfo.get(key, None)

            if key_entry is None:
                not_found.append(key)
            else:
                kwtype = key_entry.get("key_type")
                if kwtype == "CELL":
                    cell_dict.update({key: input_dict[key]})
                elif kwtype == "PARAM":
                    param_dict.update({key: input_dict[key]})
                else:
                    raise RuntimeError(
                        "Entry {} does not have key_type value".format(key))

        out_dict = {"CELL": cell_dict, "PARAM": param_dict}
        return out_dict, not_found

    def check_dict(self, input_dict, auto_fix=True, allow_flat=False):
        """
        Check input dictionary. Apply and warn about errors
        :param input_dict dictionary: a dictionary as the input, contain "CELL", "PARAM" and other keywords
        :param auto_fix bool: Whether we should fix error automatically
        :param allow_flat: Accept that the input dictionary is flat.

        :returns dict: A structured dictionary
        """
        input_dict = input_dict.copy()  # this is a shallow copy

        # construct what to be checked
        cell_dict = input_dict.pop("CELL", {})
        param_dict = input_dict.pop("PARAM", {})

        if input_dict and not auto_fix and not allow_flat:
            raise HelperCheckError("keywords: {} at top level".format(
                ", ".join(input_dict)))

        # Following functions require help info to be defined
        if self.BY_PASS:
            logger.warning("No help info found - input not checked")
            return input_dict

        # process what's left
        re_structured, not_found = self._from_flat_dict(input_dict)
        if not_found:
            suggests = [self.get_suggestion(s) for s in not_found]
            # Warnings
            not_reco = [
                "keyword '{}' is not recognized at top level".format(s)
                for s in not_found
            ]
            # Suggestions
            sugst_str = [a + "\n" + b for a, b in zip(not_reco, suggests)]
            # Combine warnings and suggestions together
            sugst_str = "\n\n".join(sugst_str)
            raise HelperCheckError(sugst_str)

        # Now construct a dictionary
        cell_dict.update(re_structured["CELL"])
        param_dict.update(re_structured["PARAM"])

        # Construct an new dictionary
        input_dict = dict(CELL=cell_dict, PARAM=param_dict)

        # Check the restructed dictionary
        invalid, wrong = self._check_dict(input_dict)

        if invalid:
            suggests = [self.get_suggestion(s) for s in invalid]
            not_founds = [
                "keyword '{}' is not found".format(s) for s in invalid
            ]
            sugst_str = [a + "\n" + b for a, b in zip(not_founds, suggests)]
            sugst_str = "\n\n".join(sugst_str)

            raise HelperCheckError(sugst_str)

        # if there are still wrong keywords, fix them
        if wrong:
            if auto_fix is True:
                for key, should_be in wrong:
                    if should_be == "PARAM":
                        logger.warning("Key {} moved to PARAM".format(key))
                        value = input_dict["CELL"].pop(key)
                        input_dict["PARAM"].update({key: value})
                    else:
                        logger.warning("Key {} moved to CELL".format(key))
                        value = input_dict["PARAM"].pop(key)
                        input_dict["CELL"].update({key: value})
            else:
                raise HelperCheckError(
                    "Keywords: {} are in "
                    "the wrong sub-dictionary".format(", ".join(
                        ["'{}'".format(k[0]) for k in wrong])))

        # Check incompatible keys
        incomp = check_incompatible(input_dict["PARAM"], incompatible_keys)
        if incomp:
            raise HelperCheckError("Incompatible keys found: {}"
                                   " - only one is allowed.".format(incomp))
        return input_dict

    def get_suggestion(self, string):
        """
        Return string for suggestion of the string
        """
        return _get_suggestion(string, list(self.help_dict.keys()))


def _get_suggestion(provided_string, allowed_strings):
    """
    Given a string and a list of allowed_strings, it returns a string to print
    on screen, with sensible text depending on whether no suggestion is found,
    or one or more than one suggestions are found.

    :param provided_string: the string to compare

    :param allowed_strings: a list of valid strings

    :return: A string to print on output, to suggest to the user
    a possible valid value.
    """
    import difflib

    similar_kws = difflib.get_close_matches(provided_string, allowed_strings)
    if len(similar_kws) == 1:
        return "(Maybe you wanted to specify {0}?)".format(similar_kws[0])
    elif len(similar_kws) > 1:
        return "(Maybe you wanted to specify one of these: {0}?)".format(
            ", ".join(similar_kws))
    else:
        return "(No similar keywords found...)"


def load_json(path):
    """
    Load a json via file path
    """
    with open(path) as f:
        res = json.load(f)
    return res


def check_incompatible(dict_in, inc_list):
    """
    Check any conflicting keys in the dictionary
    """
    for incomp in inc_list:
        c = 0
        for k in dict_in:
            if k in incomp:
                c += 1
        if c > 1:
            return incomp
    return


def find_help_info():
    """
    Return possible paths of helper dict
    The diction should store in the format:
    castep_help_info_<version>.json
    Search paths are the module path and the $HOME

    :returns: [path1, path2], [verion1, version2].....
    """
    home = os.getenv("HOME")
    default_file_path = os.path.join(module_path, "*castep_help_info_*.json")
    others = glob(os.path.join(home, ".*castep_help_info_*.json"))
    paths = glob(default_file_path) + others
    vs = []
    for p in paths:
        tmp = p.replace(".json", "").split("_")
        if len(tmp) == 4:
            vs.append(float(tmp[-1]))
        else:
            vs.append(0)
    comb = list(zip(paths, vs))
    comb.sort(key=lambda x: x[1], reverse=True)
    return comb
