# Source code for pymatgen.alchemy.transmuters

    # coding: utf-8
    # Copyright (c) Pymatgen Development Team.
    # Distributed under the terms of the MIT License.

    """
    This module implements various transmuter classes.
    Transmuters are essentially classes that generate TransformedStructures from
    various data sources. They enable the high-throughput generation of new
    structures and input files.

    It also includes the helper function, batch_write_vasp_input to generate an
    entire directory of vasp input files for running.
    """

    __author__ = "Shyue Ping Ong, Will Richards"
    __copyright__ = "Copyright 2012, The Materials Project"
    __version__ = "0.1"
    __maintainer__ = "Shyue Ping Ong"
    __email__ = "shyuep@gmail.com"
    __date__ = "Mar 4, 2012"

    import os
    import re

    from multiprocessing import Pool
    from pymatgen.alchemy.materials import TransformedStructure
    from pymatgen.io.vasp.sets import MPRelaxSet

    class StandardTransmuter:
        """
        An example of a Transmuter object, which performs a sequence of
        transformations on many structures to generate TransformedStructures.

        .. attribute: transformed_structures

            List of all transformed structures.
        """

        def __init__(self, transformed_structures, transformations=None,
                     extend_collection=0, ncores=None):
            """
            Initializes a transmuter from an initial list of
            :class:`pymatgen.alchemy.materials.TransformedStructure`.

            Args:
                transformed_structures ([TransformedStructure]): Input transformed
                    structures
                transformations ([Transformations]): New transformations to be
                    applied to all structures.
                extend_collection (int): Whether to use more than one output
                    structure from one-to-many transformations. extend_collection
                    can be an int, which determines the maximum branching for each
                    transformation.
                ncores (int): Number of cores to use for applying transformations.
                    Uses multiprocessing.Pool. Default is None, which implies
                    serial.
            """

            self.transformed_structures = transformed_structures
            self.ncores = ncores
            if transformations is not None:
                for trans in transformations:
                    self.append_transformation(trans,
                                               extend_collection=extend_collection)

        def __getitem__(self, index):
            return self.transformed_structures[index]

        def __getattr__(self, name):
            return [getattr(x, name) for x in self.transformed_structures]

        def undo_last_change(self):
            """
            Undo the last transformation in the TransformedStructure.

            Raises:
                IndexError if already at the oldest change.
            """
            for x in self.transformed_structures:
                x.undo_last_change()

        def redo_next_change(self):
            """
            Redo the last undone transformation in the TransformedStructure.

            Raises:
                IndexError if already at the latest change.
            """
            for x in self.transformed_structures:
                x.redo_next_change()

        def __len__(self):
            return len(self.transformed_structures)

        def append_transformation(self, transformation, extend_collection=False,
                                  clear_redo=True):
            """
            Appends a transformation to all TransformedStructures.

            Args:
                transformation: Transformation to append
                extend_collection: Whether to use more than one output structure
                    from one-to-many transformations. extend_collection can be a
                    number, which determines the maximum branching for each
                    transformation.
                clear_redo (bool): Whether to clear the redo list. By default,
                    this is True, meaning any appends clears the history of
                    undoing. However, when using append_transformation to do a
                    redo, the redo list should not be cleared to allow multiple
                    redos.

            Returns:
                List of booleans corresponding to initial transformed structures
                each boolean describes whether the transformation altered the
                structure
            """
            if self.ncores and transformation.use_multiprocessing:
                p = Pool(self.ncores)
                # need to condense arguments into single tuple to use map
                z = map(
                    lambda x: (x, transformation, extend_collection, clear_redo),
                    self.transformed_structures)
                new_tstructs = p.map(_apply_transformation, z, 1)
                self.transformed_structures = []
                for ts in new_tstructs:
                    self.transformed_structures.extend(ts)
            else:
                new_structures = []
                for x in self.transformed_structures:
                    new = x.append_transformation(transformation,
                                                  extend_collection,
                                                  clear_redo=clear_redo)
                    if new is not None:
                        new_structures.extend(new)
                self.transformed_structures.extend(new_structures)

        def extend_transformations(self, transformations):
            """
            Extends a sequence of transformations to the TransformedStructure.

            Args:
                transformations: Sequence of Transformations
            """
            for t in transformations:
                self.append_transformation(t)

        def apply_filter(self, structure_filter):
            """
            Applies a structure_filter to the list of TransformedStructures
            in the transmuter.

            Args:
                structure_filter: StructureFilter to apply.
            """

            def test_transformed_structure(ts):
                return structure_filter.test(ts.final_structure)

            self.transformed_structures = list(filter(test_transformed_structure,
                                                      self.transformed_structures))
            for ts in self.transformed_structures:
                ts.append_filter(structure_filter)

        def write_vasp_input(self, **kwargs):
            r"""
            Batch write vasp input for a sequence of transformed structures to
            output_dir, following the format output_dir/{formula}_{number}.

            Args:
                \\*\\*kwargs: All kwargs supported by batch_write_vasp_input.
            """
            batch_write_vasp_input(self.transformed_structures, **kwargs)

        def set_parameter(self, key, value):
            """
            Add parameters to the transmuter. Additional parameters are stored in
            the as_dict() output.

            Args:
                key: The key for the parameter.
                value: The value for the parameter.
            """
            for x in self.transformed_structures:
                x.other_parameters[key] = value

        def add_tags(self, tags):
            """
            Add tags for the structures generated by the transmuter.

            Args:
                tags: A sequence of tags. Note that this should be a sequence of
                    strings, e.g., ["My awesome structures", "Project X"].
            """
            self.set_parameter("tags", tags)

        def __str__(self):
            output = ["Current structures", "------------"]
            for x in self.transformed_structures:
                output.append(str(x.final_structure))
            return "\n".join(output)

        def append_transformed_structures(self, tstructs_or_transmuter):
            """
            Method is overloaded to accept either a list of transformed structures
            or transmuter, it which case it appends the second transmuter"s
            structures.

            Args:
                tstructs_or_transmuter: A list of transformed structures or a
                    transmuter.
            """
            if isinstance(tstructs_or_transmuter, self.__class__):
                self.transformed_structures.extend(tstructs_or_transmuter
                                                   .transformed_structures)
            else:
                for ts in tstructs_or_transmuter:
                    assert isinstance(ts, TransformedStructure)
                self.transformed_structures.extend(tstructs_or_transmuter)

        @staticmethod
        def from_structures(structures, transformations=None, extend_collection=0):
            """
            Alternative constructor from structures rather than
            TransformedStructures.

            Args:
                structures: Sequence of structures
                transformations: New transformations to be applied to all
                    structures
                extend_collection: Whether to use more than one output structure
                    from one-to-many transformations. extend_collection can be a
                    number, which determines the maximum branching for each
                    transformation.

            Returns:
                StandardTransmuter
            """
            tstruct = [TransformedStructure(s, []) for s in structures]
            return StandardTransmuter(tstruct, transformations, extend_collection)

    class CifTransmuter(StandardTransmuter):
        """
        Generates a Transmuter from a cif string, possibly containing multiple
        structures.
        """

        def __init__(self, cif_string, transformations=None, primitive=True,
                     extend_collection=False):
            """
            Generates a Transmuter from a cif string, possibly
            containing multiple structures.

            Args:
                cif_string: A string containing a cif or a series of cifs
                transformations: New transformations to be applied to all
                    structures
                primitive: Whether to generate the primitive cell from the cif.
                extend_collection: Whether to use more than one output structure
                    from one-to-many transformations. extend_collection can be a
                    number, which determines the maximum branching for each
                    transformation.
            """
            transformed_structures = []
            lines = cif_string.split("\n")
            structure_data = []
            read_data = False
            for line in lines:
                if re.match(r"^\s*data", line):
                    structure_data.append([])
                    read_data = True
                if read_data:
                    structure_data[-1].append(line)
            for data in structure_data:
                tstruct = TransformedStructure.from_cif_string("\n".join(data), [],
                                                               primitive)
                transformed_structures.append(tstruct)
            super().__init__(transformed_structures, transformations,
                             extend_collection)

        @staticmethod
        def from_filenames(filenames, transformations=None, primitive=True,
                           extend_collection=False):
            """
            Generates a TransformedStructureCollection from a cif, possibly
            containing multiple structures.

            Args:
                filenames: List of strings of the cif files
                transformations: New transformations to be applied to all
                    structures
                primitive: Same meaning as in __init__.
                extend_collection: Same meaning as in __init__.
            """

            allcifs = []
            for fname in filenames:
                with open(fname, "r") as f:
                    allcifs.append(f.read())
            return CifTransmuter("\n".join(allcifs), transformations,
                                 primitive=primitive,
                                 extend_collection=extend_collection)

    class PoscarTransmuter(StandardTransmuter):
        """
        Generates a transmuter from a sequence of POSCARs.
        """

        def __init__(self, poscar_string, transformations=None,
                     extend_collection=False):
            """
            Args:
                poscar_string: List of POSCAR strings
                transformations: New transformations to be applied to all
                    structures.
                extend_collection: Whether to use more than one output structure
                    from one-to-many transformations.
            """
            tstruct = TransformedStructure.from_poscar_string(poscar_string, [])
            super().__init__([tstruct], transformations,
                             extend_collection=extend_collection)

        @staticmethod
        def from_filenames(poscar_filenames, transformations=None,
                           extend_collection=False):
            """
            Convenient constructor to generates a POSCAR transmuter from a list of
            POSCAR filenames.

            Args:
                poscar_filenames: List of POSCAR filenames
                transformations: New transformations to be applied to all
                    structures.
                extend_collection:
                    Same meaning as in __init__.
            """
            tstructs = []
            for filename in poscar_filenames:
                with open(filename, "r") as f:
                    tstructs.append(TransformedStructure
                                    .from_poscar_string(f.read(), []))
            return StandardTransmuter(tstructs, transformations,
                                      extend_collection=extend_collection)

    def batch_write_vasp_input(transformed_structures, vasp_input_set=MPRelaxSet,
                               output_dir=".", create_directory=True,
                               subfolder=None,
                               include_cif=False, **kwargs):
        """
        Batch write vasp input for a sequence of transformed structures to
        output_dir, following the format output_dir/{group}/{formula}_{number}.

        Args:
            transformed_structures: Sequence of TransformedStructures.
            vasp_input_set: pymatgen.io.vaspio_set.VaspInputSet to creates
                vasp input files from structures.
            output_dir: Directory to output files
            create_directory (bool): Create the directory if not present.
                Defaults to True.
            subfolder: Function to create subdirectory name from
                transformed_structure.
                e.g., lambda x: x.other_parameters["tags"][0] to use the first
                tag.
            include_cif (bool): Boolean indication whether to output a CIF as
                well. CIF files are generally better supported in visualization
                programs.
        """
        for i, s in enumerate(transformed_structures):
            formula = re.sub(r"\s+", "", s.final_structure.formula)
            if subfolder is not None:
                subdir = subfolder(s)
                dirname = os.path.join(output_dir, subdir,
                                       "{}_{}".format(formula, i))
            else:
                dirname = os.path.join(output_dir, "{}_{}".format(formula, i))
            s.write_vasp_input(vasp_input_set, dirname,
                               create_directory=create_directory, **kwargs)
            if include_cif:
                from pymatgen.io.cif import CifWriter

                writer = CifWriter(s.final_structure)
                writer.write_file(os.path.join(dirname, "{}.cif".format(formula)))

    def _apply_transformation(inputs):
        """
        Helper method for multiprocessing of apply_transformation. Must not be
        in the class so that it can be pickled.

        Args:
            inputs: Tuple containing the transformed structure, the transformation
                to be applied, a boolean indicating whether to extend the
                collection, and a boolean indicating whether to clear the redo

        Returns:
            List of output structures (the modified initial structure, plus
            any new structures created by a one-to-many transformation)
        """
        ts, transformation, extend_collection, clear_redo = inputs
        new = ts.append_transformation(transformation, extend_collection,
                                       clear_redo=clear_redo)
        o = [ts]
        if new:
            o.extend(new)
        return o