"""
A module of base class for CASTEP calculations
"""
from __future__ import print_function

import copy
import os
import time
from textwrap import TextWrapper

from aiida.backends.utils import get_current_profile
from aiida.common.datastructures import CalcInfo, CodeInfo
from aiida.common.exceptions import InputValidationError, MultipleObjectsError
from aiida.common.utils import classproperty
from aiida.orm import DataFactory
from aiida_castep.calculations.datastructure import CellFile, ParamFile
from aiida_castep.calculations.helper import CastepHelper
from aiida_castep.data import OTFGData, UspData, get_pseudos_from_structure

from .utils import get_castep_ion_line, _lowercase_dict, _uppercase_dict

from .._version import calc_parser_version
__version__ = calc_parser_version

StructureData = DataFactory("structure")
ParameterData = DataFactory("parameter")
KpointsData = DataFactory("array.kpoints")
UpfData = DataFactory("upf")
RemoteData = DataFactory("remote")


class BaseCastepInputGenerator(object):
    """
    Baseclass for generating CASTEP inputs
    """
    # This may not apply to CASTEP runs
    _PSEUDO_SUBFOLDER = "./pseudo/"
    # Castep output is in the
    _PARENT_CALC_SUBFOLDER = "./parent/"
    _PREFIX = 'aiida'
    _INPUT_FILE_NAME = "aiida.cell"
    _OUTPUT_FILE_NAME = "aiida.castep"
    _SEED_NAME = 'aiida'

    # Additional files that should always be retrieved if they exist
    _internal_retrieve_list = [
        "*.err", "*.den_fmt", "*-out.cell", "*.pdos_bin",
    ]

    # in restarts, will not copy but use symlinks
    _default_symlink_usage = True

    # in restarts, it will copy from the parent the following
    _restart_copy_from = "./"

    # in restarts, it will copy the previous folder in the following one
    _restart_copy_to = _PARENT_CALC_SUBFOLDER

    # Default verbosity; change in subclasses
    _default_verbosity = 1

    # whether we are automatically validating the input parameters
    _auto_input_validation = True

    retrieve_dict = {
        "phonon": [".phonon"],
        "magres": [".magres"],
        "transitionstatesearch": [".ts"],
        "molecular dynamics": [".md"],
        "moleculardynamics": [".md"],
        "geometryoptimisation": [".geom"],
        "geometryoptimization": [".geom"],
        "spectral": [".ome_bin", ".dome_bin"],
    }

    @classproperty
    def _baseclass_use_methods(cls):
        """
        This will be manually added to the _use_methods in each subclass
        """
        return {
            "structure": {
                'valid_types': StructureData,
                'additional_parameter': None,
                'linkname': 'structure',
                'docstring': "Choose the input structure to use",
            },
            "settings": {
                'valid_types': ParameterData,
                'additional_parameter': None,
                'linkname': 'settings',
                'docstring': "Use an additional node for special settings",
            },
            "parameters": {
                'valid_types':
                ParameterData,
                'additional_parameter': None,
                'linkname': 'parameters',
                'docstring':
                ("Use a node that specifies the input parameters.",
                 "Usually needs to contain PARAM and CELL keys.")
            },
            "parent_folder": {
                'valid_types':
                RemoteData,
                'additional_parameter':
                None,
                'linkname':
                'parent_calc_folder',
                'docstring': ("Use a remote folder as parent folder (for "
                              "restarts and similar"),
            },
            "pseudo": {
                'valid_types': (UpfData, OTFGData, UspData),
                'additional_parameter':
                "kind",
                'linkname':
                cls._get_linkname_pseudo,
                'docstring':
                ("Use a node for the  pseudopotential of one of "
                 "the elements in the structure. You have to pass "
                 "an additional parameter ('kind') specifying the "
                 "name of the structure kind (i.e., the name of "
                 "the species) for which you want to use this "
                 "pseudo. You can pass either a string, or a "
                 "list of strings if more than one kind uses the "
                 "same pseudo"),
            },
        }

    def _generate_header_lines(self, other_nodes=None):
        """
        Generate header lines to go into param and cell files
        """

        wrapper = TextWrapper(initial_indent="# ", subsequent_indent="# ")
        time_str = time.strftime("%H:%M:%S %d/%m/%Y %Z")
        lines = [
            "##### Generated by aiida_castep {} #####".format(time_str),
            "#         author: Bonan Zhu (bz240@cam.ac.uk)",
            "# "
            "# AiiDA User: {}".format(self.get_user().get_full_name()),
            "# AiiDA profile: {}".format(get_current_profile()),
            "# Information of the calculation node",
            "# type: {}".format(self._query_type_string[:-1]),
            "# pk: {}".format(self.pk),
            "# uuid: {}".format(self.uuid),
            "# label: {}".format(self.label),
            "# description:",
        ]

        # Fix for SQLA backend - default of node.description is None
        if self.description:
            lines.extend(wrapper.wrap(self.description))
        lines.append("")

        # additional information of the input nodes
        if other_nodes:
            lines.append("## Information of input nodes used:")
        for node in other_nodes:
            if node:
                node_lines = [
                    "# ", "# type: {}".format(node._query_type_string[:-1]),
                    "# pk: {}".format(node.pk), "# uuid: {}".format(node.uuid),
                    "# label: {}".format(node.label), "# description:"
                ]
                _desc = node.description
                if _desc:
                    node_lines.extend(wrapper.wrap(_desc))
                node_lines.append("")
                lines.extend(node_lines)

        lines.append("# END OF HEADER")

        return lines

    def _generate_CASTEPinputdata(self,
                                  parameters,
                                  structure,
                                  pseudos,
                                  settings_dict,
                                  settings=None,
                                  kpoints=None,
                                  **kwargs):
        """
        This method creates the content of an input file in the
        CASTEP format
        Generated input here should be generic to all castep calculations

        :param parameters, ParameterData: Input goes to the .cell or .param file
        :param settings_dict: A dictionary of the settings used for generation
        :param pseudos: A dictionary of pseudo potential Data for each kind
        :param structure: A StructureData instance
        :param kpoints: A KpointsData node, optional
        """
        local_copy_list_to_append = []

        # The input dictionary should be {"CELL": {"fix_all_cell": True},
        # "PARAM":{"opt_strategy" : "speed"}}

        # The following lines enforces the case of names
        # EACH entry is written as in lower case

        input_params = _uppercase_dict(
            parameters.get_dict(), dict_name="parameters")
        input_params = {
            k: _lowercase_dict(v, dict_name=k)
            for k, v in input_params.iteritems()
        }

        # Check if there are keywords that need to be blocked

        # Set verbosity to 1.
        # Parser may not work if verbosity is not 1
        input_params["PARAM"]["iprint"] = input_params["PARAM"].get(
            "iprint", self._default_verbosity)

        # Set run_time using define value for this calcualtion
        run_time = self.get_max_wallclock_seconds()
        if run_time:
            n_seconds = run_time * 0.95
            n_seconds = (
                n_seconds // 60) * 60  # Round down to the nearest minutes
            # Do not do any thing if calculated time is less than 1 hour
            if n_seconds < 3600:
                pass
            elif "run_time" not in input_params["PARAM"]:
                input_params["PARAM"]["run_time"] = int(n_seconds)

        # Set the default comment using the label of this calculation
        comment_str = self.label
        if "comment" not in input_params["PARAM"]:
            input_params["PARAM"]["comment"] = comment_str

        # ========= Start to prepare input site data ======

        cellfile = CellFile()
        paramfile = ParamFile()
        # --------- CELL ----------
        cell_vector_list = []
        for vector in structure.cell:
            cell_vector_list.append(("{0:18.10f} {1:18.10f} "
                                     "{2:18.10f}".format(*vector)))

        cellfile["LATTICE_CART"] = cell_vector_list

        # --------- ATOMIC POSITIONS---------
        # for kind in structure.kinds:
        atomic_position_list = []
        mixture_count = 0
        # deal with initial spins
        spin_list = settings_dict.pop("SPINS", None)
        label_list = settings_dict.pop("LABELS", None)

        for i, site in enumerate(structure.sites):
            # get  the kind of the site
            kind = structure.get_kind(site.kind_name)

            # Position is always needed
            pos = site.position
            try:
                name = kind.symbol
            # If we are dealing with mixed atoms
            except ValueError:
                name = kind.symbols
                mixture_count += 1

            if spin_list:
                spin = spin_list[i]
            else:
                spin = None

            # deal with labels
            if label_list:
                label = label_list[i]
            else:
                label = None
            # Get the line of positions_abs block
            line = get_castep_ion_line(
                name,
                pos,
                label=label,
                spin=spin,
                occupation=kind.weights,
                mix_num=mixture_count)

            # Append the line to the list
            atomic_position_list.append(line)

        # End of the atomic position block
        cellfile["POSITIONS_ABS"] = atomic_position_list

        # Check the consistency of spin in parameters
        if spin_list:
            total_spin = sum(s for s in spin_list if s)
            param_spin = input_params["PARAM"].get("spin", None)
            if param_spin is not None:
                # If spin is specified - check consistency
                if param_spin != total_spin:
                    raise InputValidationError(
                        "Inconsistent spin in cell and param files."
                        "Total spin: {} in cell file but {} in param file".
                        format(total_spin, param_spin))
            else:
                # If no spin specified, do it automatically
                # Note that we don't check if spin polarized calculation is
                # requested in the first place
                input_params["PARAM"]["spin"] = total_spin

        # --------- KPOINTS ---------
        if self._use_kpoints:
            try:
                mesh, offset = kpoints.get_kpoints_mesh()
                has_mesh = True

            except AttributeError:
                try:
                    kpoints_list = kpoints.get_kpoints()
                    num_kpoints = len(kpoints_list)
                    has_mesh = False
                    if num_kpoints == 0:
                        raise InputValidationError(
                            "At least one k points must be provided")
                except AttributeError:
                    raise InputValidationError(
                        "No valid kpoints have been found")

                try:
                    _, weights = kpoints.get_kpoints(also_weights=True)

                except AttributeError:
                    import numpy as np
                    weights = np.ones(num_kpoints, dtype=float) / num_kpoints

            kpoints_line_list = []
            if has_mesh is True:
                cellfile["kpoints_mp_grid"] = "{} {} {}".format(*mesh)
                if offset != [0., 0., 0.]:
                    cellfile["kpoints_mp_offset"] = "{} {} {}".format(*offset)
            else:
                for kpoint, weight in zip(kpoints_list, weights):
                    kpoints_line_list.append("{:18.10f} {:18.10f} "
                                             "{:18.10f} {:18.10f}".format(
                                                 kpoint[0], kpoint[1],
                                                 kpoint[2], weight))
                cellfile["KPOINTS_LIST"] = kpoints_line_list

        # --------- PSEUDOPOTENTIALS --------
        # Check if we are using UPF pseudos
        # Now only support simple elemental pseudopotentials
        if pseudos:
            symbols = set()  # All of the symbols
            species_pot_list = []
            for kind in structure.kinds:
                for s in kind.symbols:
                    symbols.add(s)

            # Make symbols unique
            for s in symbols:
                ps = pseudos[s]  # Get the pseudopotential object

                # If we are dealing with a UpfData object
                if isinstance(ps, (UpfData, UspData)):
                    species_pot_list.append("{:5} {}".format(s, ps.filename))
                    # Add to the copy list
                    local_copy_list_to_append.append((ps.get_file_abs_path(),
                                                      ps.filename))

                # If we are using OTFG, just add the string property of it
                if isinstance(ps, OTFGData):
                    species_pot_list.append("{:5} {}".format(s, ps.string))

            cellfile["SPECIES_POT"] = species_pot_list

        # --------- PARAMETERS in cell file---------
        for key, value in input_params["CELL"].iteritems():

            if "species_pot" in key:
                if pseudos:
                    raise MultipleObjectsError(
                        "Both species_pot and pseudos are provided")
                self.logger.warning(
                    "Pseudopotentials directly defined in CELL dictionary")

            # Constructing block keywrods
            # We identify the key should be treated as a block it is not a string and has len() > 0
            cellfile[key] = value

        # Parameters for PARAM files
        paramfile.update(input_params["PARAM"])
        param_header = self._generate_header_lines([parameters])
        paramfile.header = param_header

        # Added header to .cell file
        cell_file_header = self._generate_header_lines(
            [parameters, structure, kpoints, settings])
        cellfile.header = cell_file_header

        return cellfile, paramfile, local_copy_list_to_append

    def _prepare_for_submission(self, tempfolder, inputdict):
        """
        Routine to be called when create the input files and other stuff

        :param tempfolder: a aiida.common.folders.Folder subclass where
                           the plugin should put all its files.
        :param inputdict: a dictionary with the input nodes, as they would
                be returned by get_inputs_dict (without the Code!)
        """

        local_copy_list = []

        # TODO implement remote copy for restart calculations
        remote_copy_list = []
        remote_symlink_list = []

        try:
            parameters = inputdict.pop(self.get_linkname('parameters'))
            # Validate the parameters
            if self._auto_input_validation is True:
                self.check_castep_input(parameters.get_dict(), auto_fix=False)

        except KeyError:
            raise InputValidationError(
                "No parameters specified for this calculation")
        if not isinstance(parameters, ParameterData):
            raise InputValidationError(
                "parameters is not of type ParameterData")

        try:
            structure = inputdict.pop(self.get_linkname('structure'))
        except KeyError:
            raise InputValidationError(
                "No structure specified for this calculation")
        if not isinstance(structure, StructureData):
            raise InputValidationError(
                "structure is not of type StructureData")

        if self._use_kpoints:
            try:
                kpoints = inputdict.pop(self.get_linkname('kpoints'))
            except KeyError:
                raise InputValidationError(
                    "No kpoints specified for this calculation")
            if not isinstance(kpoints, KpointsData):
                raise InputValidationError(
                    "kpoints is not of type KpointsData")
        else:
            kpoints = None

        # Settings can be undefined, and defaults to an empty dictionary
        settings = inputdict.pop(self.get_linkname('settings'), None)
        if settings is None:
            settings_dict = {}
        else:
            if not isinstance(settings, ParameterData):
                raise InputValidationError(
                    "settings, if specified, must be of "
                    "type ParameterData")
            # Settings converted to uppercase
            settings_dict = _uppercase_dict(
                settings.get_dict(), dict_name='settings')

        # Set up pseodo potentials
        # The existance of pseudopoential is optional
        pseudos = {}
        # I create here a dictionary that associates each kind name to a pseudo
        for link in inputdict.keys():
            if link.startswith(self._get_linkname_pseudo_prefix()):
                kindstring = link[len(self._get_linkname_pseudo_prefix()):]
                kinds = kindstring.split('_')
                the_pseudo = inputdict.pop(link)
                if not isinstance(the_pseudo, (UpfData, UspData, OTFGData)):
                    raise InputValidationError(
                        "Pseudo for kind(s) {} is not of "
                        "supported ".format(",".join(kinds)))
                for kind in kinds:
                    if kind in pseudos:
                        raise InputValidationError(
                            "Pseudo for kind {} passed "
                            "more than one time".format(kind))
                    pseudos[kind] = the_pseudo

        # Check parent calc folder
        parent_calc_folder = inputdict.pop(
            self.get_linkname('parent_folder'), None)
        if parent_calc_folder is not None:
            if not isinstance(parent_calc_folder, RemoteData):
                raise InputValidationError("parent_calc_folder, if specified, "
                                           "must be of type RemoteData")

        # If requested to reuse, check if the parent_calc_folder is defined
        param_dict = parameters.get_dict()["PARAM"]
        require_parent = False
        for k in param_dict:
            if str(k).lower() in ["reuse", "continuation"]:
                require_parent = True
                break
        if parent_calc_folder is None and require_parent:
            raise InputValidationError(
                "No parent calculation folder passed"
                " for restart calculation using reuse/continuation")

        # Check if a code is specified
        try:
            code = inputdict.pop(self.get_linkname('code'))
        except KeyError:
            raise InputValidationError(
                "No code specified for this calculation")

        # Here, there should be no more parameters...
        # But in case there is, check if this is something not implemented
        # in this base class
        if inputdict:
            for key in inputdict:
                if key not in self._use_methods:
                    raise InputValidationError(
                        "The following input data nodes are "
                        "unrecognised: {}".format(inputdict.keys()))
        ##############################
        # END OF INITIAL INPUT CHECK #
        ##############################

        # Generate input file
        cellfile, paramfile, pseudo_copy_list = self._generate_CASTEPinputdata(
            parameters,
            structure,
            pseudos,
            settings_dict,
            settings=settings,
            kpoints=kpoints,
            **inputdict)

        local_copy_list.extend(pseudo_copy_list)

        cell_input_filename = tempfolder.get_abs_path(self._SEED_NAME +
                                                      ".cell")

        param_input_filename = tempfolder.get_abs_path(self._SEED_NAME +
                                                       ".param")

        with open(cell_input_filename, "w") as incell:
            incell.write(cellfile.get_string())

        with open(param_input_filename, "w") as inparam:
            inparam.write(paramfile.get_string())

        # IMPLEMENT OPERATIONS FOR RESTART
        symlink = settings_dict.pop('PARENT_FOLDER_SYMLINK',
                                    self._default_symlink_usage)

        if parent_calc_folder is not None:
            if symlink:
                remote_list = remote_symlink_list
            else:
                remote_list = remote_copy_list
            remote_list.append(
                (parent_calc_folder.get_computer().uuid,
                 os.path.join(parent_calc_folder.get_remote_path(),
                              self._restart_copy_from), self._restart_copy_to))

        calcinfo = CalcInfo()

        calcinfo.uuid = self.uuid

        # COPY/SYMLINK LISTS
        calcinfo.local_copy_list = local_copy_list
        calcinfo.remote_copy_list = remote_copy_list
        calcinfo.remote_symlink_list = remote_symlink_list

        # SET UP extra CMDLINE arguments
        cmdline_params = settings_dict.pop("CMDLINE", [])

        # Extra parameters are added after the seed for CASTEP
        calcinfo.cmdline_params = [self._SEED_NAME] + list(cmdline_params)

        # CASTEP don't have any STDOUT etc when running calculations
        # Error is shown in the *.err file

        # Construct codeinfo instance
        codeinfo = CodeInfo()
        codeinfo.cmdline_params = [self._SEED_NAME] + list(cmdline_params)
        codeinfo.code_uuid = code.uuid
        calcinfo.codes_info = [codeinfo]

        # Retrieve by default the .castep file and the bands file
        calcinfo.retrieve_list = []

        calcinfo.retrieve_list.append(self._SEED_NAME + ".castep")
        calcinfo.retrieve_list.append(self._SEED_NAME + ".bands")

        settings_retrieve_list = settings_dict.pop("ADDITIONAL_RETRIEVE_LIST",
                                                   [])
        calcinfo.retrieve_list.extend(settings_retrieve_list)

        # If we are doing geometryoptimisation retrieved the geom file and -out.cell file
        calculation_mode = parameters.get_dict().get("PARAM", {}).get(
            "task", "")

        # dictionary for task specific file retrieve
        task_extra = self.retrieve_dict.get(calculation_mode.lower(), [])
        for suffix in task_extra:
            settings_retrieve_list.append(self._SEED_NAME + suffix)

        # Retrieve output cell file if requested
        if parameters.get_dict().get("PARAM", {}).get("write_cell_structure"):
            settings_retrieve_list.append(self._SEED_NAME + "-out.cell")

        calcinfo.retrieve_list += settings_retrieve_list
        calcinfo.retrieve_list += self._internal_retrieve_list

        # Remove parser options in the setting dictionary
        # At the moment parser options are not used here
        try:
            Parserclass = self.get_paaserclass()
            parser = Parserclass(self)
            parser_opts = parser.get_parser_settings_key()
            settings_dict.pop(parser_opts)
        except (KeyError, AttributeError):
            pass

        if settings_dict:
            raise InputValidationError(
                "The following keys have been found in "
                "the settings input node, but were not understood: {}".format(
                    ",".join(settings_dict.keys())))

        return calcinfo

    def _set_parent_remotedata(self, remotedata):
        """
        Used to set a parent remotefolder
        """
        from aiida.common.exceptions import ValidationError

        if not isinstance(remotedata, RemoteData):
            raise ValueError('remotedata must be a RemoteData')

        # complain if another remotedata is already found
        input_remote = self.get_inputs(node_type=RemoteData)
        if input_remote:
            raise ValidationError("Cannot set several parent calculation to a "
                                  "{} calculation".format(
                                      self.__class__.__name__))

        self.use_parent_folder(remotedata)

    def create_restart(self,
                       ignore_state=False,
                       restart_type="restart",
                       reuse=False,
                       use_symlink=None,
                       use_output_structure=False,
                       use_castep_bin=False,
                       param_update=None,
                       param_delete=None):
        """
        A method to help restart the calculation by creating a new one.
        Return a new calculation with all the essential input nodes in the unstored stated.

        CASTEP has two modes of 'restart', activated by setting CONTINUATION or REUSE keywords in .param file:

        CONTINUATION
          Restart from the end of the last run.
          Only limited set of parameters can be modified.
          If unmodifiable parameters were changed, they are ignored.
          E.g changes of task, nextra_bands, cut_off_energy will be ignored. This is often used for geometry optimisation or md runs.

        REUSE
          Essentially making a new calculation with parameters read from .cell and .param file.
          Data from ``castep_bin`` or ``check`` will be reused to initialise the model of the new run when applicable. This is often used for bandstructure, dos, spectral calculation.

        Note both castep_bin and check file may be used.
        They are almost the same except castep_bin does not have wavefunctions stored.


        :param bool ignore_state: Ignore the state of parent calculation
        :param str restart_type: "continuation" or "restart". If set to continuation the child calculation has keyword 'continuation' set.
        :param bool reuse: Wether we want to reuse the previous calculation.
            only applies for "restart" run
        :param bool parent_folder_symlink: if True, symlink are used instead of hard copies of the files. Default given be cin._default_symlink_usage
        :param bool use_output_structure: if True, the output structure of parent calculation is used as the input of the child calculation.
            This is useful for photon/bs calculation.
        :param calc_class: Class of the new calculation, e.g can restart into a CastepBSCalculation.
        :param param_update: A dictionary of the parameter keywords to be updated.
            The key can be put into the top level and will be assigned accordingly.
        :param param_delete: A list of the keys to be removed from param file or cell file
        :param keep_input_model: If True, keeps the input model of the parent calculation.
            i.e. use input check file rather that the check file written.

        ..see also:: :py:func:`aiida_castep.calculations.base.create_restart_`
        """

        cout = create_restart_(
            self,
            ignore_state=ignore_state,
            restart_type=restart_type,
            reuse=reuse,
            use_symlink=use_symlink,
            use_output_structure=use_output_structure,
            use_castep_bin=use_castep_bin,
            param_update=param_update,
            param_delete=param_delete)
        return cout

    @classmethod
    def continue_from(cls, cin, *args, **kwargs):
        """
        Create a new calculation as a continuation from a given calculation.
        This is effectively an "restart" for CASTEP and a lot of the parameters
        can be tweaked. For example, conducting bandstructure calculation from
        finished geometry optimisations.
        CASTEP has two modes of 'restart', activated by setting CONTINUATION or REUSE keywords in .param file.

        ..see also:: :py:func:`aiida_castep.calculations.base.create_restart_`
        """
        cout = create_restart_(cin, calc_class=cls, *args, **kwargs)

        return cout

    @classmethod
    def check_castep_input(cls, input_dict, auto_fix=False):
        """
        Validate if the keywords are valid castep keywords
        Also try to convert the parameter diction in a
        "standardized" form
        """
        helper = cls.get_input_helper()
        out = helper.check_dict(input_dict, auto_fix)
        return out

    @classmethod
    def get_input_helper(cls, version=None):
        """
        Get a CASTEP help object.
        :param version: Specify a specific version of the help
        :type version: str

        :return: A CastepHelper object
        """
        helper = CastepHelper(version)
        return helper

    @classmethod
    def _get_linkname_pseudo_prefix(cls):
        """
        The prefix for the name of the link used for each pseudo before the kind name
        """
        return "pseudo_"

    @classmethod
    def _get_linkname_pseudo(cls, kind):
        """
        The name of the link used for the pseudo for kind 'kind'.
        It appends the pseudo name to the pseudo_prefix, as returned by the
        _get_linkname_pseudo_prefix() method.

        :note: if a list of strings is given, the elements are appended
          in the same order, separated by underscores

        :param kind: a string (or list of strings) for the atomic kind(s) for
            which we want to get the link name
        """
        # If it is a list of strings, and not a single string: join them
        # by underscore
        if isinstance(kind, (tuple, list)):
            suffix_string = "_".join(kind)
        elif isinstance(kind, basestring):
            suffix_string = kind
        else:
            raise TypeError("The parameter 'kind' of _get_linkname_pseudo can "
                            "only be a string or a list of strings")
        return "{}{}".format(cls._get_linkname_pseudo_prefix(), suffix_string)

    @classmethod
    def get_restart_file_relative_path(cls,
                                       param_data_dict,
                                       use_castep_bin=False):
        """
        Returns a relative path of the restart file
        """
        parent_check_name = param_data_dict["PARAM"].get("check_point", None)
        if parent_check_name is None:
            if use_castep_bin:
                parent_check_name = cls._SEED_NAME + ".castep_bin"
            else:
                parent_check_name = cls._SEED_NAME + ".check"

        return os.path.join(cls._PARENT_CALC_SUBFOLDER, parent_check_name)

    def use_pseudos_from_family(self, family_name):
        """
        Set the pseudo to use for all atomic kinds, picking pseudos from the
        family with name family_name.


        :note: The structure must already be set.

        :param family_name: the name of the group containing the pseudos
        """
        from collections import defaultdict

        try:
            structure = self._get_reference_structure()
        except AttributeError:
            raise ValueError(
                "Structure is not set yet! Therefore, the method "
                "use_pseudos_from_family cannot automatically set "
                "the pseudos")

        # A dict {kind_name: pseudo_object}
        # But we want to run with use_pseudo(pseudo, kinds)

        kind_pseudo_dict = get_pseudos_from_structure(structure, family_name)

        # Group the species by pseudo
        # pseudo_dict will just map PK->pseudo_object
        pseudo_dict = {}
        # Will contain a list of all species of the pseudo with given PK
        pseudo_species = defaultdict(list)

        for kindname, pseudo in kind_pseudo_dict.iteritems():
            pseudo_dict[pseudo.pk] = pseudo
            pseudo_species[pseudo.pk].append(kindname)

        # Finally call the use_pseudo method
        for pseudo_pk in pseudo_dict:
            pseudo = pseudo_dict[pseudo_pk]
            kinds = pseudo_species[pseudo_pk]
            # I set the pseudo for all species, sorting alphabetically
            self.use_pseudo(pseudo, sorted(kinds))

    def _get_reference_structure(self):
        """
        Used to get the reference structure to obtain which
        pseudopotentials to use from a given family using
        use_pseudos_from_family.

        :note: this method can be redefined in a given subclass
               to specify which is the reference structure to consider.
        """
        return self.get_inputs_dict()[self.get_linkname('structure')]


def create_restart_(cin,
                    ignore_state=False,
                    restart_type="restart",
                    reuse=False,
                    use_symlink=None,
                    use_output_structure=False,
                    use_castep_bin=False,
                    calc_class=None,
                    param_update=None,
                    param_delete=None,
                    keep_input_model=False):
    """
    A function to help restart the calculation by creating a new one.
    Return a new calculation with all the essential input nodes in the unstored stated.

    CASTEP has two modes of 'restart', activated by setting CONTINUATION or REUSE keywords in .param file:

    CONTINUATION
      Restart from the end of the last run.
      Only limited set of parameters can be modified.
      If unmodifiable parameters were changed, they are ignored.
      E.g changes of task, nextra_bands, cut_off_energy will be ignored. This is often used for geometry optimisation or md runs.

    REUSE
      Essentially making a new calculation with parameters read from .cell and .param file.
      Data from ``castep_bin`` or ``check`` will be reused to initialise the model of the new run when applicable. This is often used for bandstructure, dos, spectral calculation.

    Note both castep_bin and check file may be used.
    They are almost the same except castep_bin does not have wavefunctions stored.

    :param bool ignore_state: Ignore the state of parent calculation
    :param str restart_type: "continuation" or "restart". If set to continuation the child calculation has keyword 'continuation' set.
    :param bool reuse: Wether we want to reuse the previous calculation.
        only applies for "restart" run
    :param bool parent_folder_symlink: if True, symlink are used instead of hard copies of the files.
        Default given be cin._default_symlink_usage
    :param bool use_output_structure: if True, the output structure of parent calculation is used as the input of the child calculation.
        This is useful for photon/bs calculation.
    :param calc_class: Class of the new calculation. E.g can restart into a CastepBSCalculation.
    :param param_update: A dictionary of the parameter keywords to be updated.
        The key can be put into the top level and will be assigned accordingly.
    :param param_delete: A list of the keys to be removed from param file or cell file
    :param keep_input_model: If True, keeps the input model of the parent calculation.
        i.e. use input check file rather that the check file written.
    """

    from aiida.common.datastructures import calc_states

    if cin.get_state(from_attribute=True) != calc_states.FINISHED:
        if not ignore_state:
            raise InputValidationError(
                "Calculation to be restarted must be in the {} state."
                "use ignore_state keyword to override this".format(
                    calc_states.FINISHED))

    if restart_type == "continuation":
        # If we do a conitniuation we actually have to re-use the file from previous run
        reuse = True
    elif restart_type == "restart":
        pass
    else:
        raise RuntimeError("Unkown type of restart {}".format(restart_type))

    if use_symlink is None:
        use_symlink = cin._default_symlink_usage

    # Duplicate the calculation
    if calc_class is None:
        cout = cin.duplicate()
    else:
        cout = calc_class()
        cout.set_computer(cin.get_computer())

    calc_inp = cin.get_inputs_dict()  # Input nodes of parent calculation

    # Find the RemoteData of this calculation
    if reuse:
        remote_folders = [
            o for o in cin.get_outputs() if isinstance(o, RemoteData)
        ]
        if len(remote_folders) > 1:
            raise InputValidationError("More than one output RemoteData found "
                                       "in calculation {}".format(cin.pk))
        elif not remote_folders:
            raise InputValidationError("No output RemoteData found "
                                       "in calculation {}".format(cin.pk))
        else:
            remote_folder = remote_folders[0]

    # We just keep the existing link
    elif keep_input_model:
        remote_folder = cin.get_inputs_dict().get(cin.get_linkname("parent_folder"))
    else:
        remote_folder = None

    if remote_folder:
        cout._set_parent_remotedata(remote_folder)

    # Use the out_put structure if required
    if use_output_structure:
        try:
            cout.use_structure(cin.out.output_structure)
        except AttributeError:
            cout.logger.warning(
                "Warning: No output structure found. Fallback to input structure"
            )
            cout.use_structure(calc_inp[cin.get_linkname('structure')])
    else:
        cout.use_structure(calc_inp[cin.get_linkname('structure')])

    # Copy the kpoints
    # NOTE this need to be changed for restarting BS with bs_kpoints etc
    # Move to another method to allow subclass modification
    if cin._use_kpoints:
        cout.use_kpoints(calc_inp[cin.get_linkname('kpoints')])
    cout.use_code(calc_inp[cin.get_linkname('code')])

    # copy the settings dictionary
    try:
        old_settings_dict = calc_inp[cin.get_linkname('settings')].get_dict()
    except KeyError:
        old_settings_dict = {}

    # Use deep copy to ensure two dictionaries are independent
    new_settings = copy.deepcopy(old_settings_dict)

    if use_symlink != cout._default_symlink_usage:
        new_settings['PARENT_FOLDER_SYMLINK'] = use_symlink

    if new_settings:
        if new_settings != old_settings_dict:
            # Link to an new settings
            settings = ParameterData(dict=new_settings)
            cout.use_settings(settings)
        else:
            # Nothing changed, just use the old settings
            cout.use_settings(calc_inp[cin.get_linkname('settings')])

    # SETUP the keyword in PARAM file
    parent_param = calc_inp[cin.get_linkname('parameters')]

    in_param_dict = parent_param.get_dict()
    if reuse:
        if restart_type == "restart":
            # Set keyword reuse, pop any continuation keywords
            in_param_dict['PARAM'].pop('continuation', None)
            # Define the name of reuse here
            in_param_dict['PARAM'][
                "reuse"] = cin.get_restart_file_relative_path(
                    in_param_dict, use_castep_bin)
        elif restart_type == "continuation":
            # Do the opposite
            in_param_dict['PARAM'].pop('reuse', None)
            in_param_dict['PARAM'][
                'continuation'] = cin.get_restart_file_relative_path(
                    in_param_dict, use_castep_bin)
        else:
            raise RuntimeError("Invalid restart type: " + restart_type)

    elif keep_input_model:
        # No need to change the continuation/reuse field
        pass
    else:
        # In this case we simply create a identical calculation
        # But we should discard reuse / continuation
        in_param_dict['PARAM'].pop('continuation', None)
        in_param_dict['PARAM'].pop('reuse', None)

    # UPDATE the key words
    if param_update:
        from .helper import HelperCheckError
        helper = cout.get_input_helper()
        dict_update, not_found = helper._from_flat_dict(param_update)
        if not_found:
            suggest = [helper.get_suggestion(i) for i in not_found]
            error_string = "Following keys are invalid -- "
            for error_key, sug in zip(not_found, suggest):
                error_string += "{}: {}; ".format(error_key, sug)
            raise HelperCheckError(error_string)
        else:
            in_param_dict["PARAM"].update(dict_update["PARAM"])
            in_param_dict["CELL"].update(dict_update["CELL"])

    # Remove keywords
    if param_delete:
        for key in param_delete:
            tmp1 = in_param_dict["PARAM"].pop(key, None)
            tmp2 = in_param_dict["CELL"].pop(key, None)
            if (tmp1 is None) and (tmp2 is None):
                raise RuntimeError("Key {} not found".format(key))

    # If we have not changed anything, link the old dict
    if in_param_dict == parent_param.get_dict():
        cout.use_parameters(parent_param)
    else:
        cout.use_parameters(ParameterData(dict=in_param_dict))

    # Use exactly the same pseudopotential data
    for linkname, input_node in cin.get_inputs_dict().iteritems():
        if isinstance(input_node, (UpfData, UspData, OTFGData)):
            cout.add_link_from(input_node, label=linkname)

    return cout
