"""Module to define various calculation types as Enums for VASP."""

from __future__ import annotations
from importlib.resources import files as import_resource_files
from itertools import product

from ruamel.yaml import YAML

_BASE_ENUM_PATH = {
    "vasp": import_resource_files("emmet.core.vasp"),
    "qchem": import_resource_files("emmet.core.qchem"),
}

for code_base in _BASE_ENUM_PATH:
    _BASE_ENUM_PATH[code_base] /= "calc_types"


def get_enum_source(
    enum_name: str, doc: str, members: dict, enum_class: str = "ValueEnum"
) -> str:
    """Write python-format enum from a dict of members and metadata.

    Parameters
    -----------
    enum_name : str
        Name of the enum.
    doc : str
        Enum docstr
    members : dict
        The key-value pair indexed members of the enum.
    enum_class : str = "ValueEnum"
        The name of the enum class.

    Returns
    --------
    str
    """
    items = [f"class {enum_name}({enum_class}):", f'    """ {doc} """\n']
    items += [f'    {const} = "{val}"' for const, val in members.items()]
    return "\n".join(items)


def string_bulk_replace(string: str, rules: dict[str, str]) -> str:
    """Perform multiple string replacements subject to a set of rules.

    Parameters
    -----------
    string : str
        input string
    rules : dict[str,str]
        A dict of string replacements, with the characters to replace
        as keys, and their replacements as values.

    Returns
    --------
    str
    """
    for targ_char, rep_char in rules.items():
        string = string.replace(targ_char, rep_char)
    return string


def generate_vasp_enums_file(enum_file_name: str | None = None) -> None:
    """
    Generate VASP enum members from reference yaml data.

    Parameters
    -----------
    enum_file_name : str
        Name of the file to write the enums to.
        Defaults to _BASE_ENUM_PATH / vasp_enums.json.gz
    """

    with open(_BASE_ENUM_PATH["vasp"] / "calc_types.yaml", "r") as config:
        _RUN_TASK_TYPE_DATA = YAML().load(config)

    _TASK_TYPES = _RUN_TASK_TYPE_DATA.get("TASK_TYPES")

    _RUN_TYPES = set(
        rt
        for functionals in _RUN_TASK_TYPE_DATA.get("RUN_TYPES", {}).values()
        for rt in functionals
    )
    _RUN_TYPES.update(set(f"{rt}+U" for rt in _RUN_TYPES))

    _ENUMS = {
        "RunType": {
            "_".join(rt.split()).replace("+", "_").replace("-", "_"): rt
            for rt in _RUN_TYPES
        },
        "TaskType": {"_".join(tt.split()): tt for tt in _TASK_TYPES},
        "CalcType": {
            f"{'_'.join(rt.split()).replace('+','_').replace('-','_')}"
            f"_{'_'.join(tt.split())}": f"{rt} {tt}"
            for rt, tt in product(_RUN_TYPES, _TASK_TYPES)
        },
    }

    docstr = {}
    for enum_name in _ENUMS:
        rtc_type = enum_name.split("Calc")[-1].split("Type")[0].lower()
        if len(rtc_type) > 0:
            rtc_type += " "
        docstr[enum_name] = f"VASP calculation {rtc_type}types."

    enum_file_name = enum_file_name or str(_BASE_ENUM_PATH["vasp"] / "enums.py")
    with open(enum_file_name, "w+") as f:
        f.write(
            """\"\"\"
Autogenerated Enums for VASP RunType, TaskType, and CalcType.

Do not edit this by hand to add or remove enums.
Instead, edit
    dev_scripts/generate_enums.py
and/or
    emmet/core/vasp/calc_types/calc_types.yaml
\"\"\"
from emmet.core.types.enums import ValueEnum, IgnoreCaseEnum

"""
        )
        enum_order = (
            "RunType",
            "TaskType",
            "CalcType",
        )
        for ienum, enum_name in enumerate(enum_order):
            sorted_enums = {k: _ENUMS[enum_name][k] for k in sorted(_ENUMS[enum_name])}
            f.write(
                get_enum_source(
                    enum_name,
                    docstr[enum_name],
                    sorted_enums,
                    enum_class=(
                        "IgnoreCase"
                        if enum_name in ["RunType", "CalcType"]
                        else "Value"
                    )
                    + "Enum",
                )
            )
            f.write("\n\n" if ienum < (len(enum_order) - 1) else "\n")


def generate_qchem_enum_file(enum_file_name: str | None = None) -> None:
    """
    Generate QChem enum members from reference yaml data.

    Original author, Evan Spotte-Smith <ewcspottesmith@lbl.gov>

    Parameters
    -----------
    enum_file_name : str
        Name of the file to write the enums to.
        Defaults to _BASE_ENUM_PATH / qchem_enums.json.gz
    """

    with open(_BASE_ENUM_PATH["qchem"] / "calc_types.yaml", "r") as config:
        _calc_type_meta = YAML().load(config)

    _calc_type_meta["FUNCTIONALS"] = [
        rt
        for functionals in _calc_type_meta["FUNCTIONAL_CLASSES"].values()
        for rt in functionals
    ]

    _LOTS = list()

    for funct in _calc_type_meta["FUNCTIONALS"]:
        for basis in _calc_type_meta["BASIS_SETS"]:
            for solv_model in _calc_type_meta["SOLVENT_MODELS"]:
                _LOTS.append(f"{funct}/{basis}/{solv_model}")

    _lot_str_replacements = {
        "+": "_",
        "-": "_",
        "(": "_",
        ")": "_",
        "/": "_",
        "*": "_d",
    }

    _ENUMS = {
        "LevelOfTheory": {
            "_".join(string_bulk_replace(lot, _lot_str_replacements).split()): lot
            for lot in _LOTS
        },
        "TaskType": {
            "_".join(tt.split()).replace("-", "_"): tt
            for tt in _calc_type_meta["TASK_TYPES"]
        },
        "CalcType": {
            (
                "_".join(string_bulk_replace(lot, _lot_str_replacements).split())
                + f"_{'_'.join(tt.split()).replace('-', '_')}"
            ): f"{lot} {tt}"
            for lot, tt in product(_LOTS, _calc_type_meta["TASK_TYPES"])
        },
    }

    docstr = {
        "LevelOfTheory": "Levels of theory for calculations in Q-Chem.",
        "TaskType": "Calculation task types for Q-Chem.",
        "CalcType": "Calculation types (LOT + task type) for Q-Chem.",
    }

    enum_file_name = enum_file_name or str(_BASE_ENUM_PATH["qchem"] / "enums.py")

    with open(enum_file_name, "w+") as f:
        f.write(
            """\"\"\"
Autogenerated Enums for Q-Chem LevelOfTheory, TaskType, and CalcType.

Do not edit this by hand to add or remove enums.
Instead, edit
    dev_scripts/generate_enums.py
and/or
    emmet/core/qchem/calc_types/calc_types.yaml
\"\"\"
from emmet.core.types.enums import ValueEnum

"""
        )
        enum_order = (
            "LevelOfTheory",
            "TaskType",
            "CalcType",
        )
        for ienum, enum_name in enumerate(enum_order):
            sorted_enums = {k: _ENUMS[enum_name][k] for k in sorted(_ENUMS[enum_name])}
            f.write(get_enum_source(enum_name, docstr[enum_name], sorted_enums))
            f.write("\n\n" if ienum < (len(enum_order) - 1) else "\n")


if __name__ == "__main__":
    generate_vasp_enums_file()
    generate_qchem_enum_file()
