# -*- coding: utf-8 -*-
"""`CalcJob` implementation for DFTK."""
import io
import os
import json
import typing as ty

from aiida import orm
from aiida.common import datastructures, exceptions
from aiida.engine import CalcJob, ExitCode
from aiida_pseudo.data.pseudo import UpfData
from pymatgen.core import units


_AIIDA_DFTK_VERSION_SPEC = "0.2.0"


class DftkCalculation(CalcJob):
    """`CalcJob` implementation for DFTK."""

    INPUT_FILENAME = 'run_dftk.json'
    LOGFILE = 'run_dftk.log'
    SCFRES_SUMMARY_NAME = 'self_consistent_field.json'
    # TODO: don't limit postscf
    _SUPPORTED_POSTSCF = ['compute_forces_cart', 'compute_stresses_cart', 'compute_bands']
    _PSEUDO_SUBFOLDER = './pseudo/'
    _MIN_OUTPUT_BUFFER_TIME = 60

    @staticmethod
    def _merge_dicts(dict1, dict2):
        """Recursively merge dict2 into dict1."""
        for key, value in dict2.items():
            if key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict):
                DftkCalculation._merge_dicts(dict1[key], value)
            else:
                dict1[key] = value

    @classmethod
    def define(cls, spec):
        """Define the process specification."""
        super().define(spec)
        # Inputs
        spec.input('structure', valid_type=orm.StructureData, help='structure')
        spec.input_namespace('pseudos', valid_type=UpfData, help='The pseudopotentials.', dynamic=True)
        spec.input('pseudo_rcut', valid_type=orm.Float, required=False, help='rcut to use for the pseudopotentials, in atomic units')
        spec.input('kpoints', valid_type=orm.KpointsData, help='kpoint mesh or kpoint path')
        spec.input('parameters', valid_type=orm.Dict, help='input parameters')
        spec.input('parent_folder', valid_type=orm.RemoteData, required=False, help='A remote folder used for restarts.')

        options = spec.inputs['metadata']['options']

        options['parser_name'].default = 'dftk'
        options['input_filename'].default = cls.INPUT_FILENAME
        options['max_wallclock_seconds'].default = 1800

        # TODO: Why is this here?
        options['resources'].default = {'num_machines': 1, 'num_mpiprocs_per_machine': 1}
        options['withmpi'].default = True

        # Exit codes
        # TODO: Codes 1xx are already used in the super class!
        spec.exit_code(101, 'ERROR_MISSING_SCFRES_FILE', message='The output file containing SCF results is missing.')
        spec.exit_code(102, 'ERROR_MISSING_FORCES_FILE', message='The output file containing forces is missing.')
        spec.exit_code(103, 'ERROR_MISSING_STRESSES_FILE', message='The output file containing stresses is missing.')
        spec.exit_code(104, 'ERROR_MISSING_BANDS_FILE',message='The output file containing bands is missing.')
        spec.exit_code(500, 'ERROR_SCF_CONVERGENCE_NOT_REACHED', message='The SCF minimization cycle did not converge, and the POSTSCF functions were not executed.')
        spec.exit_code(501, 'ERROR_SCF_OUT_OF_WALLTIME',message='The SCF was interuptted due to out of walltime. Non-recovarable error.')
        spec.exit_code(502, 'ERROR_POSTSCF_OUT_OF_WALLTIME',message='The POSTSCF was interuptted due to out of walltime.')
        spec.exit_code(503, 'ERROR_BANDS_CONVERGENCE_NOT_REACHED', message='The BANDS minimization cycle did not converge.')
        # Significant errors but calculation can be used to restart
        spec.exit_code(400, 'ERROR_PACKAGE_IMPORT_FAILED', message="Failed to import AiiDA DFTK or write first log message. Typically indicates an environment issue.")

        # Outputs
        spec.output('output_parameters', valid_type=orm.Dict, help='output parameters')
        # TODO: doesn't seem to be used?
        spec.output('output_structure', valid_type=orm.Dict, required=False, help='output structure')
        # TODO: doesn't seem to be used?
        spec.output(
            'output_kpoints', valid_type=orm.KpointsData, required=False, help='kpoints array, if generated by DFTK'
        )
        spec.output('output_forces', valid_type=orm.ArrayData, required=False, help='forces array')
        spec.output('output_stresses', valid_type=orm.ArrayData, required=False, help='stresses array')
        spec.output('output_bands', valid_type=orm.BandsData, required=False, help='bandstructure')

        # TODO: bands and DOS implementation required on DFTK side
        # spec.output('output_bands', valid_type=orm.BandsData, required=False,
        #     help='eigenvalues array')
        # spec.output('output_dos')

        spec.default_output_node = 'output_parameters'

    def _validate_options(self):
        """Validate the options input.
        
        Check that the wihmpi option is set to True if the number of mpiprocs is greater than 1.
        Check max_wallclock_seconds is greater than the min_output_buffer_time.
        """
        options = self.inputs.metadata.options
        if options.withmpi is False and options.resources.get('num_mpiprocs_per_machine', 1) > 1:
            # TODO: does aiida not already check this?
            raise exceptions.InputValidationError('MPI is required when num_mpiprocs_per_machine > 1.')
        if options.max_wallclock_seconds < self._MIN_OUTPUT_BUFFER_TIME:
            raise exceptions.InputValidationError(
                f'max_wallclock_seconds must be greater than {self._MIN_OUTPUT_BUFFER_TIME}.'
            )

    def _validate_inputs(self):
        """Validate input parameters."""
        parameters = self.inputs.parameters.get_dict()
        if 'postscf' in parameters:
            for postscf in parameters['postscf']:
                if postscf['$function'] not in self._SUPPORTED_POSTSCF:
                    raise exceptions.InputValidationError(f"Unsupported postscf function: {postscf['$function']}")

        # We want the option to be set for `verdi calcjob inputcat` to work,
        # but we don't allow overriding it because it would affect the name of the log file.
        if self.metadata.options.input_filename != self.INPUT_FILENAME:
            raise exceptions.InputValidationError(
                f"Input filename must be {self.INPUT_FILENAME}. Was: {self.metadata.options.input_filename}"
            )

    def _validate_pseudos(self):
        """Validate the pseudopotentials.

        Check that there is a one-to-one map of kinds in the structure to pseudopotentials.
        """
        kinds = set(kind.name for kind in self.inputs.structure.kinds)
        pseudos = set(self.inputs.pseudos.keys())
        if kinds != pseudos:
            raise exceptions.InputValidationError(
                'Mismatch between the defined pseudos and the list of kinds of the structure.\n'
                f'Pseudos: {pseudos};\nKinds:{kinds}'
            )

    def _validate_kpoints(self):
        """Validate the k-points intput.

        Check that the input k-points provide a k-points mesh.
        """
        try:
            self.inputs.kpoints.get_kpoints_mesh()
        except AttributeError as exc:
            raise exceptions.InputValidationError('The kpoints input does not have a valid mesh set.') from exc

    def _generate_inputdata(
        self, parameters: orm.Dict, structure: orm.StructureData, pseudos: dict, kpoints: orm.KpointsData
    ) -> ty.Tuple[dict, list]:
        """Generate the input dict (json) for DFTK.

        :param parameters: a dict defines the calculation parameters for DFTK
        :param structre: a StructureDate define the crystal
        :param pseudos: a dict contains the pseudos
        :param kpoints: a KpointData
        :return: dict for the DFTK json input
        :return: list of a pseudos needed to be copied
        """

        local_copy_pseudo_list = []

        data = {
            'periodic_system': {},
            'pseudopotentials': {},
            'model_kwargs': {},
            'basis_kwargs': {},
            'scf': {},
            'postscf': [],
        }
        data['periodic_system']['bounding_box'] = [[x * units.ang_to_bohr for x in vec] for vec in structure.cell]
        data['periodic_system']['atoms'] = []
        for site in structure.sites:
            data['periodic_system']['atoms'].append({
                'symbol': site.kind_name,
                'position': [X * units.ang_to_bohr for X in list(site.position)],
            })
        for symbol, pseudo in pseudos.items():
            data['pseudopotentials'][symbol] = f'{self._PSEUDO_SUBFOLDER}{pseudo.filename}'
            local_copy_pseudo_list.append((pseudo.uuid, pseudo.filename, f'{self._PSEUDO_SUBFOLDER}{pseudo.filename}'))
        if 'pseudo_rcut' in self.inputs:
            data['pseudopotentials']['$kwargs'] = {'rcut': self.inputs.pseudo_rcut.value}
        data['basis_kwargs']['kgrid'], data['basis_kwargs']['kshift'] = kpoints.get_kpoints_mesh()

        # set the maxtime for the SCF cycle, with a margin of _MIN_OUTPUT_BUFFER_TIME and 10%, whichever leads to a larger margin
        maxtime = min(
            self.inputs.metadata.options.max_wallclock_seconds - self._MIN_OUTPUT_BUFFER_TIME,
            0.9 * self.inputs.metadata.options.max_wallclock_seconds,
        )
        data['scf']['maxtime'] = maxtime
        
        DftkCalculation._merge_dicts(data, parameters.get_dict())

        return data, local_copy_pseudo_list

    def _generate_retrieve_list(self, parameters: orm.Dict) -> list:
        """Generate the list of files to retrieve based on the type of calculation requested in the input parameters.

        :param parameters: input parameters
        :returns: list of files to retreive
        """
        parameters = parameters.get_dict()
        # Retrieve the postscf files, all function.hdf5 except compute_bands.json
        retrieve_list = [
            f"{item['$function']}.json" if item['$function'] == 'compute_bands' else f"{item['$function']}.hdf5"
            for item in parameters['postscf']
        ]
        retrieve_list.append(self.LOGFILE)
        retrieve_list.append('timings.json')
        retrieve_list.append(f'{self.SCFRES_SUMMARY_NAME}')
        return retrieve_list

    def prepare_for_submission(self, folder):
        """Create the input file(s) from the input nodes.

        :param folder: an `aiida.common.folders.Folder` where the plugin should temporarily place all files needed by
            the calculation.
        :return: `aiida.common.datastructures.CalcInfo` instance
        """

        self._validate_options()
        self._validate_inputs()
        self._validate_pseudos()
        self._validate_kpoints()

        # Create lists which specify files to copy and symlink
        remote_copy_list = []
        remote_symlink_list = []

        # Generate the input file content
        input_filecontent, local_copy_list = self._generate_inputdata(self.inputs.parameters, self.inputs.structure, self.inputs.pseudos, self.inputs.kpoints)

        # write input file
        input_filename = folder.get_abs_path(self.metadata.options.input_filename)
        with io.open(input_filename, 'w', encoding='utf-8') as stream:
            json.dump(input_filecontent, stream, indent=4)

        # List the files (scfres.jld2) to copy or symlink in the case of a restart
        if 'parent_folder' in self.inputs:
            # Symlink if on the same computer, otherwise copy
            same_computer = self.inputs.code.computer.uuid == self.inputs.parent_folder.computer.uuid
            checkpointfile_info = (
                self.inputs.parent_folder.computer.uuid,
                os.path.join(self.inputs.parent_folder.get_remote_path(), self.inputs.parameters['scf']['checkpointfile']),
                self.inputs.parameters['scf']['checkpointfile']
            )
            if same_computer:
                remote_symlink_list.append(checkpointfile_info)
            else:
                remote_copy_list.append(checkpointfile_info)

        # prepare command line parameters
        cmdline_params = [
            # Precompilation under MPI generally deadlocks. Make sure everything is already precompiled.
            '--compiled-modules=strict',
            '-e', 'using AiidaDFTK; AiidaDFTK.run(inputfile="{}", allowed_versions="{}")'.format(
                self.metadata.options.input_filename,
                _AIIDA_DFTK_VERSION_SPEC,
            ),
        ]

        # prepare retrieve list
        retrieve_list = self._generate_retrieve_list(self.inputs.parameters)

        # Set up the `CodeInfo` to pass to `CalcInfo`
        codeinfo = datastructures.CodeInfo()
        codeinfo.code_uuid = self.inputs.code.uuid
        codeinfo.cmdline_params = cmdline_params

        # Set up the `CalcInfo` so AiiDA knows what to do with everything
        calcinfo = datastructures.CalcInfo()
        calcinfo.codes_info = [codeinfo]
        calcinfo.retrieve_list = retrieve_list
        calcinfo.remote_symlink_list = remote_symlink_list
        calcinfo.remote_copy_list = remote_copy_list
        calcinfo.local_copy_list = local_copy_list

        return calcinfo

class PrecompileCalculation(CalcJob):
    """Calcjob implementation to precompile AiidaDFTK."""

    _SUCCESS_PRINT = "Import successful"

    @classmethod
    def define(cls, spec):
        """Define the process specification."""
        super().define(spec)

        options = spec.inputs['metadata']['options']
        options['resources'].default = {'num_machines': 1, 'num_mpiprocs_per_machine': 1}

    def prepare_for_submission(self, folder):
        # Set up the `CodeInfo` to pass to `CalcInfo`
        codeinfo = datastructures.CodeInfo()
        codeinfo.code_uuid = self.inputs.code.uuid
        codeinfo.cmdline_params = [
            '-e', f'using Pkg; Pkg.precompile(; strict=true); using AiidaDFTK; println("{self._SUCCESS_PRINT}")'
        ]
        codeinfo.withmpi = False

        # Set up the `CalcInfo` so AiiDA knows what to do with everything
        calcinfo = datastructures.CalcInfo()
        calcinfo.codes_info = [codeinfo]

        return calcinfo

    # Easier to override the parse method than to write a parser.
    def parse(self, *args, **kwargs):
        exit_code = super().parse(*args, **kwargs)
        if exit_code.status != 0:
            return exit_code

        retrieved = self.node.outputs.retrieved
        filename_stdout = self.node.get_option('scheduler_stdout')

        if filename_stdout is None:
            self.report('could not determine `stdout` filename because `scheduler_stdout` option was not set.')
        else:
            try:
                scheduler_stdout = retrieved.base.repository.get_object_content(filename_stdout, mode='r')
                if self._SUCCESS_PRINT in scheduler_stdout:
                    return ExitCode(0)
            except FileNotFoundError:
                self.report(f'could not parse scheduler output: the `{filename_stdout}` file is missing')

        return self.exit_codes.ERROR_UNSPECIFIED
