# -*- coding: utf-8 -*-
"""Workchain to compute a band structure for a given structure using Quantum ESPRESSO pw.x.
This is a modified version of the original `PwBandsWorkChain` in aiida-quantumespresso in order to
support caching of the results wisely. The PW calculation will be re-run if the the parent folder is cleaned and
but required by the following PH calculation.
"""
from aiida import orm
from aiida.common import AttributeDict
from aiida.engine import ToContext, WorkChain, if_, while_
from aiida.plugins import WorkflowFactory
from aiida_quantumespresso.calculations.functions.seekpath_structure_analysis import (
    seekpath_structure_analysis,
)
from aiida_quantumespresso.utils.mapping import prepare_process_inputs
from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin

PwBaseWorkChain = WorkflowFactory("quantumespresso.pw.base")
PwRelaxWorkChain = WorkflowFactory("quantumespresso.pw.relax")


def validate_inputs(inputs, ctx=None):  # pylint: disable=unused-argument
    """Validate the inputs of the entire input namespace."""
    # pylint: disable=no-member

    if "nbands_factor" in inputs and "nbnd" in inputs["bands"]["pw"][
        "parameters"
    ].base.attributes.get("SYSTEM", {}):
        return PwBandsWorkChain.exit_codes.ERROR_INVALID_INPUT_NUMBER_OF_BANDS.message

    # Cannot specify both `bands_kpoints` and `bands_kpoints_distance`
    if all(key in inputs for key in ["bands_kpoints", "bands_kpoints_distance"]):
        return PwBandsWorkChain.exit_codes.ERROR_INVALID_INPUT_KPOINTS.message


class PwBandsWorkChain(ProtocolMixin, WorkChain):
    """Workchain to compute a band structure for a given structure using Quantum ESPRESSO pw.x.

    The logic for the computation of various parameters for the BANDS step is as follows:

    Number of bands:
        One can specify the number of bands to be used in the BANDS step either directly through the input parameters
        `bands.pw.parameters.SYSTEM.nbnd` or through `nbands_factor`. Note that specifying both is not allowed. When
        neither is specified nothing will be set by the work chain and the default of Quantum ESPRESSO will end up being
        used. If the `nbands_factor` is specified the maximum value of the following values will be used:

        * `nbnd` of the preceding SCF calculation
        * 0.5 * nelectrons * nbands_factor
        * 0.5 * nelectrons + 4

    Kpoints:
        There are three options; specify either an existing `KpointsData` through `bands_kpoints`, or specify the
        `bands_kpoint_distance`, or specify neither. For the former those exact kpoints will be used for the BANDS step.
        In the two other cases, the structure will first be normalized using SeekPath and the path along high-symmetry
        k-points will be generated on that structure. The distance between kpoints for the path will be equal to that
        of `bands_kpoints_distance` or the SeekPath default if not specified.
    """

    @classmethod
    def define(cls, spec):
        """Define the process specification."""
        # yapf: disable
        super().define(spec)
        spec.expose_inputs(PwRelaxWorkChain, namespace='relax', exclude=('clean_workdir', 'structure'),
            namespace_options={'required': False, 'populate_defaults': False,
            'help': 'Inputs for the `PwRelaxWorkChain`, if not specified at all, the relaxation step is skipped.'})
        spec.expose_inputs(PwBaseWorkChain, namespace='scf',
            exclude=('clean_workdir', 'pw.structure'),
            namespace_options={'help': 'Inputs for the `PwBaseWorkChain` for the SCF calculation.'})
        spec.expose_inputs(PwBaseWorkChain, namespace='bands',
            exclude=('clean_workdir', 'pw.structure', 'pw.kpoints', 'pw.kpoints_distance', 'pw.parent_folder'),
            namespace_options={'help': 'Inputs for the `PwBaseWorkChain` for the BANDS calculation.'})
        spec.input('structure', valid_type=orm.StructureData, help='The inputs structure.')
        spec.input('clean_workdir', valid_type=orm.Bool, default=lambda: orm.Bool(False),
            help='If `True`, work directories of all called calculation will be cleaned at the end of execution.')
        spec.input('nbands_factor', valid_type=orm.Float, required=False,
            help='The number of bands for the BANDS calculation is that used for the SCF multiplied by this factor.')
        spec.input('bands_kpoints', valid_type=orm.KpointsData, required=False,
            help='Explicit kpoints to use for the BANDS calculation. Specify either this or `bands_kpoints_distance`.')
        spec.input('bands_kpoints_distance', valid_type=orm.Float, required=False,
            help='Minimum kpoints distance for the BANDS calculation. Specify either this or `bands_kpoints`.')
        spec.inputs.validator = validate_inputs
        spec.outline(
            cls.setup,
            if_(cls.should_run_relax)(
                cls.run_relax,
                cls.inspect_relax,
            ),
            if_(cls.should_run_seekpath)(
                cls.run_seekpath,
            ),
            while_(cls.should_run_band)(
                cls.run_scf,
                cls.inspect_scf,
                cls.run_bands,
                cls.inspect_bands,
            ),
            cls.results,
        )
        spec.exit_code(201, 'ERROR_INVALID_INPUT_NUMBER_OF_BANDS',
            message='Cannot specify both `nbands_factor` and `bands.pw.parameters.SYSTEM.nbnd`.')
        spec.exit_code(202, 'ERROR_INVALID_INPUT_KPOINTS',
            message='Cannot specify both `bands_kpoints` and `bands_kpoints_distance`.')
        spec.exit_code(401, 'ERROR_SUB_PROCESS_FAILED_RELAX',
            message='The PwRelaxWorkChain sub process failed')
        spec.exit_code(402, 'ERROR_SUB_PROCESS_FAILED_SCF',
            message='The scf PwBasexWorkChain sub process failed')
        spec.exit_code(403, 'ERROR_SUB_PROCESS_FAILED_BANDS',
            message='The bands PwBasexWorkChain sub process failed')
        spec.output('primitive_structure', valid_type=orm.StructureData,
            required=False,
            help='The normalized and primitivized structure for which the bands are computed.')
        spec.output('seekpath_parameters', valid_type=orm.Dict,
            required=False,
            help='The parameters used in the SeeKpath call to normalize the input or relaxed structure.')
        spec.output('scf_parameters', valid_type=orm.Dict,
            help='The output parameters of the SCF `PwBaseWorkChain`.')
        spec.output('band_parameters', valid_type=orm.Dict,
            help='The output parameters of the BANDS `PwBaseWorkChain`.')
        spec.output('band_structure', valid_type=orm.BandsData,
            help='The computed band structure.')
        # yapf: enable

    @classmethod
    def get_protocol_filepath(cls):
        """Return ``pathlib.Path`` to the ``.yaml`` file that defines the protocols."""
        from importlib_resources import files

        from ..protocols import pw as pw_protocols

        return files(pw_protocols) / "bands.yaml"

    @classmethod
    def get_builder_from_protocol(
        cls, code, structure, protocol=None, overrides=None, options=None, **kwargs
    ):
        """Return a builder prepopulated with inputs selected according to the chosen protocol.

        :param code: the ``Code`` instance configured for the ``quantumespresso.pw`` plugin.
        :param structure: the ``StructureData`` instance to use.
        :param protocol: protocol to use, if not specified, the default will be used.
        :param overrides: optional dictionary of inputs to override the defaults of the protocol.
        :param options: A dictionary of options that will be recursively set for the ``metadata.options`` input of all
            the ``CalcJobs`` that are nested in this work chain.
        :param kwargs: additional keyword arguments that will be passed to the ``get_builder_from_protocol`` of all the
            sub processes that are called by this workchain.
        :return: a process builder instance with all inputs defined ready for launch.
        """
        inputs = cls.get_protocol_inputs(protocol, overrides)

        args = (code, structure, protocol)
        relax = PwRelaxWorkChain.get_builder_from_protocol(
            *args, overrides=inputs.get("relax", None), options=options, **kwargs
        )
        scf = PwBaseWorkChain.get_builder_from_protocol(
            *args, overrides=inputs.get("scf", None), options=options, **kwargs
        )
        bands = PwBaseWorkChain.get_builder_from_protocol(
            *args, overrides=inputs.get("bands", None), options=options, **kwargs
        )

        relax.pop("structure", None)
        relax.pop("clean_workdir", None)
        relax.pop("base_final_scf", None)
        scf["pw"].pop("structure", None)
        scf.pop("clean_workdir", None)
        bands["pw"].pop("structure", None)
        bands.pop("clean_workdir", None)
        bands.pop("kpoints_distance", None)
        bands.pop("kpoints_force_parity", None)

        builder = cls.get_builder()
        builder.structure = structure
        builder.relax = relax
        builder.scf = scf
        builder.bands = bands
        builder.clean_workdir = orm.Bool(inputs["clean_workdir"])
        builder.nbands_factor = orm.Float(inputs["nbands_factor"])
        builder.bands_kpoints_distance = orm.Float(inputs["bands_kpoints_distance"])

        return builder

    def setup(self):
        """Define the current structure in the context to be the input structure."""
        self.ctx.current_structure = self.inputs.structure
        self.ctx.current_number_of_bands = None
        self.ctx.bands_kpoints = self.inputs.get("bands_kpoints", None)

        # Set to run for the first loop
        self.ctx.should_run_band = True

        # The list of nodes that are invalided the caching of the node and re-run scf calculation. When the remote folder of parent calculation is empty, the caching is invalid, but since invalid caching is applied to all_same_nodes, we need to keep track of the invalid nodes and bring caching valid back to the node so the next run will still using the caching, otherwise the next parent calculation will be re-run again, which is not what we want.
        # Warning!! experimental feature and will potentially cause "race condition" like behavior between workflows.
        self.ctx.cache_invalid_list = []

    def should_run_band(self):
        """The value will be set to `True` if bands calculation is failed because of the remote folder issue.
        Otherwise it keeps on trying to rerun the bands calculation until remote folder of preliminary scf is not found.
        """
        return self.ctx.should_run_band

    def should_run_relax(self):
        """If the 'relax' input namespace was specified, we relax the input structure."""
        return "relax" in self.inputs

    def should_run_seekpath(self):
        """Seekpath should only be run if the `bands_kpoints` input is not specified."""
        return "bands_kpoints" not in self.inputs

    def run_relax(self):
        """Run the PwRelaxWorkChain to run a relax PwCalculation."""
        inputs = AttributeDict(self.exposed_inputs(PwRelaxWorkChain, namespace="relax"))
        inputs.metadata.call_link_label = "relax"
        inputs.structure = self.ctx.current_structure

        running = self.submit(PwRelaxWorkChain, **inputs)

        self.report(f"launching PwRelaxWorkChain<{running.pk}>")

        return ToContext(workchain_relax=running)

    def inspect_relax(self):
        """Verify that the PwRelaxWorkChain finished successfully."""
        workchain = self.ctx.workchain_relax

        if not workchain.is_finished_ok:
            self.report(
                f"PwRelaxWorkChain failed with exit status {workchain.exit_status}"
            )
            return self.exit_codes.ERROR_SUB_PROCESS_FAILED_RELAX

        self.ctx.current_structure = workchain.outputs.output_structure
        self.ctx.current_number_of_bands = (
            workchain.outputs.output_parameters.base.attributes.get("number_of_bands")
        )

    def run_seekpath(self):
        """Run the structure through SeeKpath to get the normalized structure and path along high-symmetry k-points .

        This is only called if the `bands_kpoints` input was not specified.
        """
        inputs = {
            "reference_distance": self.inputs.get("bands_kpoints_distance", None),
            "metadata": {"call_link_label": "seekpath"},
        }
        result = seekpath_structure_analysis(self.ctx.current_structure, **inputs)
        self.ctx.current_structure = result["primitive_structure"]
        self.ctx.bands_kpoints = result["explicit_kpoints"]

        self.out("primitive_structure", result["primitive_structure"])
        self.out("seekpath_parameters", result["parameters"])

    def run_scf(self):
        """Run the PwBaseWorkChain in scf mode on the primitive cell of (optionally relaxed) input structure."""
        inputs = AttributeDict(self.exposed_inputs(PwBaseWorkChain, namespace="scf"))
        inputs.metadata.call_link_label = "scf"
        inputs.pw.structure = self.ctx.current_structure
        inputs.pw.parameters = inputs.pw.parameters.get_dict()
        inputs.pw.parameters.setdefault("CONTROL", {})["calculation"] = "scf"

        # Make sure to carry the number of bands from the relax workchain if it was run and it wasn't explicitly defined
        # in the inputs. One of the base workchains in the relax workchain may have changed the number automatically in
        #  the sanity checks on band occupations.
        if self.ctx.current_number_of_bands:
            inputs.pw.parameters.setdefault("SYSTEM", {}).setdefault(
                "nbnd", self.ctx.current_number_of_bands
            )

        inputs = prepare_process_inputs(PwBaseWorkChain, inputs)
        running = self.submit(PwBaseWorkChain, **inputs)

        self.report(f"launching PwBaseWorkChain<{running.pk}> in scf mode")

        return ToContext(workchain_scf=running)

    def inspect_scf(self):
        """Verify that the PwBaseWorkChain for the scf run finished successfully."""
        workchain = self.ctx.workchain_scf

        # duplicate of _disable_cache of _BaseEvaluateWorkChain
        # I do not want SCF calc used for caching if it is a from cached calculation
        # right after!@@!
        if workchain.is_finished:
            for child in workchain.called_descendants:
                if (
                    child.process_label == "PwCalculation"
                    and child.base.caching.is_created_from_cache
                ):
                    child.is_valid_cache = False

        if not workchain.is_finished_ok:
            self.report(
                f"scf PwBaseWorkChain failed with exit status {workchain.exit_status}"
            )
            return self.exit_codes.ERROR_SUB_PROCESS_FAILED_SCF

        self.ctx.current_folder = workchain.outputs.remote_folder
        # the current_folder is set so can bring the caching invalid back to the node
        while self.ctx.cache_invalid_list:
            node_pk = self.ctx.cache_invalid_list.pop()
            node = orm.load_node(node_pk)
            node.is_valid_cache = True

        self.ctx.current_number_of_bands = (
            workchain.outputs.output_parameters.base.attributes.get("number_of_bands")
        )

    def run_bands(self):
        """Run the PwBaseWorkChain in bands mode along the path of high-symmetry determined by seekpath."""
        inputs = AttributeDict(self.exposed_inputs(PwBaseWorkChain, namespace="bands"))
        inputs.metadata.call_link_label = "bands"
        inputs.kpoints = self.ctx.bands_kpoints
        inputs.pw.structure = self.ctx.current_structure
        inputs.pw.parent_folder = self.ctx.current_folder
        inputs.pw.parameters = inputs.pw.parameters.get_dict()
        inputs.pw.parameters.setdefault("CONTROL", {})
        inputs.pw.parameters.setdefault("SYSTEM", {})
        inputs.pw.parameters.setdefault("ELECTRONS", {})

        # The following flags always have to be set in the parameters, regardless of what caller specified in the inputs
        inputs.pw.parameters["CONTROL"]["calculation"] = "bands"

        # Only set the following parameters if not directly explicitly defined in the inputs
        inputs.pw.parameters["ELECTRONS"].setdefault("diagonalization", "cg")
        inputs.pw.parameters["ELECTRONS"].setdefault("diago_full_acc", True)

        # If `nbands_factor` is defined in the inputs we set the `nbnd` parameter
        if "nbands_factor" in self.inputs:
            factor = self.inputs.nbands_factor.value
            parameters = self.ctx.workchain_scf.outputs.output_parameters.get_dict()
            nbands = int(parameters["number_of_bands"])
            nelectron = int(parameters["number_of_electrons"])
            nbnd = max(int(0.5 * nelectron * factor), int(0.5 * nelectron) + 4, nbands)
            inputs.pw.parameters["SYSTEM"]["nbnd"] = nbnd

        # Otherwise set the current number of bands, unless explicitly set in the inputs
        else:
            inputs.pw.parameters["SYSTEM"].setdefault(
                "nbnd", self.ctx.current_number_of_bands
            )

        inputs = prepare_process_inputs(PwBaseWorkChain, inputs)
        running = self.submit(PwBaseWorkChain, **inputs)

        self.report(f"launching PwBaseWorkChain<{running.pk}> in bands mode")

        return ToContext(workchain_bands=running)

    def inspect_bands(self):
        """Verify that the PwBaseWorkChain for the bands run finished successfully."""
        workchain = self.ctx.workchain_bands

        if not workchain.is_finished_ok:
            # The inner logic to detect if the failed bands calculation is due to the empty remote folder
            # If so, we will invalid the caching of the node and re-run scf calculation.
            # !!! This empty check is not a good way to detect the remote folder issue, but it is the only way I can think of now.
            # !!! If the SSH connection is lost, the exception will be raised and the workchain will be terminated.
            # Therefore use cleaned flag to detect the remote folder issue this didn't require SSH connection. But will not work if the remote folder is cleaned by other means.
            is_cleaned = self.ctx.current_folder.base.extras.get("cleaned", False)
            if is_cleaned:
                self.logger.warning(
                    f"PhBaseWorkChain failed because the remote folder is empty with exit status {workchain.exit_status}, invalid the caching of the node and re-run scf calculation."
                )
                # invalid the caching of the node and re-run scf calculation
                workchain_scf = self.ctx.workchain_scf
                scf_pw_node = [
                    c for c in workchain_scf.called if isinstance(c, orm.CalcJobNode)
                ][0]
                all_same_nodes = scf_pw_node.base.caching.get_all_same_nodes()
                for node in all_same_nodes:
                    node.is_valid_cache = False
                    self.ctx.cache_invalid_list.append(node.pk)
                return
            else:
                self.report(
                    f"bands PwBaseWorkChain failed with exit status {workchain.exit_status}"
                )
                return self.exit_codes.ERROR_SUB_PROCESS_FAILED_BANDS

        # The bands workchain finished successfully, so we can set the `should_run_band` flag to False
        self.ctx.should_run_band = False

    def results(self):
        """Attach the desired output nodes directly as outputs of the workchain."""
        self.report("workchain succesfully completed")
        self.out("scf_parameters", self.ctx.workchain_scf.outputs.output_parameters)
        self.out("band_parameters", self.ctx.workchain_bands.outputs.output_parameters)
        self.out("band_structure", self.ctx.workchain_bands.outputs.output_band)

    def on_terminated(self):
        """Clean the working directories of all child calculations if `clean_workdir=True` in the inputs."""
        super().on_terminated()

        if self.inputs.clean_workdir.value is False:
            self.report("remote folders will not be cleaned")
            return

        cleaned_calcs = []

        for called_descendant in self.node.called_descendants:
            if isinstance(called_descendant, orm.CalcJobNode):
                try:
                    called_descendant.outputs.remote_folder._clean()  # pylint: disable=protected-access
                    cleaned_calcs.append(called_descendant.pk)
                except (IOError, OSError, KeyError):
                    pass

        if cleaned_calcs:
            self.report(
                f"cleaned remote folders of calculations: {' '.join(map(str, cleaned_calcs))}"
            )
