text: # -*- coding: utf-8 -*-
  ###############################################################################
  # Copyright (c), Forschungszentrum Jülich GmbH, IAS-1/PGI-1, Germany.         #
  #                All rights reserved.                                         #
  # This file is part of the AiiDA-FLEUR package.                               #
  #                                                                             #
  # The code is hosted on GitHub at https://github.com/JuDFTteam/aiida-fleur    #
  # For further information on the license, see the LICENSE.txt file            #
  # For further information please visit http://www.flapw.de or                 #
  # http://aiida-fleur.readthedocs.io/en/develop/                               #
  ###############################################################################
  
  """
      In this module you find the workflow 'FleurSSDispConvWorkChain' for the calculation of
      Spin Spiral energy Dispersion converging all the directions.
  """
  
  from __future__ import absolute_import
  import copy
  import six
  
  from aiida.engine import WorkChain
  from aiida.engine import calcfunction as cf
  from aiida.orm import Dict
  from aiida.common import AttributeDict
  
  from aiida_fleur.workflows.scf import FleurScfWorkChain
  
  
  class FleurSSDispConvWorkChain(WorkChain):
      """
          This workflow calculates the Spin Spiral Dispersion of a structure.
      """
  
      _workflowversion = "0.2.0"
  
      _wf_default = {
          'beta': {'all': 1.57079},
          'q_vectors': {'label': [0.0, 0.0, 0.0],
                        'label2': [0.125, 0.0, 0.0]
                        }
      }
  
      @classmethod
      def define(cls, spec):
          super(FleurSSDispConvWorkChain, cls).define(spec)
          spec.expose_inputs(FleurScfWorkChain, namespace='scf')
          spec.input("wf_parameters", valid_type=Dict, required=False)
  
          spec.outline(
              cls.start,
              cls.converge_scf,
              cls.get_results,
              cls.return_results
          )
  
          spec.output('out', valid_type=Dict)
  
          # exit codes
          spec.exit_code(230, 'ERROR_INVALID_INPUT_PARAM',
                         message="Invalid workchain parameters.")
          spec.exit_code(340, 'ERROR_ALL_QVECTORS_FAILED',
                         message="Convergence SSDisp calculation failed for all q-vectors.")
          spec.exit_code(341, 'ERROR_SOME_QVECTORS_FAILED',
                         message="Convergence SSDisp calculation failed for some q-vectors.")
  
      def start(self):
          """
          Retrieve and initialize paramters of the WorkChain
          """
          self.report('INFO: started Spin Stiffness calculation'
                      ' convergence calculation workflow version {}\n'.format(self._workflowversion))
  
          self.ctx.info = []
          self.ctx.warnings = []
          self.ctx.errors = []
          self.ctx.energy_dict = []
  
          # initialize the dictionary using defaults if no wf paramters are given
          wf_default = copy.deepcopy(self._wf_default)
          if 'wf_parameters' in self.inputs:
              wf_dict = self.inputs.wf_parameters.get_dict()
          else:
              wf_dict = wf_default
  
          extra_keys = []
          for key in wf_dict.keys():
              if key not in wf_default.keys():
                  extra_keys.append(key)
          if extra_keys:
              error = 'ERROR: input wf_parameters for SSDisp Conv contains extra keys: {}'.format(
                  extra_keys)
              self.report(error)
              return self.exit_codes.ERROR_INVALID_INPUT_PARAM
  
          # extend wf parameters given by user using defaults
          for key, val in six.iteritems(wf_default):
              wf_dict[key] = wf_dict.get(key, val)
          self.ctx.wf_dict = wf_dict
  
      def converge_scf(self):
          """
          Converge charge density with or without SOC.
          Depending on a branch of Spiral calculation, submit a single Fleur calculation to obtain
          a reference for further force theorem calculations or
          submit a set of Fleur calculations to converge charge density for all given SQAs.
          """
          inputs = {}
          for key, q_vector in six.iteritems(self.ctx.wf_dict['q_vectors']):
              inputs[key] = self.get_inputs_scf()
              inputs[key].calc_parameters['qss'] = {'x': q_vector[0],
                                                    'y': q_vector[1],
                                                    'z': q_vector[2]}
              inputs[key].calc_parameters = Dict(dict=inputs[key]['calc_parameters'])
              res = self.submit(FleurScfWorkChain, **inputs[key])
              self.to_context(**{key: res})
  
      def get_inputs_scf(self):
          """
          Initialize inputs for scf workflow:
          wf_param, options, calculation parameters, codes, structure
          """
          input_scf = AttributeDict(self.exposed_inputs(FleurScfWorkChain, namespace='scf'))
  
          if 'wf_parameters' not in input_scf:
              scf_wf_dict = {}
          else:
              scf_wf_dict = input_scf.wf_parameters.get_dict()
  
          if 'inpxml_changes' not in scf_wf_dict:
              scf_wf_dict['inpxml_changes'] = []
  
          # change beta parameter
          for key, val in six.iteritems(self.ctx.wf_dict.get('beta')):
              scf_wf_dict['inpxml_changes'].append(('set_atomgr_att_label',
                                                    {'attributedict': {'nocoParams': [('beta', val)]},
                                                     'atom_label': key}))
  
          input_scf.wf_parameters = Dict(dict=scf_wf_dict)
  
          if 'calc_parameters' in input_scf:
              calc_parameters = input_scf.calc_parameters.get_dict()
          else:
              calc_parameters = {}
          input_scf.calc_parameters = calc_parameters
  
          return input_scf
  
      def get_results(self):
          """
          Retrieve results of converge calculations
          """
          t_energydict = {}
          outnodedict = {}
          htr_to_eV = 27.21138602
  
          for label in six.iterkeys(self.ctx.wf_dict['q_vectors']):
              calc = self.ctx[label]
  
              if not calc.is_finished_ok:
                  message = ('One SCF workflow was not successful: {}'.format(label))
                  self.ctx.warnings.append(message)
                  continue
  
              try:
                  outnodedict[label] = calc.outputs.output_scf_wc_para
              except KeyError:
                  message = (
                      'One SCF workflow failed, no scf output node: {}.'
                      ' I skip this one.'.format(label))
                  self.ctx.errors.append(message)
                  continue
  
              outpara = calc.outputs.output_scf_wc_para.get_dict()
  
              t_e = outpara.get('total_energy', 'failed')
              if not isinstance(t_e, float):
                  message = (
                      'Did not manage to extract float total energy from one '
                      'SCF workflow: {}'.format(label))
                  self.ctx.warnings.append(message)
                  continue
              e_u = outpara.get('total_energy_units', 'Htr')
              if e_u == 'Htr' or 'htr':
                  t_e = t_e * htr_to_eV
              t_energydict[label] = t_e
  
          if t_energydict:
              # Find a minimal value of Spiral and count it as 0
              minenergy = min(t_energydict.values())
  
              for key in six.iterkeys(t_energydict):
                  t_energydict[key] = t_energydict[key] - minenergy
  
          self.ctx.energydict = t_energydict
  
      def return_results(self):
          """
          Retrieve results of converge calculations
          """
  
          failed_labels = []
  
          for label in six.iterkeys(self.ctx.wf_dict['q_vectors']):
              if label not in six.iterkeys(self.ctx.energydict):
                  failed_labels.append(label)
  
          out = {'workflow_name': self.__class__.__name__,
                 'workflow_version': self._workflowversion,
                 # 'initial_structure': self.inputs.structure.uuid,
                 'energies': self.ctx.energydict,
                 'q_vectors': self.ctx.wf_dict['q_vectors'],
                 'failed_labels': failed_labels,
                 'energy_units': 'eV',
                 'info': self.ctx.info,
                 'warnings': self.ctx.warnings,
                 'errors': self.ctx.errors}
  
          # create link to workchain node
          out = save_output_node(Dict(dict=out))
          self.out('out', out)
  
          if not self.ctx.energydict:
              return self.exit_codes.ERROR_ALL_QVECTORS_FAILED
          elif failed_labels:
              return self.exit_codes.ERROR_SOME_QVECTORS_FAILED
  
      def control_end_wc(self, errormsg):
          """
          Controlled way to shutdown the workchain. will initialize the output nodes
          The shutdown of the workchain will has to be done afterwards
          """
          self.report(errormsg)
          self.ctx.errors.append(errormsg)
          self.return_results()
  
  
  @cf
  def save_output_node(out):
      """
      This calcfunction saves the out dict in the db
      """
      out_wc = out.clone()
      return out_wc
  