# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.

from __future__ import print_function
from itertools import islice
import numbers
import numpy as np
import os
import pandas as pd
import posixpath
import scipy.constants
from pyiron_atomistics.atomistics.structure.structurestorage import StructureStorage
from pyiron_base import (
    state,
    GenericParameters,
    GenericJob,
    Executable,
    FlattenedStorage,
)
from pyiron_atomistics import ase_to_pyiron, Atoms
from pyiron_contrib.atomistics.ml.potentialfit import PotentialFit
from pyiron_contrib.atomistics.mlip.cfgs import savecfgs, loadcfgs, Cfg
from pyiron_contrib.atomistics.mlip.potential import MtpPotential

__author__ = "Jan Janssen"
__copyright__ = (
    "Copyright 2020, Max-Planck-Institut für Eisenforschung GmbH - "
    "Computational Materials Design (CM) Department"
)
__version__ = "1.0"
__maintainer__ = "Jan Janssen"
__email__ = "janssen@mpie.de"
__status__ = "development"
__date__ = "Sep 1, 2017"


gpa_to_ev_ang = (
    1e22 / scipy.constants.physical_constants["joule-electron volt relationship"][0]
)


class Mlip(GenericJob, PotentialFit):
    def __init__(self, project, job_name):
        super(Mlip, self).__init__(project, job_name)
        self.__version__ = "0.1.0"
        self.__name__ = "Mlip"
        self._executable_activate()
        self._job_dict = {}
        self.input = MlipParameter()
        self._command_line = CommandLine()
        self._potential = MtpPotential()

    def _executable_activate(self, enforce=False):
        if self._executable is None or enforce:
            if len(self.__module__.split(".")) > 1:
                self._executable = Executable(
                    codename=self.__name__,
                    module=self.__module__.split(".")[-2],
                    path_binary_codes=state.settings.resource_paths,
                )
            else:
                self._executable = Executable(
                    codename=self.__name__,
                    path_binary_codes=state.settings.resource_paths,
                )

    @property
    def calculation_dataframe(self):
        (
            job_id_lst,
            time_step_start_lst,
            time_step_end_lst,
            time_step_delta_lst,
        ) = self._job_dict_lst(self._job_dict)
        return pd.DataFrame(
            {
                "Job id": job_id_lst,
                "Start config": time_step_start_lst,
                "End config": time_step_end_lst,
                "Step": time_step_delta_lst,
            }
        )

    @property
    def potential_files(self):
        pot = os.path.join(self.working_directory, "Trained.mtp_")
        states = os.path.join(self.working_directory, "state.mvs")
        if os.path.exists(pot) and os.path.exists(states):
            return [pot, states]

    def _get_elements(self):
        """
        Return elements in training in insertion order, i.e. elements seen earlier get lower indices.
        """
        elements = []
        for job_id in self._job_dict:
            j = self.project.inspect(job_id)
            if j["NAME"] == "TrainingContainer":
                candidates = j.to_object().get_elements()
            else:
                candidates = j["input/structure/species"]
            for e in candidates:
                if e not in elements:
                    elements.append(e)
        return elements

    def potential_dataframe(self, elements=None):
        """
        :class:`pandas.DataFrame`: potential dataframe for lammps jobs

        .. attention::

            The `elements` argument must be given for in any case for non-unary alloys, as the scrapping code from the
            training data is currently broken!

        Args:
            elements (list of str, optional): name of elements in this potential in the order of their indices in pyiron
                                              structures, if not given use elements seen in training data
        """
        if elements is None:
            elements = self.input["species"]
            if elements is None:
                elements = self._get_elements()  # AAAH

        if self.status.finished:
            return pd.DataFrame(
                {
                    "Name": ["".join(elements)],
                    "Filename": [self.potential_files],
                    "Model": ["Custom"],
                    "Species": [elements],
                    "Config": [["pair_style mlip mlip.ini\n", "pair_coeff * *\n"]],
                }
            )
        else:
            raise ValueError(
                f"Potential only available after job is finished, not {self.status}!"
            )

    @property
    def potential(self):
        if self.status.finished and self._potential is not None:
            return self._potential
        else:
            raise ValueError("potential only available on successfully finished jobs")

    def set_input_to_read_only(self):
        """
        This function enforces read-only mode for the input classes, but it has to be implement in the individual
        classes.
        """
        self.input.read_only = True
        self._command_line.read_only = True

    def add_job_to_fitting(
        self, job_id, time_step_start=0, time_step_end=-1, time_step_delta=10
    ):
        if time_step_end == -1:
            time_step_end = (
                np.shape(self.project.inspect(int(job_id))["output/generic/cells"])[0]
                - 1
            )
        self._job_dict[job_id] = {
            "time_step_start": time_step_start,
            "time_step_end": time_step_end,
            "time_step_delta": time_step_delta,
        }

    def write_input(self):
        restart_file_lst = [
            os.path.basename(f) for f in self._restart_file_list
        ] + list(self._restart_file_dict.values())
        if "testing.cfg" not in restart_file_lst:
            species = self._write_test_set(
                file_name="testing.cfg", cwd=self.working_directory
            )
            species_count = len(species)
            self.input["species"] = species
        elif self.input["filepath"] != "auto":
            species_count = 0
        else:
            raise ValueError("speciescount not set")
        if "training.cfg" not in restart_file_lst:
            self._write_training_set(
                file_name="training.cfg", cwd=self.working_directory
            )
        self._copy_potential(
            species_count=species_count,
            file_name="start.mtp",
            cwd=self.working_directory,
        )
        if self.version != "1.0.0":
            self._command_line.activate_postprocessing()
        self._command_line[0] = self._command_line[0].replace(
            "energy_auto", str(self.input["energy-weight"])
        )
        self._command_line[0] = self._command_line[0].replace(
            "force_auto", str(self.input["force-weight"])
        )
        self._command_line[0] = self._command_line[0].replace(
            "stress_auto", str(self.input["stress-weight"])
        )
        self._command_line[0] = self._command_line[0].replace(
            "iteration_auto", str(int(self.input["iteration"]))
        )
        self._command_line.write_file(file_name="mlip.sh", cwd=self.working_directory)

        species = np.array(species)
        input_store = StructureStorage()
        input_store.add_array("energy", dtype=np.float64, shape=(), per="chunk")
        input_store.add_array("forces", dtype=np.float64, shape=(3,), per="element")
        input_store.add_array("stress", dtype=np.float64, shape=(6,), per="chunk")
        for cfg in loadcfgs(os.path.join(self.working_directory, "training.cfg")):
            struct = Atoms(
                symbols=species[np.cast[np.int64](cfg.types)],
                positions=cfg.pos,
                cell=cfg.lat,
                pbc=[True] * 3,
            )  # HACK for pbc
            input_store.add_structure(
                struct,
                identifier=cfg.desc,
                energy=cfg.energy,
                forces=cfg.forces,
                stress=cfg.stresses,
            )

        with self.project_hdf5.open("input") as hdf5_input:
            input_store.to_hdf(hdf=hdf5_input, group_name="training_data")

    def collect_logfiles(self):
        pass

    def collect_output(self):
        file_name = os.path.join(self.working_directory, "diff.cfg")
        if os.path.exists(file_name):
            _, _, _, _, _, _, _, job_id_diff_lst, timestep_diff_lst = read_cgfs(
                file_name
            )
        else:
            job_id_diff_lst, timestep_diff_lst = [], []
        file_name = os.path.join(self.working_directory, "selected.cfg")
        if os.path.exists(file_name):
            (
                _,
                _,
                _,
                _,
                _,
                _,
                _,
                job_id_new_training_lst,
                timestep_new_training_lst,
            ) = read_cgfs(file_name)
        else:
            job_id_new_training_lst, timestep_new_training_lst = [], []
        file_name = os.path.join(self.working_directory, "grades.cfg")
        if os.path.exists(file_name):
            (
                _,
                _,
                _,
                _,
                _,
                _,
                grades_lst,
                job_id_grades_lst,
                timestep_grades_lst,
            ) = read_cgfs(file_name)
        else:
            grades_lst, job_id_grades_lst, timestep_grades_lst = [], [], []
        try:
            self._potential.load(os.path.join(self.working_directory, "Trained.mtp_"))
        except:
            self.logger.warn(
                "Failed to parse potential file! job.potential will not be available."
            )
            self._potential = None

        training_store = FlattenedStorage()
        training_store.add_array("energy", dtype=np.float64, shape=(), per="chunk")
        training_store.add_array("forces", dtype=np.float64, shape=(3,), per="element")
        training_store.add_array("stress", dtype=np.float64, shape=(6,), per="chunk")
        for cfg in loadcfgs(os.path.join(self.working_directory, "training_efs.cfg")):
            training_store.add_chunk(
                len(cfg.pos),
                identifier=cfg.desc,
                energy=cfg.energy,
                forces=cfg.forces,
                stress=cfg.stresses,
            )

        testing_store = FlattenedStorage()
        testing_store.add_array("energy", dtype=np.float64, shape=(), per="chunk")
        testing_store.add_array("forces", dtype=np.float64, shape=(3,), per="element")
        testing_store.add_array("stress", dtype=np.float64, shape=(6,), per="chunk")
        for cfg in loadcfgs(os.path.join(self.working_directory, "testing_efs.cfg")):
            testing_store.add_chunk(
                len(cfg.pos),
                identifier=cfg.desc,
                energy=cfg.energy,
                forces=cfg.forces,
                stress=cfg.stresses,
            )

        with self.project_hdf5.open("output") as hdf5_output:
            hdf5_output["grades"] = grades_lst
            hdf5_output["job_id"] = job_id_grades_lst
            hdf5_output["timestep"] = timestep_grades_lst
            hdf5_output["job_id_diff"] = job_id_diff_lst
            hdf5_output["timestep_diff"] = timestep_diff_lst
            hdf5_output["job_id_new"] = job_id_new_training_lst
            hdf5_output["timestep_new"] = timestep_new_training_lst
            if self._potential is not None:
                self._potential.to_hdf(hdf=hdf5_output)
            training_store.to_hdf(hdf=hdf5_output, group_name="training_efs")
            testing_store.to_hdf(hdf=hdf5_output, group_name="testing_efs")

    def get_structure(self, iteration_step=-1):
        job = self.project.load(self["output/job_id_diff"][iteration_step])
        return job.get_structure(self["output/timestep_diff"][iteration_step])

    def to_hdf(self, hdf=None, group_name=None):
        super(Mlip, self).to_hdf(hdf=hdf, group_name=group_name)
        self.input.to_hdf(self._hdf5)
        with self._hdf5.open("input") as hdf_input:
            (
                job_id_lst,
                time_step_start_lst,
                time_step_end_lst,
                time_step_delta_lst,
            ) = self._job_dict_lst(self._job_dict)
            hdf_input["job_id_list"] = job_id_lst
            hdf_input["time_step_start_lst"] = time_step_start_lst
            hdf_input["time_step_end_lst"] = time_step_end_lst
            hdf_input["time_step_delta_lst"] = time_step_delta_lst

    def from_hdf(self, hdf=None, group_name=None):
        super(Mlip, self).from_hdf(hdf=hdf, group_name=group_name)
        self.input.from_hdf(self._hdf5)
        with self._hdf5.open("input") as hdf_input:
            for job_id, start, end, delta in zip(
                hdf_input["job_id_list"],
                hdf_input["time_step_start_lst"],
                hdf_input["time_step_end_lst"],
                hdf_input["time_step_delta_lst"],
            ):
                self.add_job_to_fitting(
                    job_id=job_id,
                    time_step_start=start,
                    time_step_end=end,
                    time_step_delta=delta,
                )
        if self.status.finished:
            with self._hdf5.open("output") as hdf_output:
                if "potential" in hdf_output.list_groups():
                    self._potential.from_hdf(hdf=hdf_output)

    def get_suggested_number_of_configuration(
        self, species_count=None, multiplication_factor=2.0
    ):
        if self.input["filepath"] == "auto":
            potential = self.input["potential"]
            if isinstance(potential, numbers.Integral):
                potential = f"{potential:02}g.mtp"
            source = self._find_potential(potential)
        else:
            source = self.input["filepath"]
        with open(source) as f:
            lines = f.readlines()
        radial_basis_size, radial_funcs_count, alpha_scalar_moments = 0.0, 0.0, 0.0
        for line in lines:
            if "radial_basis_size" in line:
                radial_basis_size = int(line.split()[2])
            if "radial_funcs_count" in line:
                radial_funcs_count = int(line.split()[2])
            if "species_count" in line and species_count is None:
                species_count = int(line.split()[2])
            if "alpha_scalar_moments" in line:
                alpha_scalar_moments = int(line.split()[2])
        return int(
            multiplication_factor
            * (
                radial_basis_size * radial_funcs_count * species_count**2
                + alpha_scalar_moments
            )
        )

    @staticmethod
    def _update_species_count(lines, species_count):
        ind = 0
        for ind, line in enumerate(lines):
            if "species_count" in line:
                break
        species_line = lines[ind].split()
        species_line[-1] = str(species_count)
        lines[ind] = " ".join(species_line) + "\n"
        return lines

    @staticmethod
    def _update_min_max_distance(lines, min_distance, max_distance):
        min_dis_updated, max_dis_updated = False, False
        for ind, line in enumerate(lines):
            if "min_dist" in line and min_distance is not None:
                min_dist_line = lines[ind].split()
                min_dist_line[2] = str(min_distance)
                lines[ind] = " ".join(min_dist_line) + "\n"
                min_dis_updated = True
            elif "max_dist" in line and max_distance is not None:
                max_dist_line = lines[ind].split()
                max_dist_line[2] = str(max_distance)
                lines[ind] = " ".join(max_dist_line) + "\n"
                max_dis_updated = True
            elif min_dis_updated and max_dis_updated:
                break
        return lines

    def _copy_potential(self, species_count, file_name, cwd=None):
        if cwd is not None:
            file_name = posixpath.join(cwd, file_name)
        if self.input["filepath"] == "auto":
            potential = self.input["potential"]
            if isinstance(potential, numbers.Integral):
                potential = f"{potential:02}g.mtp"
            source = self._find_potential(potential)
        else:
            source = self.input["filepath"]
        with open(source) as f:
            lines = f.readlines()
        if self.input["filepath"] == "auto":
            lines_modified = self._update_species_count(lines, species_count)
            lines_modified = self._update_min_max_distance(
                lines=lines_modified,
                min_distance=self.input["min_dist"],
                max_distance=self.input["max_dist"],
            )
        else:
            lines_modified = lines
        with open(file_name, "w") as f:
            f.writelines(lines_modified)

    @staticmethod
    def stress_tensor_components(stresses):
        return (
            np.array(
                [
                    stresses[0][0],
                    stresses[1][1],
                    stresses[2][2],
                    stresses[1][2],
                    stresses[0][1],
                    stresses[0][2],
                ]
            )
            * gpa_to_ev_ang
        )

    def _write_configurations(
        self, file_name="training.cfg", cwd=None, respect_step=True
    ):
        if cwd is not None:
            file_name = posixpath.join(cwd, file_name)
        (
            indices_lst,
            position_lst,
            forces_lst,
            cell_lst,
            energy_lst,
            track_lst,
            stress_lst,
        ) = ([], [], [], [], [], [], [])
        all_species = set()
        for job_id, value in self._job_dict.items():
            ham = self.project.inspect(job_id)
            if respect_step:
                start = value["time_step_start"]
                end = value["time_step_end"] + 1
                delta = value["time_step_delta"]
            else:
                start, end, delta = 0, None, 1
            time_step = start
            # HACK: until the training container has a proper HDF5 interface
            if ham.__name__ == "TrainingContainer":
                job = ham.to_object()
                all_species.update(job.get_elements())
                symbol_map = {s: i for i, s in enumerate(sorted(set(all_species)))}
                if end is None:
                    end = job.number_of_structures
                for time_step in range(start, end, delta):
                    symbols = job._container.get_array("symbols", time_step)
                    indices_lst.append([symbol_map[s] for s in symbols])
                    position_lst.append(
                        job._container.get_array("positions", time_step)
                    )
                    forces_lst.append(job._container.get_array("forces", time_step))
                    cell_lst.append(job._container.get_array("cell", time_step))
                    energy_lst.append(job._container.get_array("energy", time_step))
                    if job._container.has_array("stress"):
                        volume = np.abs(np.linalg.det(cell_lst[-1]))
                        stress_lst.append(
                            job._container.get_array("stress", time_step) * volume
                        )
                        if np.isnan(stress_lst[-1]).any():
                            stress_lst[-1] = None
                    track_lst.append(str(ham.job_id) + "_" + str(time_step))
                continue
            original_dict = {
                el: ind for ind, el in enumerate(sorted(ham["input/structure/species"]))
            }
            species_dict = {
                ind: original_dict[el]
                for ind, el in enumerate(ham["input/structure/species"])
            }
            # HACK: this does not preserve element order, but the old code is broken for structures with different
            # species anyway
            all_species.update(ham["input/structure/species"])
            if ham.__name__ in [
                "Vasp",
                "ThermoIntDftEam",
                "ThermoIntDftMtp",
                "ThermoIntVasp",
            ]:
                if len(ham["output/outcar/stresses"]) != 0:
                    for position, forces, cell, energy, stresses, volume in zip(
                        ham["output/generic/positions"][start:end:delta],
                        ham["output/generic/forces"][start:end:delta],
                        ham["output/generic/cells"][start:end:delta],
                        ham["output/generic/dft/energy_free"][start:end:delta],
                        ham["output/outcar/stresses"][start:end:delta],
                        ham["output/generic/volume"][start:end:delta],
                    ):
                        indices_lst.append(
                            [species_dict[el] for el in ham["input/structure/indices"]]
                        )
                        position_lst.append(position)
                        forces_lst.append(forces)
                        cell_lst.append(cell)
                        energy_lst.append(energy)
                        stress_lst.append(stresses * volume / gpa_to_ev_ang)
                        track_lst.append(str(ham.job_id) + "_" + str(time_step))
                        time_step += delta
                else:
                    for position, forces, cell, energy in zip(
                        ham["output/generic/positions"][start:end:delta],
                        ham["output/generic/forces"][start:end:delta],
                        ham["output/generic/cells"][start:end:delta],
                        ham["output/generic/dft/energy_free"][start:end:delta],
                    ):
                        indices_lst.append(
                            [species_dict[el] for el in ham["input/structure/indices"]]
                        )
                        position_lst.append(position)
                        forces_lst.append(forces)
                        cell_lst.append(cell)
                        energy_lst.append(energy)
                        track_lst.append(str(ham.job_id) + "_" + str(time_step))
                        time_step += delta
            elif ham.__name__ in ["Lammps", "LammpsInt2", "LammpsMlip"]:
                for position, forces, cell, energy, stresses, volume in zip(
                    ham["output/generic/positions"][start:end:delta],
                    ham["output/generic/forces"][start:end:delta],
                    ham["output/generic/cells"][start:end:delta],
                    ham["output/generic/energy_pot"][start:end:delta],
                    ham["output/generic/pressures"][start:end:delta],
                    ham["output/generic/volume"][start:end:delta],
                ):
                    indices_lst.append(
                        [species_dict[el] for el in ham["input/structure/indices"]]
                    )
                    position_lst.append(position)
                    forces_lst.append(forces)
                    cell_lst.append(cell)
                    energy_lst.append(energy)
                    stress_lst.append(self.stress_tensor_components(stresses * volume))
                    track_lst.append(str(ham.job_id) + "_" + str(time_step))
                    time_step += delta
            else:
                for position, forces, cell, energy, stresses, volume in zip(
                    ham["output/generic/positions"][start:end:delta],
                    ham["output/generic/forces"][start:end:delta],
                    ham["output/generic/cells"][start:end:delta],
                    ham["output/generic/energy_pot"][start:end:delta],
                    ham["output/generic/pressures"][start:end:delta],
                    ham["output/generic/volume"][start:end:delta],
                ):
                    indices_lst.append(
                        [species_dict[el] for el in ham["input/structure/indices"]]
                    )
                    position_lst.append(position)
                    forces_lst.append(forces)
                    cell_lst.append(cell)
                    energy_lst.append(energy)
                    stress_lst.append(self.stress_tensor_components(stresses * volume))
                    track_lst.append(str(ham.job_id) + "_" + str(time_step))
                    time_step += delta
        write_cfg(
            file_name=file_name,
            indices_lst=indices_lst,
            position_lst=position_lst,
            cell_lst=cell_lst,
            forces_lst=forces_lst,
            energy_lst=energy_lst,
            track_lst=track_lst,
            stress_lst=stress_lst,
        )
        return list(all_species)

    def _write_test_set(self, file_name="testing.cfg", cwd=None):
        return self._write_configurations(
            file_name=file_name, cwd=cwd, respect_step=False
        )

    def _write_training_set(self, file_name="training.cfg", cwd=None):
        self._write_configurations(file_name=file_name, cwd=cwd, respect_step=True)

    @staticmethod
    def _job_dict_lst(job_dict):
        job_id_lst, time_step_start_lst, time_step_end_lst, time_step_delta_lst = (
            [],
            [],
            [],
            [],
        )
        for key, value in job_dict.items():
            job_id_lst.append(key)
            time_step_start_lst.append(value["time_step_start"])
            time_step_end_lst.append(value["time_step_end"])
            time_step_delta_lst.append(value["time_step_delta"])
        return job_id_lst, time_step_start_lst, time_step_end_lst, time_step_delta_lst

    @staticmethod
    def _find_potential(potential_name):
        for resource_path in state.settings.resource_paths:
            if os.path.exists(
                os.path.join(resource_path, "mlip", "potentials", "templates")
            ):
                resource_path = os.path.join(
                    resource_path, "mlip", "potentials", "templates"
                )
            if "potentials" in resource_path and os.path.exists(
                os.path.join(resource_path, potential_name)
            ):
                return os.path.join(resource_path, potential_name)
        raise ValueError("Potential not found!")

    # PotentialFit Implementation
    def _add_training_data(self, container):
        self.add_job_to_fitting(container.id, 0, container.number_of_structures - 1, 1)

    def _get_training_data(self):
        # TODO/BUG: only works after input is written for now, instead this should go over _job_
        return self["input/training_data"].to_object()

    def _get_predicted_data(self):
        return self["output/training_efs"].to_object()

    def get_lammps_potential(self):
        return self.potential_dataframe()


class MlipParameter(GenericParameters):
    def __init__(self, separator_char=" ", comment_char="#", table_name="mlip_inp"):
        super(MlipParameter, self).__init__(
            separator_char=separator_char,
            comment_char=comment_char,
            table_name=table_name,
        )

    def load_default(self, file_content=None):
        if file_content is None:
            file_content = """\
potential 16g.mtp
filepath auto
energy-weight 1
force-weight 0.01
stress-weight 0.001
iteration 1000
min_dist 2.0
max_dist 5.0
"""
        self.load_string(file_content)


class CommandLine(GenericParameters):
    def __init__(self, input_file_name=None, **qwargs):
        super(CommandLine, self).__init__(
            input_file_name=input_file_name,
            table_name="cmd",
            comment_char="#",
            val_only=True,
        )

    def load_default(self, file_content=None):
        if file_content is None:
            file_content = """\
$MLP_COMMAND_PARALLEL train --energy-weight=energy_auto --force-weight=force_auto --stress-weight=stress_auto --max-iter=iteration_auto start.mtp training.cfg > training.log
"""
        self.load_string(file_content)

    def activate_postprocessing(self, file_content=None):
        if file_content is None:
            file_content = """\
$MLP_COMMAND_PARALLEL train --energy-weight=energy_auto --force-weight=force_auto --stress-weight=stress_auto --max-iter=iteration_auto start.mtp training.cfg > training.log
$MLP_COMMAND_SERIAL calc-grade Trained.mtp_ training.cfg testing.cfg grades.cfg --mvs-filename=state.mvs > grading.log
$MLP_COMMAND_SERIAL calc-errors Trained.mtp_ training.cfg > training.errors
$MLP_COMMAND_SERIAL calc-errors Trained.mtp_ testing.cfg > testing.errors
$MLP_COMMAND_SERIAL calc-efs Trained.mtp_ training.cfg training_efs.cfg
$MLP_COMMAND_SERIAL calc-efs Trained.mtp_ testing.cfg testing_efs.cfg
$MLP_COMMAND_SERIAL select-add Trained.mtp_ training.cfg testing.cfg diff.cfg > select.log
cp training.cfg training_new.cfg 
cat diff.cfg >> training_new.cfg
"""
        self.load_string(file_content)


def write_cfg(
    file_name,
    indices_lst,
    position_lst,
    cell_lst,
    forces_lst=None,
    energy_lst=None,
    track_lst=None,
    stress_lst=None,
):
    if stress_lst is None or len(stress_lst) == 0:
        stress_lst = [None] * len(position_lst)
    if forces_lst is None or len(forces_lst) == 0:
        forces_lst = [None] * len(position_lst)
    if track_lst is None or len(track_lst) == 0:
        track_lst = [None] * len(position_lst)
    if energy_lst is None or len(energy_lst) == 0:
        energy_lst = [None] * len(position_lst)
    cfg_lst = []
    for indices, position, forces, cell, energy, stress, track_str in zip(
        indices_lst,
        position_lst,
        forces_lst,
        cell_lst,
        energy_lst,
        stress_lst,
        track_lst,
    ):
        cfg_object = Cfg()
        cfg_object.pos = position
        cfg_object.lat = cell
        cfg_object.types = indices
        cfg_object.energy = energy
        cfg_object.forces = forces
        cfg_object.stresses = stress
        if track_str is not None:
            cfg_object.desc = "pyiron\t" + track_str
        cfg_lst.append(cfg_object)
    savecfgs(filename=file_name, cfgs=cfg_lst, desc=None)


def read_cgfs(file_name):
    cgfs_lst = loadcfgs(filename=file_name, max_cfgs=None)
    cell, positions, forces, stress, energy, indicies, grades, jobids, timesteps = (
        [],
        [],
        [],
        [],
        [],
        [],
        [],
        [],
        [],
    )
    for cgf in cgfs_lst:
        if cgf.pos is not None:
            positions.append(cgf.pos)
        if cgf.lat is not None:
            cell.append(cgf.lat)
        if cgf.types is not None:
            indicies.append(cgf.types)
        if cgf.energy is not None:
            energy.append(cgf.energy)
        if cgf.forces is not None:
            forces.append(cgf.forces)
        if cgf.stresses is not None:
            stress.append(
                [
                    [cgf.stresses[0], cgf.stresses[5], cgf.stresses[4]],
                    [cgf.stresses[5], cgf.stresses[1], cgf.stresses[3]],
                    [cgf.stresses[4], cgf.stresses[3], cgf.stresses[2]],
                ]
            )
        if cgf.grade is not None:
            grades.append(cgf.grade)
        if cgf.desc is not None:
            job_id, timestep = cgf.desc.split("_")
            jobids.append(int(job_id))
            timesteps.append(int(timestep))
    return [
        np.array(cell),
        np.array(positions),
        np.array(forces),
        np.array(stress),
        np.array(energy),
        np.array(indicies),
        np.array(grades),
        np.array(jobids),
        np.array(timesteps),
    ]
