"""
@author: JD Chodera
@author: JH Prinz
"""

import numpy as np
import mdtraj as md
import simtk.unit as u

from openpathsampling.netcdfplus import StorableObject, LoaderProxy
import openpathsampling as paths


# ==============================================================================
# TRAJECTORY
# ==============================================================================


class Trajectory(list, StorableObject):
    """
    Simulation trajectory. Essentially a python list of snapshots
    """

    engine = None

    def __init__(self, trajectory=None):
        """
        Create a simulation trajectory object

        Parameters
        ----------

        trajectory : :obj:`Trajectory` or list of :obj:`openpathsampling.engines.BaseSnapshot`
     
       if specified, make a deep copy of specified trajectory
        """

        # Initialize list.
        list.__init__(self)
        StorableObject.__init__(self)

        if trajectory is not None:
            if type(trajectory) is Trajectory:
                self.extend(trajectory
.iter_proxies())
            else:
                self.extend(trajectory)

    def extend(self, iterable):
        if type(iterable) is Trajectory:
            list.extend(self, iterable.iter_proxies())
        else:
            list.extend(self, iterable)

    def to_dict(self):
        return {
            'snapshots': self.as_proxies()
        }

    @classmethod
    def from_dict(cls, dct):
        return cls(dct['snapshots'])

    def __str__(self):
        return 'Trajectory[' + str(len(self)) + ']'

    def __repr__(self):
        return 'Trajectory[' + str(len(self)) + ']'

    def map(self, fnc, allow_fast=True):
        """
        This runs a function and tries to be fast.

        Fast here means that functions that are purely based on CVs can be
        evaluated without actually loading the real Snapshot object. This
        functions tries to do that and if it fails it does it the usual way
        and creates the snapshot object. This bears the possibility that
        the function uses the fake snapshots and returns a non-sense value.
        It is up to the user to make sure this will not happen.
        """

        if allow_fast:
            try:
                return [fnc(frame) for frame in list.__iter__(self)]
            except:
                pass

        return [fnc(frame) for frame in self]

    @property
    def reversed(self):
        """
        Returns a reversed (shallow) copy of the trajectory itself. Effectively
        creates a new Trajectory object and then fills it with shallow reversed
        copies of the contained snapshots.

        Returns
        -------
        :class:`openpathsampling.trajectory.Trajectory`
            the reversed trajectory
        """

        return Trajectory([snap for snap in reversed(self)])

    @property
    def n_snapshots(self):
        """
        Return the number of frames in the trajectory.
        
        Returns
        -------        
        length (int) - the number of frames in the trajectory

        Notes
        -----
        Might be removed in later versions for len(trajectory) is more pythonic

        See also
        --------
        len

        """

        return len(self)

    def __getattr__(self, item):
        """
        Fallback to access Snapshot properties

        """
        if len(self) > 0:
            snapshot_class = self[0].__class__
            if hasattr(snapshot_class, item) or \
                    hasattr(snapshot_class, '__features__') \
                    and item in snapshot_class.__features__.variables:
                first = getattr(self[0], item)
                if type(first) is u.Quantity:
                    inner = first._value
                    if type(inner) is np.ndarray:
                        dtype = inner.dtype

                        out = np.empty(tuple([len(self)] +
                                             list(inner.shape)), dtype=dtype)

                        for idx, s in enumerate(list.__iter__(self)):
                            np.copyto(out[idx], getattr(s, item)._value)

                        return out * first.unit
                    else:
                        out = [None] * len(self)

                        for idx, s in enumerate(list.__iter__(self)):
                            out[idx] = getattr(s, item)

                        return out
                elif type(first) is np.ndarray:
                    dtype = first.dtype

                    out = np.empty(tuple([len(self)] +
                                         list(first.shape)), dtype=dtype)

                    for idx, s in enumerate(list.__iter__(self)):
                        np.copyto(out[idx], getattr(s, item))

                    return out
                else:
                    out = [None] * len(self)

                    for idx, s in enumerate(list.__iter__(self)):
                        out[idx] = getattr(s, item)

                    return out

            else:
                std_msg = "'{0}' object has no attribute '{1}'"
                snap_msg = "Cannot delegate to snapshots. "
                snap_msg += "'{2}' has no attribute '{1}'"
                spacer = "\n                "
                msg = (std_msg + spacer + snap_msg).format(
                    str(self.__class__.__name__), 
                    item,
                    snapshot_class.__name__
                )
                raise AttributeError(msg)

        else:
            return []

    # ==========================================================================
    # LIST INHERITANCE FUNCTIONS
    # ==========================================================================

    def __getslice__(self, *args, **kwargs):
        ret = list.__getslice__(self, *args, **kwargs)
        if type(ret) is list:
            ret = Trajectory(ret)

        return ret

    # this is intuitive. hash(Trajectory(traj)) == hash(traj)
    # but hash(LoaderProxy(..., traj.__uuid__)) != hash(traj)

    def __hash__(self):
        if len(self) == 0:
            return hash(tuple())
        else:
            return hash(
                (list.__getitem__(self, 0), len(self),
                 list.__getitem__(self, -1)))

    # this might be faster, but does not allow to compare arbitrary
    # trajectories as one might expect. hash(Trajectory(traj)) != hash(traj)
    # but hash(LoaderProxy(..., traj.__uuid__)) == hash(traj)
    # it could also lead to better caching and memory behaviour

    # __hash__ = StorableObject.__hash__
    #
    # __eq__ = StorableObject.__eq__
    # __ne__ = StorableObject.__ne__
    #
    def __getitem__(self, index):
        # Allow for numpy style selection using lists
        if hasattr(index, '__iter__'):
            ret = [list.__getitem__(self, i) for i in index]
        else:
            ret = list.__getitem__(self, index)

        if type(ret) is list:
            ret = Trajectory(ret)
        elif type(ret) is LoaderProxy:
            ret = ret.__subject__

        return ret

    def __reversed__(self):
        for snap_idx in range(len(self) - 1, -1, -1):
            yield self[snap_idx].reversed

    def index_symmetric(self, value):
        """
        Return index of a snapshot or its reversed inside a trajectory

        """
        try:
            fw = self.index(value)
        except ValueError:
            fw = None
            pass

        try:
            bw = self.index(value.reversed)
        except ValueError:
            bw = None
            pass

        if fw is None:
            if bw is None:
                raise KeyError(
                    '%r or its reversed is not found in trajectory.')
            else:
                return bw
        else:
            if bw is None:
                return fw
            else:
                return min(fw, bw)

    def contains_symmetric(self, item):
        """
        Test whether a snapshot or its reversed is in a trajectory

        Returns
        -------
        bool

        """
        fw = item in self
        if not fw:
            return item.reversed in self
        