"""
Parameter related utils

Contains utils and definitions that are used together with the parameters.
"""
# pylint: disable=too-many-branches

from __future__ import annotations

import enum
from copy import deepcopy
from pathlib import Path
from typing import Any
from warnings import warn

from aiida import orm
from aiida.common.extendeddicts import AttributeDict
from yaml import safe_load

from aiida_vasp.utils.extended_dicts import update_nested_dict

_BASE_NAMESPACES: list[str] = ['electronic', 'smearing', 'charge', 'dynamics', 'bands', 'relax', 'converge']
_DEFAULT_OVERRIDE_NAMESPACE: str = 'incar'


class ChargeEnum(enum.IntEnum):
    """
    Encode values for the initial charge density.

    See: https://www.vasp.at/wiki/index.php/ICHARG
    """

    WAVE = 0
    CHARGE = 1
    ATOMIC = 2
    POTENTIAL = 4
    CONSTANT_CHARGE = 11
    CONSTANT_ATOMIC = 12


class IntSmearingEnum(enum.IntEnum):
    """
    Encode values for the smearing used during integration in reciprocal space.

    See: https://www.vasp.at/wiki/index.php/ISMEAR.
    """

    MP = 1  # pylint: disable=invalid-name
    GAUSSIAN = 0
    FERMI = -1
    PARTIAL = -2
    TETRA = -5


class OrbitEnum(enum.IntEnum):
    """
    Encode values for the projector information.

    See: https://www.vasp.at/wiki/index.php/LORBIT
    """

    ATOM = 0
    ATOM_LM = 1
    ATOM_LM_PHASE = 2
    NO_RWIGS_ATOM = 10
    NO_RWIGS_ATOM_LM = 11
    NO_RWIGS_ATOM_LM_PHASE = 12
    NO_RWIGS_ATOM_LM_PHASE_AUTO = 14
    ATOM_LM_WAVE = 5

    @classmethod
    def get_lorbit_from_combination(cls, **kwargs) -> 'OrbitEnum':
        """Get the correct mode of the projectors/decomposition."""
        combination = tuple(kwargs[i] for i in ['lm', 'phase', 'wigner_seitz_radius'])
        value_from_combinations = {
            (False, False, True): cls.ATOM,
            (True, False, True): cls.ATOM_LM,
            (True, True, True): cls.ATOM_LM_PHASE,
            (False, False, False): cls.NO_RWIGS_ATOM,
            (True, False, False): cls.NO_RWIGS_ATOM_LM,
            (True, True, False): cls.NO_RWIGS_ATOM_LM_PHASE,
            # Not supported, so also calculate lm decomposed
            (False, True, True): cls.ATOM_LM_PHASE,
            # Not supported, so also calculate lm decomposed
            (False, True, False): cls.NO_RWIGS_ATOM_LM_PHASE,
        }
        return value_from_combinations[combination]


class RelaxAlgoEnum(enum.IntEnum):
    """
    Encode values for algorithm descriptively.

    See: https://www.vasp.at/wiki/index.php/IBRION
    """

    NO_UPDATE = -1
    IONIC_RELAXATION_RMM_DIIS = 1
    IONIC_RELAXATION_CG = 2


class RelaxModeEnum(enum.IntEnum):
    """
    Encode values for degrees of freedom mode of relaxation descriptively.

    See: https://cms.mpi.univie.ac.at/wiki/index.php/ISIF
    """

    POS_ONLY = 2
    POS_SHAPE_VOL = 3
    POS_SHAPE = 4
    SHAPE_ONLY = 5
    SHAPE_VOL = 6
    VOL_ONLY = 7

    @classmethod
    def get_isif_from_dof(cls, **kwargs) -> 'RelaxModeEnum':
        """Get the correct mode of relaxation for the given degrees of freedom."""
        relax_choices = ('positions', 'shape', 'volume')  # pylint: disable=invalid-name
        dof = tuple(kwargs[i] for i in relax_choices)
        value_from_dof = {
            (True, False, False): cls.POS_ONLY,
            (True, True, True): cls.POS_SHAPE_VOL,
            (True, True, False): cls.POS_SHAPE,
            (False, True, False): cls.SHAPE_ONLY,
            (False, True, True): cls.SHAPE_VOL,
            (False, False, True): cls.VOL_ONLY,
        }
        try:
            return value_from_dof[dof]
        except KeyError as no_dof:
            raise ValueError(f'Invalid combination for degrees of freedom: {dict(zip(relax_choices, dof))}') from no_dof


class ParametersMassage:  # pylint: disable=too-many-instance-attributes
    """
    A class that contains all relevant massaging of the input parameters for VASP.

    The idea is that this class accepts the set input parameters from AiiDA (non code specifics), checks if any code
    specific parameters supplied are valid VASP input parameters (only rudimentary at this point, should also cross
    check and check types) and convert the AiiDA input parameters to VASP specific parameters. A set function needs to
    be developed for each VASP INCAR parameter that we want to set based on supplied AiiDA/AiiDA-VASP specific
    parameters. These set functions takes these parameters
    and converts it to VASP INCAR compatible tags. The parameter property should return ready to go parameters
    containing the default override namespace, the namespaces set in the `_set_extra_parameters` function and any
    additional namespaces
    that might have been set using the `additional_override_namespaces` setting entry in settings
    that can be supplied to the VaspWorkChain.

    The default override namespace (see `_DEFAULT_OVERRIDE_NAMESPACE`) should always be present when using this VASP
    plugin.
    If using additional plugins, one can for instance supply additional namespace override that can be used,
    depending on what is needed in those plugins and how you construct your workchains.
    """

    _valid_parameters_and_mtime: tuple[list[str], int] = None

    def __init__(
        self,
        parameters: Any,
        unsupported_parameters: Any = None,
        settings: Any = None,
        skip_parameters_validation: bool = False,
    ) -> None:
        self.exit_code = None

        # Flag for skipping any validations
        self._skip_validation = skip_parameters_validation

        # Check type of parameters and set
        self._parameters = check_inputs(parameters)

        # Check type of supplied unsupported_parameters and set
        self._unsupported_parameters = check_inputs(unsupported_parameters)

        # Check type of settings and set
        self._settings = check_inputs(settings)

        # Check setting for any possible supplied additional override namespace and set
        self._additional_override_namespaces = self._fetch_additional_override_namespaces()

        self._massage = AttributeDict()
        # Initialize any allowed override namespaces
        self._massage[_DEFAULT_OVERRIDE_NAMESPACE] = AttributeDict()
        for item in self._additional_override_namespaces:
            self._massage[item] = AttributeDict()

        self._check_valid_namespaces()
        # Load the valid INCAR parameters that are supported by VASP
        self._load_valid_params()
        # Establish set functions which will convert AiiDA/AiiDA-VASP specific parameters to VASP INCAR tags
        self._functions = ParameterSetFunctions(self._parameters, self._massage[_DEFAULT_OVERRIDE_NAMESPACE])
        # Convert general parameters to VASP specific ones, using the set functions
        self._set_vasp_parameters()
        # Override or set parameters that are supplied in the default override namespace
        # (should be valid VASP INCAR tags)
        self._set_override_vasp_parameters()
        # Set any extra parameters not related to INCAR
        self._set_extra_vasp_parameters()
        # Set any additional override namespace that should just be forwarded
        self._set_additional_override_parameters()
        # Finally, we validate the INCAR parameters in order to prepare it for dispatch
        self._validate_vasp_parameters()

    def _check_valid_namespaces(self) -> None:
        """Check that we do not have namespaces on the input parameters that is unsupported."""
        for key in self._parameters.keys():
            if key not in list(_BASE_NAMESPACES + self._additional_override_namespaces + [_DEFAULT_OVERRIDE_NAMESPACE]):
                raise ValueError(f'The supplied namespace: {key} is not supported.')

    def _load_valid_params(self) -> None:
        """Import a list of valid parameters for VASP. This is generated from the manual."""

        # Cache the valid parameters as class attributes to void repeated loading
        filepath = Path(__file__).parent / 'parameters.yml'
        mtime = Path(filepath).stat().st_mtime

        if (
            ParametersMassage._valid_parameters_and_mtime is None
            or mtime != ParametersMassage._valid_parameters_and_mtime[1]
        ):
            with open(filepath, 'r', encoding='utf8') as handler:
                tags_data = safe_load(handler)
                ParametersMassage._valid_parameters_and_mtime = (list(tags_data.keys()), mtime)

        # Now add any unsupported parameter to the list
        self._valid_parameters = deepcopy(ParametersMassage._valid_parameters_and_mtime[0])
        for key, _ in self._unsupported_parameters.items():
            if key.lower() not in self._valid_parameters:
                self._valid_parameters.append(key.lower())

    def _fetch_additional_override_namespaces(self) -> list[str]:
        """Check the settings for any additional supplied override namespace and return it."""
        try:
            override_namespaces = self._settings.additional_override_namespaces
        except AttributeError:
            override_namespaces = []

        return override_namespaces

    def _set_vasp_parameters(self) -> None:
        """Iterate over the valid parameters and call the set function associated with that parameter."""
        for key in self._valid_parameters:
            self._set(key)

    def _set_override_vasp_parameters(self) -> None:
        """Set the any supplied override parameters."""
        if _DEFAULT_OVERRIDE_NAMESPACE not in self._parameters:
            return

        for key, item in self._parameters[_DEFAULT_OVERRIDE_NAMESPACE].items():
            # Sweep the override input parameters (only care about the ones in the default override namespace)
            # to check if they are valid VASP tags.
            if self._valid_vasp_parameter(key.lower()):
                # Add or override in the default override namespace
                self._massage[_DEFAULT_OVERRIDE_NAMESPACE][key.lower()] = item
            else:
                raise ValueError(f'The supplied key: {key} is not a support VASP parameter.')

    def _set_extra_vasp_parameters(self) -> None:
        """
        Find if there are any extra parameters that are not part of the INCAR that needs to be set.

        One example is the dynamic namespace which handles for instance flags for selective dynamics.
        These flags are more connected to a calculation than a StructureData and thus it was necessary
        to make sure it was valid input to the VASP workchain.

        """
        # return if dynamics is not supplied
        if 'dynamics' not in self._parameters:
            return

        if self._parameters.dynamics:
            self._massage.dynamics = AttributeDict()

        for key_, item in self._parameters.dynamics.items():
            key = key_.lower()
            if key in ['positions_dof']:
                self._massage.dynamics[key] = item
            else:
                warn(f"Key {key} is not supported for 'dynamics' input.")

    def _set_additional_override_parameters(self) -> None:
        """Set any customized parameter namespace, including its content on the massaged container."""
        parameters_keys = self._parameters.keys()
        for item in self._additional_override_namespaces:
            if item in parameters_keys:
                # Only add if namespace exists in parameters
                self._massage[item] = AttributeDict(self._parameters[item])

    def _valid_vasp_parameter(self, key: str) -> bool:
        """Make sure a key are recognized as a valid VASP input parameter."""
        if self._skip_validation or (key in self._valid_parameters):
            return True

        return False

    def _validate_vasp_parameters(self) -> None:
        """Make sure all the massaged values are recognized as valid VASP input parameters."""

        for key in self._massage[_DEFAULT_OVERRIDE_NAMESPACE]:
            if not self._valid_vasp_parameter(key.lower()):
                raise ValueError(f'The supplied key: {key} is not a support VASP parameter.')

    def _set(self, key: str) -> None:
        """Call the necessary function to set each parameter."""
        try:
            getattr(self._functions, 'set_' + key)()
        except AttributeError:
            # We have no setter function for the valid key, meaning there is no general parameter that is linked
            # to this key (INCAR tag). These tags have to be supplied in the default override namespace for now.
            pass

    @property
    def parameters(self) -> AttributeDict:
        """Return the massaged parameter set ready to go in VASP format."""
        return self._massage


class ParameterSetFunctions:
    """Container for the set functions that converts an AiiDA parameters to a default override specific one."""

    def __init__(self, parameters: AttributeDict, incar: AttributeDict) -> None:
        self._parameters = parameters
        self._incar = incar

    def set_encut(self) -> None:
        """
        Set which plane wave cutoff to use.

        See https://www.vasp.at/wiki/index.php/ENCUT
        """

        try:
            self._incar.encut = self._parameters.electronic.pwcutoff
        except AttributeError:
            pass

    def set_ibrion(self) -> None:
        """
        Set which algorithm to use for ionic movements.

        See: https://www.vasp.at/wiki/index.php/IBRION
        """

        if self._relax():
            try:
                if self._parameters.relax.algo == 'cg':
                    self._incar.ibrion = RelaxAlgoEnum.IONIC_RELAXATION_CG.value
                elif self._parameters.relax.algo == 'rd':
                    self._incar.ibrion = RelaxAlgoEnum.IONIC_RELAXATION_RMM_DIIS.value
                else:
                    raise ValueError(f'The supplied relax.algo: {self._parameters.relax.algo} is not supported')
            except AttributeError:
                pass

    def set_ediffg(self) -> None:
        """
        Set the cutoff to use for relaxation.

        See: https://www.vasp.at/wiki/index.php/EDIFFG
        """

        if not self._relax():
            # This flag is only valid if you have enabled relaxation
            return
        energy_cutoff = False
        try:
            if self._parameters.relax.energy_cutoff is not None:
                self._incar.ediffg = self._parameters.relax.energy_cutoff
                energy_cutoff = True
        except AttributeError:
            pass
        try:
            self._incar.ediffg = -abs(self._parameters.relax.force_cutoff)
            if energy_cutoff:
                raise ValueError('User supplied both a force and an energy cutoff for the relaxation. Please select.')
        except AttributeError:
            pass

    def set_nsw(self) -> None:
        """
        Set the number of ionic steps to perform.

        See: https://www.vasp.at/wiki/index.php/NSW
        """

        if self._relax():
            try:
                self._set_simple('nsw', self._parameters.relax.steps)
            except AttributeError:
                pass

    def set_isif(self) -> None:
        """
        Set relaxation mode according to the chosen degrees of freedom.

        See: https://www.vasp.at/wiki/index.php/ISIF
        """

        if self._relax():
            positions = self._parameters.get('relax', {}).get('positions', False)
            shape = self._parameters.get('relax', {}).get('shape', False)
            volume = self._parameters.get('relax', {}).get('volume', False)
            try:
                self._incar.isif = RelaxModeEnum.get_isif_from_dof(
                    positions=positions, shape=shape, volume=volume
                ).value
            except AttributeError:
                pass

    def set_ismear(self) -> None:
        """
        Make sure we do not supply invalid integration methods when running explicit k-point grids.

        See: https://www.vasp.at/wiki/index.php/ISMEAR
        """

        try:
            if self._parameters.smearing.gaussian:
                self._set_simple('ismear', IntSmearingEnum.GAUSSIAN.value)
        except AttributeError:
            pass
        try:
            if self._parameters.smearing.fermi:
                self._set_simple('ismear', IntSmearingEnum.FERMI.value)
        except AttributeError:
            pass
        try:
            if self._parameters.smearing.mp:
                self._set_simple('ismear', IntSmearingEnum.MP.value * abs(int(self._parameters.smearing.mp)))
        except AttributeError:
            pass
        try:
            if self._parameters.smearing.tetra:
                self._set_simple('ismear', IntSmearingEnum.TETRA.value)
        except AttributeError:
            pass

    def set_icharg(self) -> None:
        """
        Set the flag to start from input charge density and keep it constant.

        See: https://www.vasp.at/wiki/index.php/ICHARG
        """

        try:
            if self._parameters.charge.from_wave:
                self._set_simple('icharg', ChargeEnum.WAVE.value)
        except AttributeError:
            pass
        try:
            if self._parameters.charge.from_charge:
                self._set_simple('icharg', ChargeEnum.CHARGE.value)
        except AttributeError:
            pass
        try:
            if self._parameters.charge.from_atomic:
                self._set_simple('icharg', ChargeEnum.ATOMIC.value)
        except AttributeError:
            pass
        try:
            if self._parameters.charge.from_potential:
                self._set_simple('icharg', ChargeEnum.POTENTIAL.value)
        except AttributeError:
            pass
        try:
            if self._parameters.charge.constant_charge:
                self._set_simple('icharg', ChargeEnum.CONSTANT_CHARGE.value)
        except AttributeError:
            pass
        try:
            if self._parameters.charge.constant_atomic:
                self._set_simple('icharg', ChargeEnum.CONSTANT_ATOMIC.value)
        except AttributeError:
            pass

    def set_lorbit(self) -> None:
        """
        Set the flag that controls the projectors/decomposition onto orbitals.

        See: https://www.vasp.at/wiki/index.php/LORBIT
        """
        self._set_wigner_seitz_radius()
        try:
            if self._parameters.bands.decompose_bands:
                if self._parameters.bands.decompose_wave:
                    # Issue a warning that one can only use either or
                    raise ValueError('Only projections/decompositions on the bands or the wave function are allowed.')
                wigner_seitz_radius = False
                try:
                    if abs(self._incar.rwigs[0]) > 1e-8:
                        wigner_seitz_radius = True
                except AttributeError:
                    pass
                if self._parameters.bands.decompose_auto:
                    self._set_simple('lorbit', OrbitEnum.NO_RWIGS_ATOM_LM_PHASE_AUTO.value)
                else:
                    try:
                        lm = self._parameters.bands.lm  # pylint: disable=invalid-name
                    except AttributeError:
                        lm = False  # pylint: disable=invalid-name
                    try:
                        phase = self._parameters.bands.phase
                    except AttributeError:
                        phase = False
                    lorbit = OrbitEnum.get_lorbit_from_combination(
                        lm=lm, phase=phase, wigner_seitz_radius=wigner_seitz_radius
                    ).value
                    self._set_simple('lorbit', lorbit)
            else:
                try:
                    if self._parameters.bands.decompose_wave:
                        self._set_simple('lorbit', OrbitEnum.ATOM_LM_WAVE.value)
                except AttributeError:
                    pass
        except AttributeError:
            try:
                if self._parameters.bands.decompose_wave:
                    self._set_simple('lorbit', OrbitEnum.ATOM_LM_WAVE.value)
            except AttributeError:
                pass

    def _set_wigner_seitz_radius(self) -> None:
        """
        Set the Wigner Seitz radius that is used to project/decompose.

        See: https://www.vasp.at/wiki/index.php/RWIGS
        """
        try:
            wigner_seitz_radius = self._parameters.bands.wigner_seitz_radius
            # Check that it is defined as a list
            if isinstance(wigner_seitz_radius, list):
                if wigner_seitz_radius[0]:
                    self._set_simple('rwigs', wigner_seitz_radius)
            else:
                raise ValueError(
                    'The parameter wigner_seitz_radius should be supplied as a list of floats bigger than zero.'
                )
        except AttributeError:
            pass

    def _relax(self) -> bool:
        """Check if we have enabled relaxation."""
        return bool(
            self._parameters.get('relax', {}).get('positions')
            or self._parameters.get('relax', {}).get('shape')
            or self._parameters.get('relax', {}).get('volume')
        )

    def _set_simple(self, target: str, value: Any) -> None:
        """Set basic parameter."""
        try:
            self._incar[target] = value
        except AttributeError:
            pass


def check_inputs(supplied_inputs: None | AttributeDict | orm.Dict | dict) -> AttributeDict:
    """Check that the inputs are of some correct type and returned as AttributeDict."""
    inputs = None
    if supplied_inputs is None:
        inputs = AttributeDict()
    elif isinstance(supplied_inputs, orm.Dict):
        inputs = AttributeDict(supplied_inputs.get_dict())
    elif isinstance(supplied_inputs, dict):
        inputs = AttributeDict(supplied_inputs)
    elif isinstance(supplied_inputs, AttributeDict):
        inputs = supplied_inputs
    else:
        raise ValueError(
            f'The supplied type {type(inputs)} of inputs is not supported. Supply a dict, Dict or an AttributeDict.'
        )

    return inputs


def inherit_and_merge_parameters(inputs: dict[str, Any]) -> AttributeDict:
    """
    Goes trough the inputs namespaces and the namespaces in the inputs.parameters and merge them.

    Note that parameters specified in the inputs.parameters will override what is supplied as workchain input,
    in case there is overlap.
    """
    parameters = AttributeDict()
    namespaces = _BASE_NAMESPACES

    # We start with a clean parameters and first set the allowed namespaces and its content from the inputs of the
    # workchain
    for namespace in namespaces:  # pylint: disable=too-many-nested-blocks
        parameters[namespace] = AttributeDict()
        try:
            for key, item in inputs[namespace].items():
                if isinstance(item, orm.ArrayData):
                    # Only allow one array per input
                    if len(item.get_arraynames()) > 1:
                        raise IndexError(
                            f'The input array with a key {key} contains more than one array. '
                            'Please make sure an input only contains one array.'
                        )
                    for array in item.get_arraynames():
                        parameters[namespace][key] = item.get_array(array)
                elif isinstance(item, orm.Dict):
                    parameters[namespace][key] = item.get_dict()
                elif isinstance(item, orm.List):
                    parameters[namespace][key] = item.get_list()
                else:
                    parameters[namespace][key] = item.value
        except KeyError:
            pass

    # Then obtain the inputs.parameters.
    # Here we do not do any checks for valid parameters, that is done later when reaching the ParameterMassager.
    try:
        input_parameters = AttributeDict(inputs.parameters.get_dict())
    except AttributeError:
        # Inputs might not have parameters
        input_parameters = AttributeDict()

    # Now the namespace and content of the workchain inputs and the inputs.parameters are merged.
    # Any supplied namespace in the parameters (i.e. inputs.parameters.somekey) will override what
    # is supplied to the workchain input namespace (i.e. inputs.somekey).
    # We cannot use regular update here, as we only want to replace each key if it exists, if a key
    # contains a new dict we need to traverse that, hence we have a function to perform this update.
    update_nested_dict(parameters, input_parameters)

    return parameters


# class ParserSettingsChecker:
#     """
#     Check for inconsistency between the parser settings and the INCAR tags.
#     """

#     def __init__(self, parameters: dict, settings: dict, default_settings=None):
#         """
#         Instantiate an ParserSettingsChecker object.

#         :param parameters: A dictionary of INCAR tags to be used
#         :param settings: The settings dictionary
#         :param default_option: The default options to be used for instantiating the ParserSetting object.

#         """

#         if default_settings is None:
#             default_settings = DEFAULT_SETTINGS

#         # Obtain the expected quantities names to be parsed
#         self.settings = ParserSettings(settings, default_settings)
#         self.quantities = settings.quantity_names_to_parse
#         self.parameters = parameters

#     def check(self):
#         """
#         Check the consistency between the supplied INCAR tags the requested quantities to be parsed.
#         """

#         self.check_maximum_stress()
#         self.check_wavecar_chgcar()
#         self.check_born_charges()
#         self.check_dynmat()

#     def check_maximum_stress(self):
#         """Check the the maximum_stress can be parsed"""

#         if 'maximum_stress' not in self.quantities and 'stress' not in self.quantities:
#             return

#         isif = self.parameters.get('isif')
#         ibrion = self.parameters.get('ibrion', -1)
#         lhfcalc = self.parameters.get('lhfcalc', False)
#         if isif is None:
#             if ibrion == 0:
#                 isif = 0
#             if lhfcalc:
#                 isif = 0

#         if isif == 0:
#             raise InputValidationError(
#                 'Requested to parse <maximum_stress> but it would not be calculated due to ISIF settings.'
#             )

#     def check_wavecar_chgcar(self):
#         """Check if WAVECAR CHGCAR are set to be written"""

#         if 'wavecar' in self.quantities and not self.parameters.get('lwave', True):
#             raise InputValidationError('Requested to retrieve <WAVECAR> but it not set to be written.')

#         if 'chgcar' in self.quantities and not self.parameters.get('lwave', True):
#             raise InputValidationError('Requested to retrieve <CHGCAR> but it not set to be written.')

#     def check_born_charges(self):
#         """Check if projectors can be parsed"""
#         if 'born_charges' not in self.quantities and 'dielectrics' not in self.quantities:
#             return

#         lepsilon = self.parameters.get('lepsilon', False)
#         if not lepsilon:
#             raise InputValidationError(
#                 'Requested to parse "born_charges"/"dielectrics" but they are not going to be calculated.'
#             )

#     def check_dynmat(self):
#         if 'hessian' not in self.quantities and 'dynmat' not in self.quantities:
#             return
#         ibrion = self.parameters.get('ibrion', -1)
#         if ibrion not in [5, 6, 7, 8]:
#             raise InputValidationError('Requsted to parse "hessian"/"dynmat" '
#                                      'but they are not going to be calculated.')
