"""Trajectory"""
import contextlib
import io
import warnings
from typing import Tuple

import numpy as np

from ase import __version__
from ase.atoms import Atoms
from ase.calculators.calculator import PropertyNotImplementedError
from ase.calculators.singlepoint import SinglePointCalculator, all_properties
from ase.io.formats import is_compressed
from ase.io.jsonio import decode, encode
from ase.io.pickletrajectory import PickleTrajectory
from ase.parallel import world
from ase.utils import tokenize_version

__all__ = ['Trajectory', 'PickleTrajectory']


def Trajectory(filename, mode='r', atoms=None, properties=None, master=None,
               comm=world):
    """A Trajectory can be created in read, write or append mode.

    Parameters:

    filename: str
        The name of the file.  Traditionally ends in .traj.
    mode: str
        The mode.  'r' is read mode, the file should already exist, and
        no atoms argument should be specified.
        'w' is write mode.  The atoms argument specifies the Atoms
        object to be written to the file, if not given it must instead
        be given as an argument to the write() method.
        'a' is append mode.  It acts as write mode, except that
        data is appended to a preexisting file.
    atoms: Atoms object
        The Atoms object to be written in write or append mode.
    properties: list of str
        If specified, these calculator properties are saved in the
        trajectory.  If not specified, all supported quantities are
        saved.  Possible values: energy, forces, stress, dipole,
        charges, magmom and magmoms.
    master: bool
        Controls which process does the actual writing. The
        default is that process number 0 does this.  If this
        argument is given, processes where it is True will write.
    comm: Communicator object
        Communicator to handle parallel file reading and writing.

    The atoms, properties and master arguments are ignored in read mode.
    """
    if mode == 'r':
        return TrajectoryReader(filename)
    return TrajectoryWriter(filename, mode, atoms, properties, master=master,
                            comm=comm)


class TrajectoryWriter:
    """Writes Atoms objects to a .traj file."""

    def __init__(self, filename, mode='w', atoms=None, properties=None,
                 master=None, comm=world):
        """A Trajectory writer, in write or append mode.

        Parameters:

        filename: str
            The name of the file.  Traditionally ends in .traj.
        mode: str
            The mode.  'r' is read mode, the file should already exist, and
            no atoms argument should be specified.
            'w' is write mode.  The atoms argument specifies the Atoms
            object to be written to the file, if not given it must instead
            be given as an argument to the write() method.
            'a' is append mode.  It acts as write mode, except that
            data is appended to a preexisting file.
        atoms: Atoms object
            The Atoms object to be written in write or append mode.
        properties: list of str
            If specified, these calculator properties are saved in the
            trajectory.  If not specified, all supported quantities are
            saved.  Possible values: energy, forces, stress, dipole,
            charges, magmom and magmoms.
        master: bool
            Controls which process does the actual writing. The
            default is that process number 0 does this.  If this
            argument is given, processes where it is True will write.
        comm: MPI communicator
            MPI communicator for this trajectory writer, by default world.
            Passing a different communicator facilitates writing of
            different trajectories on different MPI ranks.
        """
        if master is None:
            master = comm.rank == 0

        self.filename = filename
        self.mode = mode
        self.atoms = atoms
        self.properties = properties
        self.master = master
        self.comm = comm

        self.description = {}
        self.header_data = None
        self.multiple_headers = False

        self._open(filename, mode)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        self.close()

    def set_description(self, description):
        self.description.update(description)

    def _open(self, filename, mode):
        import ase.io.ulm as ulm
        if mode not in 'aw':
            raise ValueError('mode must be "w" or "a".')
        if self.master:
            self.backend = ulm.open(filename, mode, tag='ASE-Trajectory')
            if len(self.backend) > 0 and mode == 'a':
                with Trajectory(filename) as traj:
                    atoms = traj[0]
                self.header_data = get_header_data(atoms)
        else:
            self.backend = ulm.DummyWriter()

    def write(self, atoms=None, **kwargs):
        """Write the atoms to the file.

        If the atoms argument is not given, the atoms object specified
        when creating the trajectory object is used.

        Use keyword arguments to add extra properties::

            writer.write(atoms, energy=117, dipole=[0, 0, 1.0])
        """
        if atoms is None:
            atoms = self.atoms

        for image in atoms.iterimages():
            self._write_atoms(image, **kwargs)

    def _write_atoms(self, atoms, **kwargs):
        b = self.backend

        if self.header_data is None:
            b.write(version=1, ase_version=__version__)
            if self.description:
                b.write(description=self.description)
            # Atomic numbers and periodic boundary conditions are written
            # in the header in the beginning.
            #
            # If an image later on has other numbers/pbc, we write a new
            # header.  All subsequent images will then have their own header
            # whether or not their numbers/pbc change.
            self.header_data = get_header_data(atoms)
            write_header = True
        else:
            if not self.multiple_headers:
                header_data = get_header_data(atoms)
                self.multiple_headers = not headers_equal(self.header_data,
                                                          header_data)
            write_header = self.multiple_headers

        write_atoms(b, atoms, write_header=write_header)

        calc = atoms.calc

        if calc is None and len(kwargs) > 0:
            calc = SinglePointCalculator(atoms)

        if calc is not None:
            if not hasattr(calc, 'get_property'):
                calc = OldCalculatorWrapper(calc)
            c = b.child('calculator')
            c.write(name=calc.name)
            if hasattr(calc, 'todict'):
                c.write(parameters=calc.todict())
            for prop in all_properties:
                if prop in kwargs:
                    x = kwargs[prop]
                else:
                    if self.properties is not None:
                        if prop in self.properties:
                            x = calc.get_property(prop, atoms)
                        else:
                            x = None
                    else:
                        try:
                            x = calc.get_property(prop, atoms,
                                                  allow_calculation=False)
                        except (PropertyNotImplementedError, KeyError):
                            # KeyError is needed for Jacapo.
                            # XXX We can perhaps remove this.
                            x = None
                if x is not None:
                    if prop in ['stress', 'dipole']:
                        x = x.tolist()
                    c.write(prop, x)

        info = {}
        for key, value in atoms.info.items():
            try:
                encode(value)
            except TypeError:
                warnings.warn(f'Skipping "{key}" info.')
            else:
                info[key] = value
        if info:
            b.write(info=info)

        b.sync()

    def close(self):
        """Close the trajectory file."""
        self.backend.close()

    def __len__(self):
        return self.comm.sum_scalar(len(self.backend))


class TrajectoryReader:
    """Reads Atoms objects from a .traj file."""

    def __init__(self, filename):
        """A Trajectory in read mode.

        The filename traditionally ends in .traj.
        """
        self.filename = filename
        self.numbers = None
        self.pbc = None
        self.masses = None

        self._open(filename)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        self.close()

    def _open(self, filename):
        import ase.io.ulm as ulm
        self.backend = ulm.open(filename, 'r')
        self._read_header()

    def _read_header(self):
        b = self.backend
        if b.get_tag() != 'ASE-Trajectory':
            raise OSError('This is not a trajectory file!')

        if len(b) > 0:
            self.pbc = b.pbc
            self.numbers = b.numbers
            self.masses = b.get('masses')
            self.constraints = b.get('constraints', '[]')
            self.description = b.get('description')
            self.version = b.version
            self.ase_version = b.get('ase_version')

    def close(self):
        """Close the trajectory file."""
        self.backend.close()

    def __getitem__(self, i=-1):
        if isinstance(i, slice):
            return SlicedTrajectory(self, i)
        b = self.backend[i]
        if 'numbers' in b:
            # numbers and other header info was written alongside the image:
            atoms = read_atoms(b, traj=self)
        else:
            # header info was not written because they are the same:
            atoms = read_atoms(b,
                               header=[self.pbc, self.numbers, self.masses,
                                       self.constraints],
                               traj=self)
        if 'calculator' in b:
            results = {}
            implemented_properties = []
            c = b.calculator
            for prop in all_properties:
                if prop in c:
                    results[prop] = c.get(prop)
                    implemented_properties.append(prop)
            calc = SinglePointCalculator(atoms, **results)
            calc.name = b.calculator.name
            calc.implemented_properties = implemented_properties

            if 'parameters' in c:
                calc.parameters.update(c.parameters)
            atoms.calc = calc

        return atoms

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

    def __iter__(self):
        for i in range(len(self)):
            yield self[i]


class SlicedTrajectory:
    """Wrapper to return a slice from a trajectory without loading
    from disk. Initialize with a trajectory (in read mode) and the
    desired slice object."""

    def __init__(self, trajectory, sliced):
        self.trajectory = trajectory
        self.map = range(len(self.trajectory))[sliced]

    def __getitem__(self, i):
        if isinstance(i, slice):
            # Map directly to the original traj, not recursively.
            traj = SlicedTrajectory(self.trajectory, slice(0, None))
            traj.map = self.map[i]
            return traj
        return self.trajectory[self.map[i]]

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


def get_header_data(atoms):
    return {'pbc': atoms.pbc.copy(),
            'numbers': atoms.get_atomic_numbers(),
            'masses': atoms.get_masses() if atoms.has('masses') else None,
            'constraints': list(atoms.constraints)}


def headers_equal(headers1, headers2):
    assert len(headers1) == len(headers2)
    eq = True
    for key in headers1:
        eq &= np.array_equal(headers1[key], headers2[key])
    return eq


class VersionTooOldError(Exception):
    pass


def read_atoms(backend,
               header: Tuple = None,
               traj: TrajectoryReader = None,
               _try_except: bool = True) -> Atoms:
    from ase.constraints import dict2constraint

    if _try_except:
        try:
            return read_atoms(backend, header, traj, False)
        except Exception as ex:
            if (traj is not None and tokenize_version(__version__) <
                    tokenize_version(traj.ase_version)):
                msg = ('You are trying to read a trajectory file written '
                       f'by ASE-{traj.ase_version} from ASE-{__version__}. '
                       'It might help to update your ASE')
                raise VersionTooOldError(msg) from ex
            else:
                raise

    b = backend
    if header:
        pbc, numbers, masses, constraints = header
    else:
        pbc = b.pbc
        numbers = b.numbers
        masses = b.get('masses')
        constraints = b.get('constraints', '[]')

    atoms = Atoms(positions=b.positions,
                  numbers=numbers,
                  cell=b.cell,
                  masses=masses,
                  pbc=pbc,
                  info=b.get('info'),
                  constraint=[dict2constraint(d)
                              for d in decode(constraints)],
                  momenta=b.get('momenta'),
                  magmoms=b.get('magmoms'),
                  charges=b.get('charges'),
                  tags=b.get('tags'))
    return atoms


def write_atoms(backend, atoms, write_header=True):
    b = backend

    if write_header:
        b.write(pbc=atoms.pbc.tolist(),
                numbers=atoms.numbers)
        if atoms.constraints:
            if all(hasattr(c, 'todict') for c in atoms.constraints):
                b.write(constraints=encode(atoms.constraints))

        if atoms.has('masses'):
            b.write(masses=atoms.get_masses())

    b.write(positions=atoms.get_positions(),
            cell=atoms.get_cell().tolist())

    if atoms.has('tags'):
        b.write(tags=atoms.get_tags())
    if atoms.has('momenta'):
        b.write(momenta=atoms.get_momenta())
    if atoms.has('initial_magmoms'):
        b.write(magmoms=atoms.get_initial_magnetic_moments())
    if atoms.has('initial_charges'):
        b.write(charges=atoms.get_initial_charges())


def read_traj(fd, index):
    trj = TrajectoryReader(fd)
    for i in range(*index.indices(len(trj))):
        yield trj[i]


@contextlib.contextmanager
def defer_compression(fd):
    """Defer the file compression until all the configurations are read."""
    # We do this because the trajectory and compressed-file
    # internals do not play well together.
    # Be advised not to defer compression of very long trajectories
    # as they use a lot of memory.
    if is_compressed(fd):
        with io.BytesIO() as bytes_io:
            try:
                # write the uncompressed data to the buffer
                yield bytes_io
            finally:
                # write the buffered data to the compressed file
                bytes_io.seek(0)
                fd.write(bytes_io.read())
    else:
        yield fd


def write_traj(fd, images):
    """Write image(s) to trajectory."""
    if isinstance(images, Atoms):
        images = [images]
    with defer_compression(fd) as fd_uncompressed:
        trj = TrajectoryWriter(fd_uncompressed)
        for atoms in images:
            trj.write(atoms)


class OldCalculatorWrapper:
    def __init__(self, calc):
        self.calc = calc
        try:
            self.name = calc.name
        except AttributeError:
            self.name = calc.__class__.__name__.lower()

    def get_property(self, prop, atoms, allow_calculation=True):
        try:
            if (not allow_calculation and
                    self.calc.calculation_required(atoms, [prop])):
                return None
        except AttributeError:
            pass

        method = 'get_' + {'energy': 'potential_energy',
                           'magmom': 'magnetic_moment',
                           'magmoms': 'magnetic_moments',
                           'dipole': 'dipole_moment'}.get(prop, prop)
        try:
            result = getattr(self.calc, method)(atoms)
        except AttributeError:
            raise PropertyNotImplementedError
        return result


def convert(name):
    import os
    t = TrajectoryWriter(name + '.new')
    for atoms in PickleTrajectory(name, _warn=False):
        t.write(atoms)
    t.close()
    os.rename(name, name + '.old')
    os.rename(name + '.new', name)


def main():
    import optparse
    parser = optparse.OptionParser(usage='python -m ase.io.trajectory '
                                   'a1.traj [a2.traj ...]',
                                   description='Convert old trajectory '
                                   'file(s) to new format. '
                                   'The old file is kept as a1.traj.old.')
    _opts, args = parser.parse_args()
    for name in args:
        convert(name)


if __name__ == '__main__':
    main()