text: import glob
  import os
  import shutil
  
  from fireworks.core.firework import FiretaskBase, FWAction
  from fireworks.utilities.fw_utilities import explicit_serialize
  from pymatgen.analysis.diffusion.neb.io import (
      MVLCINEBSet,
      get_endpoint_dist,
      get_endpoints_from_index,
  )
  from pymatgen.core import Structure
  from pymatgen.io.vasp import Incar, Kpoints, Poscar, Potcar
  
  from atomate.utils.utils import get_logger
  
  """
  NEB workflow firetasks.
  """
  
  __author__ = "Hanmei Tang, Iek-Heng Chu"
  __email__ = "hat003@eng.ucsd.edu, ihchu@eng.ucsd.edu"
  
  
  VASP_NEB_OUTPUT_FILES = {"INCAR", "KPOINTS", "POTCAR", "vasprun.xml"}
  VASP_NEB_OUTPUT_SUB_FILES = {
      "CHG",
      "CHGCAR",
      "CONTCAR",
      "DOSCAR",
      "EIGENVAL",
      "IBZKPT",
      "PCDAT",
      "POSCAR",
      "PROCAR",
      "OSZICAR",
      "OUTCAR",
      "REPORT",
      "WAVECAR",
      "XDATCAR",
  }
  
  
  @explicit_serialize
  class TransferNEBTask(FiretaskBase):
      """
      This class transfers NEB outputs from current directory to destination directory. "label" is
      used to determine the step of calculation and hence the final path. The corresponding structure
      will be updated in fw_spec before files transferring.
  
      Required params:
          label (str): Type of calculation outputs being transferred, choose from "parent", "ep0",
              "ep1", "neb1", "neb2" and etc..
      Optional params:
          d_img (float): Distance between neighbouring images, used to determine the number of images
              if "IMAGES" not provided in user_incar_settings, in Angstrom.
      """
  
      required_params = ["label"]
      optional_params = ["d_img"]
  
      def run_task(self, fw_spec):
  
          label = self["label"]
          assert label in ["parent", "ep0", "ep1"] or "neb" in label, "Unknown label!"
  
          d_img = float(self.get("d_img", 0.7))  # Angstrom
          wf_name = fw_spec["wf_name"]
          src_dir = os.path.abspath(".")
          dest_dir = os.path.join(fw_spec["_fw_env"]["run_dest_root"], wf_name, label)
          shutil.copytree(src_dir, dest_dir)
  
          # Update fw_spec based on the type of calculations.
          if "neb" in label:
              # Update all relaxed images.
              subs = glob.glob("[0-2][0-9]")
              nimages = len(subs)
              concar_list = [f"{i:02}/CONTCAR" for i in range(nimages)[1:-1]]
              images = [Structure.from_file(contcar) for contcar in concar_list]
  
              # Update the two ending "images".
              images.insert(0, Structure.from_file("00/POSCAR"))
              images.append(Structure.from_file(f"{nimages - 1:02}/POSCAR"))
              images = [s.as_dict() for s in images]
              neb = fw_spec.get("neb")
              neb.append(images)
              update_spec = {
                  "neb": neb,
                  "_queueadapter": {
                      "nnodes": str(len(images) - 2),
                      "nodes": str(len(images) - 2),
                  },
              }
              # Use neb walltime if it is in fw_spec
              if fw_spec["neb_walltime"] is not None:
                  update_spec["_queueadapter"].update(
                      {"walltime": fw_spec.get("neb_walltime")}
                  )
  
          elif label in ["ep0", "ep1"]:
              # Update relaxed endpoint structures.
              file = glob.glob("CONTCAR*")[0]
              ep = Structure.from_file(file, False)  # One endpoint
  
              if fw_spec.get("incar_images"):  # "incar_images": pre-defined image number.
                  update_spec = {
                      label: ep.as_dict(),
                      "_queueadapter": {
                          "nnodes": fw_spec["incar_images"],
                          "nodes": fw_spec["incar_images"],
                      },
                  }
              else:
                  # Calculate number of images if "IMAGES" tag is not provided.
                  index = int(label[-1])
                  ep_1_dict = fw_spec.get(f"ep{1 - index}")  # Another endpoint
                  try:
                      ep_1 = Structure.from_dict(ep_1_dict)
                  except Exception:
                      ep_1 = ep_1_dict
  
                  max_dist = max(get_endpoint_dist(ep, ep_1))
                  nimages = round(max_dist / d_img) or 1
                  update_spec = {
                      label: ep,
                      "_queueadapter": {"nnodes": int(nimages), "nodes": int(nimages)},
                  }
              # Use neb walltime if it is in fw_spec
              if fw_spec["neb_walltime"] is not None:
                  update_spec["_queueadapter"].update(
                      {"walltime": fw_spec.get("neb_walltime")}
                  )
  
          else:  # label == "parent"
              f = glob.glob("CONTCAR*")[0]
              s = Structure.from_file(f, False)
              ep0, ep1 = get_endpoints_from_index(s, fw_spec["site_indices"])
  
              update_spec = {
                  "parent": s.as_dict(),
                  "ep0": ep0.as_dict(),
                  "ep1": ep1.as_dict(),
              }
  
          # Clear current directory.
          for d in os.listdir(src_dir):
              try:
                  os.remove(os.path.join(src_dir, d))
              except Exception:
                  shutil.rmtree(os.path.join(src_dir, d))
  
          return FWAction(update_spec=update_spec)
  
  
  @explicit_serialize
  class RunNEBVaspFake(FiretaskBase):
      """
      Vasp Emulator for NEB, which has a different file arrangement. Similar to RunVaspFake class.
  
      Required params:
          ref_dir (string): Path to reference vasp run directory with input files in the folder named
             'inputs' and output files in the folder named 'outputs'.
  
      Optional params:
          params_to_check (list): optional list of incar parameters to check.
      """
  
      required_params = ["ref_dir"]
      optional_params = ["params_to_check"]
  
      def run_task(self, fw_spec):
          self._get_params()
          self._verify_inputs()
          self._clear_inputs()
          self._generate_outputs()
  
      def _get_params(self):
          """Define some convenient variables."""
  
          self.user_dir = os.getcwd()
          self.ref_dir_input = os.path.join(self["ref_dir"], "inputs")
          self.ref_dir_output = os.path.join(self["ref_dir"], "outputs")
  
          user_sdir = glob.glob(os.path.join(os.getcwd(), "[0-9][0-9]"))
          ref_sdir_input = glob.glob(
              os.path.join(self["ref_dir"], "inputs", "[0-9][0-9]")
          )
          ref_sdir_output = glob.glob(
              os.path.join(self["ref_dir"], "outputs", "[0-9][0-9]")
          )
          user_sdir.sort()
          ref_sdir_input.sort()
          ref_sdir_output.sort()
  
          # Check sub-folders consistence.
          if len(user_sdir) != len(ref_sdir_input):
              raise ValueError(
                  "Sub-folder numbers are inconsistent! "
                  "Paths are:\n{}\n{}".format(self.user_dir, self.ref_dir_input)
              )
          self.user_sdir = user_sdir
          self.ref_sdir_input = ref_sdir_input
          self.ref_sdir_output = ref_sdir_output
  
      def _verify_inputs(self):
          """Validation of input files under user NEB directory."""
          user_incar = Incar.from_file(os.path.join(self.user_dir, "INCAR"))
          ref_incar = Incar.from_file(os.path.join(self.ref_dir_input, "INCAR"))
  
          # Check INCAR
          params_to_check = self.get("params_to_check", [])
          defaults = {"ICHAIN": 0, "LCLIMB": True}
          for p in params_to_check:
              if user_incar.get(p, defaults.get(p)) != ref_incar.get(p, defaults.get(p)):
                  raise ValueError(f"INCAR value of {p} is inconsistent!")
  
          # Check KPOINTS
          user_kpoints = Kpoints.from_file(os.path.join(self.user_dir, "KPOINTS"))
          ref_kpoints = Kpoints.from_file(os.path.join(self.ref_dir_input, "KPOINTS"))
          if (
              user_kpoints.style != ref_kpoints.style
              or user_kpoints.num_kpts != ref_kpoints.num_kpts
          ):
              raise ValueError(
                  "KPOINT files are inconsistent! "
                  "Paths are:\n{}\n{} with kpts = {} {}".format(
                      self.user_dir, self.ref_dir_input, user_kpoints, ref_kpoints
                  )
              )
  
          # Check POTCAR
          user_potcar = Potcar.from_file(os.path.join(self.user_dir, "POTCAR"))
          ref_potcar = Potcar.from_file(os.path.join(self.ref_dir_input, "POTCAR"))
          if user_potcar.symbols != ref_potcar.symbols:
              raise ValueError(
                  "POTCAR files are inconsistent! "
                  "Paths are:\n{}\n{}".format(self.user_dir, self.ref_dir_input)
              )
  
          # Check POSCARs
          for u, r in zip(self.user_sdir, self.ref_sdir_input):
              user_poscar = Poscar.from_file(os.path.join(u, "POSCAR"))
              ref_poscar = Poscar.from_file(os.path.join(r, "POSCAR"))
              if (
                  user_poscar.natoms != ref_poscar.natoms
                  or user_poscar.site_symbols != ref_poscar.site_symbols
              ):
                  raise ValueError(f"POSCAR files are inconsistent! Paths are:\n{u}\n{r}")
  
      def _clear_inputs(self):
          """
          Remove all input files from user NEB directory.
          """
          # Clear neb directory
          for x in VASP_NEB_OUTPUT_FILES:
              p = os.path.join(os.getcwd(), x)
              if os.path.exists(p):
                  os.remove(p)
  
          # Clear neb sub-directory
          for d in self.user_sdir:
              for x in VASP_NEB_OUTPUT_SUB_FILES:
                  p = os.path.join(d, x)
                  if os.path.exists(p):
                      os.remove(p)
  
      def _generate_outputs(self):
          """Copy valid outputs from a reference folder to user NEB directory."""
          # Copy NEB files.
          for file_name in os.listdir(self.ref_dir_output):
              full_file_name = os.path.join(self.ref_dir_output, file_name)
              if os.path.isfile(full_file_name):
                  shutil.copy(full_file_name, os.getcwd())
  
          # Copy NEB sub-files.
          for u_dir, r_dir in zip(self.user_sdir, self.ref_sdir_output):
              for file_name in os.listdir(r_dir):
                  full_file_name = os.path.join(r_dir, file_name)
                  if os.path.isfile(full_file_name):
                      shutil.copy(full_file_name, u_dir)
  
  
  @explicit_serialize
  class WriteNEBFromImages(FiretaskBase):
      """
      Generate CI-NEB input sets using given images and endpoints structures. The structures
      correspond to structures nested in subfolders ("00", "01", "02", etc.)
  
      Required parameters:
          neb_label (str): "1", "2", etc., labels the running sequence of NEB
      Optional parameters:
          user_incar_settings (dict): Additional INCAR settings.
          user_kpoints_settings (dict): Additional KPOINTS settings.
      """
  
      required_params = ["neb_label"]
      optional_params = ["user_incar_settings", "user_kpoints_settings"]
  
      def run_task(self, fw_spec):
          user_incar_settings = self.get("user_incar_settings", {})
          user_kpoints_settings = self.get("user_kpoints_settings", {})
          neb_label = self.get("neb_label")
          assert neb_label.isdigit() and int(neb_label) >= 1
          images = fw_spec["neb"][int(neb_label) - 1]
          try:
              images = [Structure.from_dict(i) for i in images]
          except Exception:
              images = images
          vis = MVLCINEBSet(
              images,
              user_incar_settings=user_incar_settings,
              user_kpoints_settings=user_kpoints_settings,
          )
          vis.write_input(".")
  
  
  @explicit_serialize
  class WriteNEBFromEndpoints(FiretaskBase):
      """
      Generate NEB input sets using endpoint structures, default writes CI-NEB input.
      MVLCINEBSet is the only vasp_input_set supported now.
  
      The number of images:
          1) search in "user_incar_settings";
          2) otherwise, calculated using "d_img".
      Required parameters:
          user_incar_settings (dict): additional INCAR settings.
  
      Optional parameters:
          user_kpoints_settings (dict): additional KPOINTS settings.
          sort_tol (float): Distance tolerance (in Angstrom) used to match the atomic indices between
              start and end structures. If it is set 0, then no sorting will be performed.
          d_img (float): distance in Angstrom, used in calculating number of images. Default 0.7
              Angstrom.
          interpolation_type (str): method to do image interpolation from two endpoints.
                              Choose from ["IDPP", "linear"], default "IDPP"
      """
  
      required_params = ["user_incar_settings"]
      optional_params = [
          "user_kpoints_settings",
          "sort_tol",
          "d_img",
          "interpolation_type",
      ]
  
      def run_task(self, fw_spec):
          user_incar_settings = self["user_incar_settings"]
          interpolation_type = self.get("interpolation_type", "IDPP")
          idpp_species = fw_spec.get("idpp_species")
          user_kpoints_settings = self.get("user_kpoints_settings")
          try:
              ep0 = Structure.from_dict(fw_spec["ep0"])
              ep1 = Structure.from_dict(fw_spec["ep1"])
          except Exception:
              ep0 = fw_spec["ep0"]
              ep1 = fw_spec["ep1"]
  
          # Get number of images.
          nimages = user_incar_settings.get("IMAGES", self._get_nimages(ep0, ep1))
          if interpolation_type == "IDPP":
              from pymatgen.analysis.diffusion.neb.pathfinder import IDPPSolver
  
              obj = IDPPSolver.from_endpoints([ep0, ep1], nimages=nimages)
              images = obj.run(species=idpp_species)
              images_dic_list = [image.as_dict() for image in images]
          elif interpolation_type == "linear":
              images = self._get_images_by_linear_interp(nimages, ep0, ep1)
              images_dic_list = [i.as_dict() for i in images]
          else:
              raise ValueError(
                  "The interpolation method must either be 'linear' or 'IDPP'!"
              )
  
          write = WriteNEBFromImages(
              neb_label="1",
              user_incar_settings=user_incar_settings,
              user_kpoints_settings=user_kpoints_settings,
          )
          fw_spec["neb"] = [images_dic_list]
          write.run_task(fw_spec=fw_spec)
  
      def _get_nimages(self, ep0, ep1):
          """
          Calculate the number of images using "d_img", which can be
          overwritten in a optional_params list.
  
          Returns:
              nimages (int): number of images.
          """
          # Check endpoints consistence.
          assert ep0.atomic_numbers == ep1.atomic_numbers, "Endpoints are inconsistent!"
          d_img = self.get("d_img") or 0.7
  
          # Assume dilute diffusion
          max_dist = 0
          for s_0, s_1 in zip(ep0, ep1):
              site_dist = ep0.lattice.get_distance_and_image(
                  s_0.frac_coords, s_1.frac_coords
              )[0]
              if site_dist > max_dist:
                  max_dist = site_dist
  
          # Number of images must more than one.
          nimages = round(max_dist / d_img) or 1
  
          return nimages
  
      def _get_images_by_linear_interp(self, nimages, ep0, ep1):
          logger = get_logger(__name__)
          sort_tol = self.get("sort_tol", 0.0)
  
          try:
              images = ep0.interpolate(ep1, nimages=nimages + 1, autosort_tol=sort_tol)
          except Exception as e:
              if "Unable to reliably match structures " in str(e):
                  logger.warning(
                      "Auto sorting is turned off because it is unable to match the "
                      "end-point structures!"
                  )
                  images = ep0.interpolate(ep1, nimages=nimages + 1, autosort_tol=0)
              else:
                  raise e
  
          return images
  