# Source code for pymatgen.core.trajectory

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

    """
    This module provides classes used to define a MD trajectory.
    """

    import itertools
    import os
    import warnings
    from fnmatch import fnmatch
    from typing import List, Union, Sequence

    import numpy as np
    from monty.io import zopen
    from monty.json import MSONable
    from pymatgen.core.structure import Structure, Lattice, Element, Specie, DummySpecie, Composition
    from pymatgen.io.vasp.outputs import Xdatcar, Vasprun

    __author__ = "Eric Sivonxay, Shyam Dwaraknath"
    __version__ = "0.0"
    __date__ = "Jan 25, 2019"

    class Trajectory(MSONable):
        """
        Trajectory object that stores structural information related to a MD simulation.
        Provides basic functions such as slicing trajectory or obtaining displacements.
        """

        def __init__(self, lattice: Union[List, np.ndarray, Lattice],
                     species: List[Union[str, Element, Specie, DummySpecie, Composition]],
                     frac_coords: List[Sequence[Sequence[float]]],
                     time_step: float = 2,
                     site_properties: dict = None,
                     frame_properties: dict = None,
                     constant_lattice: bool = True,
                     coords_are_displacement: bool = False,
                     base_positions: Sequence[Sequence[float]] = None):
            """
            Create a trajectory object
            Args:
                lattice: The lattice as any 2D array. Each row should correspond to a lattice
                    vector. E.g., [[10,0,0], [20,10,0], [0,0,30]] specifies a
                    lattice with lattice vectors [10,0,0], [20,10,0] and [0,0,30].
                species: List of species on each site. Can take in flexible input,
                    including:
                    i.  A sequence of element / specie specified either as string
                        symbols, e.g. ["Li", "Fe2+", "P", ...] or atomic numbers,
                        e.g., (3, 56, ...) or actual Element or Specie objects.
                    ii. List of dict of elements/species and occupancies, e.g.,
                        [{"Fe" : 0.5, "Mn":0.5}, ...]. This allows the setup of
                        disordered structures.
                frac_coords (MxNx3 array): list of fractional coordinates of
                    each species
                time_step (int, float): Timestep of simulation in femtoseconds. Defaults to 2fs.
                site_properties (list): Properties associated with the sites as a list of
                    dicts of sequences, e.g., [{"magmom":[5,5,5,5]}, {"magmom":[5,5,5,5]}]. The sequences
                    have to be the same length as the atomic species and fractional_coords. Number of supplied
                    dicts should match number of frames in trajectory
                    Defaults to None for no properties.
                frame_properties (dict): Properties of the trajectory such as energy, pressure, etc. each property should
                    have a length equal to the trajectory length. eg: {'energy': [#, #, #, #], 'pressure': [0, 0.1, 0 0.02]}
                constant_lattice (bool): Whether the lattice changes during the simulation, such as in an NPT MD simulation.
                coords_are_displacement (bool): Whether supplied coordinates are given in displacements (True) or
                    positions (False)
                base_positions (Nx3 array): The starting positions of all atoms in trajectory. Used to reconstruct positions
                    when converting from displacements to positions. Only needs to be specified if
                    coords_are_displacement=True. Defaults to first index of frac_coords if coords_are_displacement=False.
            """
            # To support from_dict and as_dict
            if isinstance(frac_coords, list):
                frac_coords = np.array(frac_coords)

            if isinstance(lattice, Lattice):
                lattice = lattice.matrix

            if isinstance(lattice, list):
                lattice = np.array(lattice)

            self.frac_coords = frac_coords
            if coords_are_displacement:
                if base_positions is None:
                    warnings.warn("Without providing an array of starting positions, \
                                   the positions for each time step will not be available")
                self.base_positions = base_positions
            else:
                self.base_positions = frac_coords[0]
            self.coords_are_displacement = coords_are_displacement

            if not constant_lattice and np.shape(lattice) == (3, 3):
                self.lattice = [lattice for i in range(np.shape(self.frac_coords)[0])]
            else:
                self.lattice = lattice
            self.constant_lattice = constant_lattice
            self.species = species
            self.site_properties = site_properties
            self.frame_properties = frame_properties
            self.time_step = time_step

        def get_structure(self, i):
            """
            Returns structure at specified index
            Args:
                i (int): Index of structure
            Returns:
                (Structure) pymatgen structure object
            """
            return self[i]

        def to_positions(self):
            """
            Converts fractional coordinates of trajectory into positions
            """
            if self.coords_are_displacement:
                cumulative_displacements = np.cumsum(self.frac_coords, axis=0)
                positions = self.base_positions + cumulative_displacements
                self.frac_coords = positions
                self.coords_are_displacement = False

        def to_displacements(self):
            """
            Converts position coordinates of trajectory into displacements between consecutive frames
            """
            if not self.coords_are_displacement:
                displacements = np.subtract(self.frac_coords, np.roll(self.frac_coords, 1, axis=0))
                displacements[0] = np.zeros(np.shape(self.frac_coords[0]))
                # Deal with PBC
                displacements = [np.subtract(item, np.round(item)) for item in displacements]

                self.frac_coords = displacements
                self.coords_are_displacement = True

        def extend(self, trajectory):
            """
            Concatenate another trajectory
            Args:
                trajectory (Trajectory): Trajectory to add
            """
            if self.time_step != trajectory.time_step:
                raise ValueError('Trajectory not extended: Time steps of trajectories is incompatible')

            if len(self.species) != len(trajectory.species) and self.species != trajectory.species:
                raise ValueError('Trajectory not extended: species in trajectory do not match')

            # Ensure both trajectories are in positions before combining
            self.to_positions()
            trajectory.to_positions()

            self.site_properties = self._combine_site_props(self.site_properties, trajectory.site_properties,
                                                            np.shape(self.frac_coords)[0],
                                                            np.shape(trajectory.frac_coords)[0])
            self.frame_properties = self._combine_frame_props(self.frame_properties, trajectory.frame_properties,
                                                              np.shape(self.frac_coords)[0],
                                                              np.shape(trajectory.frac_coords)[0])

            self.frac_coords = np.concatenate((self.frac_coords, trajectory.frac_coords), axis=0)
            self.lattice, self.constant_lattice = self._combine_lattice(self.lattice, trajectory.lattice,
                                                                        np.shape(self.frac_coords)[0],
                                                                        np.shape(trajectory.frac_coords)[0])

        def __iter__(self):
            for i in range(np.shape(self.frac_coords)[0]):
                yield self[i]

        def __len__(self):
            return np.shape(self.frac_coords)[0]

        def __getitem__(self, frames):
            """
            Gets a subset of the trajectory if a slice is given, if an int is given, return a structure
            Args:
                frames (int, slice): int or slice of trajectory to return
            Return:
                (Trajectory, Structure) Subset of trajectory
            """
            # If trajectory is in displacement mode, return the displacements at that frame
            if self.coords_are_displacement:
                if isinstance(frames, int):
                    if frames >= np.shape(self.frac_coords)[0]:
                        raise ValueError('Selected frame exceeds trajectory length')
                    # For integer input, return the displacements at that timestep
                    return self.frac_coords[frames]
                if isinstance(frames, slice):
                    # For slice input, return a list of the displacements
                    start, stop, step = frames.indices(len(self))
                    return [self.frac_coords[i] for i in range(start, stop, step)]
                if isinstance(frames, (list, np.ndarray)):
                    # For list input, return a list of the displacements
                    pruned_frames = [i for i in frames if i < len(self)]  # Get rid of frames that exceed trajectory length
                    if len(pruned_frames) < len(frames):
                        warnings.warn('Some or all selected frames exceed trajectory length')
                    return [self.frac_coords[i] for i in pruned_frames]
                raise Exception('Given accessor is not of type int, slice, list, or array')

            # If trajectory is in positions mode, return a structure for the given frame or trajectory for the given frames
            if isinstance(frames, int):
                if frames >= np.shape(self.frac_coords)[0]:
                    raise ValueError('Selected frame exceeds trajectory length')
                # For integer input, return the structure at that timestep
                lattice = self.lattice if self.constant_lattice else self.lattice[frames]
                site_properties = self.site_properties[frames] if self.site_properties else None
                site_properties = self.site_properties[frames] if self.site_properties else None
                return Structure(Lattice(lattice), self.species, self.frac_coords[frames],
                                 site_properties=site_properties,
                                 to_unit_cell=True)
            if isinstance(frames, slice):
                # For slice input, return a trajectory of the sliced time
                start, stop, step = frames.indices(len(self))
                pruned_frames = range(start, stop, step)
                lattice = self.lattice if self.constant_lattice else [self.lattice[i] for i in pruned_frames]
                frac_coords = [self.frac_coords[i] for i in pruned_frames]
                if self.site_properties is not None:
                    site_properties = [self.site_properties[i] for i in pruned_frames]
                else:
                    site_properties = None
                if self.frame_properties is not None:
                    frame_properties = {}
                    for key, item in self.frame_properties.items():
                        frame_properties[key] = [item[i] for i in pruned_frames]
                else:
                    frame_properties = None
                return Trajectory(lattice, self.species, frac_coords, time_step=self.time_step,
                                  site_properties=site_properties, frame_properties=frame_properties,
                                  constant_lattice=self.constant_lattice, coords_are_displacement=False,
                                  base_positions=self.base_positions)
            if isinstance(frames, (list, np.ndarray)):
                # For list input, return a trajectory of the specified times
                pruned_frames = [i for i in frames if i < len(self)]  # Get rid of frames that exceed trajectory length
                if len(pruned_frames) < len(frames):
                    warnings.warn('Some or all selected frames exceed trajectory length')
                lattice = self.lattice if self.constant_lattice else [self.lattice[i] for i in pruned_frames]
                frac_coords = [self.frac_coords[i] for i in pruned_frames]
                if self.site_properties is not None:
                    site_properties = [self.site_properties[i] for i in pruned_frames]
                else:
                    site_properties = None
                if self.frame_properties is not None:
                    frame_properties = {}
                    for key, item in self.frame_properties.items():
                        frame_properties[key] = [item[i] for i in pruned_frames]
                else:
                    frame_properties = None
                return Trajectory(lattice, self.species, frac_coords, time_step=self.time_step,
                                  site_properties=site_properties, frame_properties=frame_properties,
                                  constant_lattice=self.constant_lattice, coords_are_displacement=False,
                                  base_positions=self.base_positions)
            raise Exception('Given accessor is not of type int, slice, tuple, list, or array')

        def copy(self):
            """
            :return: Copy of Trajectory.
            """
            return Trajectory(self.lattice, self.species, self.frac_coords, time_step=self.time_step,
                              site_properties=self.site_properties, frame_properties=self.frame_properties,
                              constant_lattice=self.constant_lattice, coords_are_displacement=False,
                              base_positions=self.base_positions)

        @classmethod
        def from_structures(cls, structures, constant_lattice=True, **kwargs):
            """
            Convenience constructor to obtain trajectory from a list of structures.
            Note: Assumes no atoms removed during simulation

            Args:
                structures (list): list of pymatgen Structure objects.
                constant_lattice (bool): Whether the lattice changes during the simulation, such as in an NPT MD
                    simulation. True results in
            Returns:
                (Trajectory)
            """
            frac_coords = [structure.frac_coords for structure in structures]
            if constant_lattice:
                lattice = structures[0].lattice.matrix
            else:
                lattice = [structure.lattice.matrix for structure in structures]
            site_properties = {}

            site_properties = [structure.site_properties for structure in structures]
            return cls(lattice, structures[0].species, frac_coords, site_properties=site_properties,
                       constant_lattice=constant_lattice, **kwargs)

        @classmethod
        def from_file(cls, filename, constant_lattice=True, **kwargs):
            """
            Convenience constructor to obtain trajectory from XDATCAR or vasprun.xml file
            Args:
                filename (str): The filename to read from.
                constant_lattice (bool): Whether the lattice changes during the simulation, such as in an NPT MD
                    simulation. True results in
            Returns:
                (Trajectory)
            """
            # TODO: Support other filetypes

            fname = os.path.basename(filename)
            if fnmatch(fname, "*XDATCAR*"):
                structures = Xdatcar(filename).structures
            elif fnmatch(fname, "vasprun*.xml*"):
                structures = Vasprun(filename).structures
            else:
                raise ValueError("Unsupported file")

            return cls.from_structures(structures, constant_lattice=constant_lattice, **kwargs)

        def as_dict(self):
            """
            :return: MSONAble dict.
            """
            d = {"@module": self.__class__.__module__,
                 "@class": self.__class__.__name__,
                 "species": self.species, "time_step": self.time_step,
                 "site_properties": self.site_properties,
                 "frame_properties": self.frame_properties,
                 "constant_lattice": self.constant_lattice,
                 "coords_are_displacement": self.coords_are_displacement,
                 "base_positions": self.base_positions}
            d["lattice"] = self.lattice.tolist()
            d["frac_coords"] = self.frac_coords.tolist()

            return d

        @staticmethod
        def _combine_lattice(attr_1, attr_2, len_1, len_2):
            """
            Helper function to combine trajectory properties such as site_properties or lattice
            """
            if np.shape(attr_1) == (3, 3) and np.shape(attr_2) == (3, 3):
                attribute = attr_1
                attribute_constant = True
            elif np.shape(attr_1) == 3 and np.shape(attr_2) == 3:
                attribute = np.concatenate((attr_1, attr_2), axis=0)
                attribute_constant = False
            else:
                attribute = [attr_1.copy()] * len_1 if isinstance(attr_1, list) else attr_1.copy()
                attribute.extend([attr_2.copy()] * len_2 if isinstance(attr_2, list) else attr_2.copy())
                attribute_constant = False
            return attribute, attribute_constant

        @staticmethod
        def _combine_site_props(attr_1, attr_2, len_1, len_2):
            """
            Helper function to combine site properties of 2 trajectories
            """
            if attr_1 is None and attr_2 is None:
                return None
            if attr_1 is None or attr_2 is None:
                new_site_properties = []
                if attr_1 is None:
                    new_site_properties.extend([None for i in range(len_1)])
                elif len(attr_1) == 1:
                    new_site_properties.extend([attr_1[0] for i in range(len_1)])
                elif len(attr_1) > 1:
                    new_site_properties.extend(attr_1)

                if attr_2 is None:
                    new_site_properties.extend([None for i in range(len_2)])
                elif len(attr_2) == 1:
                    new_site_properties.extend([attr_2[0] for i in range(len_2)])
                elif len(attr_2) > 1:
                    new_site_properties.extend(attr_2)

                return new_site_properties
            if len(attr_1) == 1 and len(attr_2) == 1:
                # If both properties lists are do not change within their respective trajectory
                if attr_1 == attr_2:
                    # If both site_properties are the same, only store one
                    return attr_1
                new_site_properties = [attr_1[0] for i in range(len_1)]
                new_site_properties.extend([attr_2[0] for i in range(len_2)])
                return new_site_properties
            if len(attr_1) > 1 and len(attr_2) > 1:
                # Both properties have site properties that change within the trajectory, concat both together
                return [*attr_1, *attr_2]

            new_site_properties = []
            if attr_1 is None:
                new_site_properties.extend([None for i in range(len_1)])
            elif len(attr_1) == 1:
                new_site_properties.extend([attr_1[0] for i in range(len_1)])
            elif len(attr_1) > 1:
                new_site_properties.extend(attr_1)

            if attr_2 is None:
                new_site_properties.extend([None for i in range(len_2)])
            elif len(attr_2) == 1:
                new_site_properties.extend([attr_2[0] for i in range(len_2)])
            elif len(attr_2) > 1:
                new_site_properties.extend(attr_2)

            return new_site_properties

        @staticmethod
        def _combine_frame_props(attr_1, attr_2, len_1, len_2):
            """
            Helper function to combine frame properties such as energy or pressure
            """
            if attr_1 is None and attr_2 is None:
                return None

            # Find all common keys
            all_keys = set(attr_1.keys()).union(set(attr_2.keys()))

            # Initialize dict with the common keys
            new_frame_props = dict(zip(all_keys, [[] for i in all_keys]))
            for key in all_keys:
                if key in attr_1.keys():
                    new_frame_props[key].extend(attr_1[key])
                else:
                    # If key doesn't exist in the first trajectory, append None for each index
                    new_frame_props[key].extend([None for i in range(len_1)])

                if key in attr_2.keys():
                    new_frame_props[key].extend(attr_2[key])
                else:
                    # If key doesn't exist in the second trajectory, append None for each index
                    new_frame_props[key].extend([None for i in range(len_2)])

            return new_frame_props

        def write_Xdatcar(self, filename="XDATCAR", system=None, significant_figures=6):
            """
            Writes Xdatcar to a file. The supported kwargs are the same as those for
            the Xdatcar_from_structs.get_string method and are passed through directly.

            Args:
                filename (str): name of file (It's prudent to end the filename with 'XDATCAR',
                    as most visualization and analysis software require this for autodetection)
                system (str): Description of system
                significant_figures (int): Significant figures in the output file
            """

            # Ensure trajectory is in position form
            self.to_positions()

            if system is None:
                system = f'{self[0].composition.reduced_formula}'

            lines = []
            format_str = "{{:.{0}f}}".format(significant_figures)
            syms = [site.specie.symbol for site in self[0]]
            site_symbols = [a[0] for a in itertools.groupby(syms)]
            syms = [site.specie.symbol for site in self[0]]
            natoms = [len(tuple(a[1])) for a in itertools.groupby(syms)]

            for si, frac_coords in enumerate(self.frac_coords):
                # Only print out the info block if
                if self.constant_lattice and si == 0:
                    lines.extend([system, "1.0"])

                    if self.constant_lattice:
                        _lattice = self.lattice
                    else:
                        _lattice = self.lattice[si]

                    for latt_vec in _lattice:
                        lines.append(f'{" ".join([str(el) for el in latt_vec])}')

                    lines.append(" ".join(site_symbols))
                    lines.append(" ".join([str(x) for x in natoms]))

                lines.append(f"Direct configuration=     {str(si + 1)}")

                for (frac_coord, specie) in zip(frac_coords, self.species):
                    coords = frac_coord
                    line = f'{" ".join([format_str.format(c) for c in coords])} {specie}'
                    lines.append(line)

            xdatcar_string = "\n".join(lines) + "\n"

            with zopen(filename, "wt") as f:
                f.write(xdatcar_string)