from openpathsampling.ensemble import SlicedTrajectoryEnsemble, SequentialEnsemble, AllInXEnsemble, AllOutXEnsemble, LengthEnsemble
from openpathsampling.collectivevariable import FunctionCV
from openpathsampling.volume import CVDefinedVolume
from openpathsampling.engines import Trajectory

# This is a hack to easily create test sequences that act as "trajectories" for us
from openpathsampling.tests.test_helpers import CallIdentity
op = CallIdentity()
vol = CVDefinedVolume(op, -0.5, 0.5)

ens = SequentialEnsemble([
    AllInXEnsemble(vol), 
    AllOutXEnsemble(vol), 
    AllInXEnsemble(vol) & LengthEnsemble(1)
])

inV = 0.0
outV = 1.0

"""
Explanation: Combining SequentialEnsembles with SlicedTrajectoryEnsembles
The ensemble creation features in OpenPathSampling are very powerful, but the ways that some of them interact with each other can be quite confusing. In this document, we explore what happens when one combines SequentialEnsembles and SlicedTrajectoryEnsembles.
First, let's summarize what each of these ensembles do. SequentialEnsembles take a list of ensembles and applies them in order, so a trajectory satisfies the first ensemble for as long as it can, then it switches to the next ensemble, and so forth. SlicedTrajectoryEnsembles modify the trajectory seen by a given ensemble by applying a Python slice object to select only certain frames from the ensemble.
There are several ways of slicing SequentialEnsembles; the purpose of this document is to explain the differences.

To make discussion easy, we'll use a fake order parameter and fake trajectories in this example. We'll also define all our ensembles in terms of AllInXEnsembles and AllOutXEnsembles, which are easier to visualize.
When dealing with SequentialEnsembles, we should distinguish between two types of trajectories. There's the total trajectory for the ensemble, which the normal trajectory including all frames. However, each subensemble of the SequentialEnsemble applies to a subtrajectory. The key to the subtleties of how SequentialEnsembles interact with SlicedTrajectoryEnsembles is to consider the which of these trajectories we really want to slice.
First, we import the necessary things and set up our order parameter and our initial ensemble:
End of explanation
"""


even_slice = slice(None,None,2)
total_sliced_whole = SlicedTrajectoryEnsemble(ens, even_slice)

"""
Explanation: Slicing the global trajectory for the whole SequentialEnsemble
One approach we might need is to apply the whole SequentialEnsemble on some appropriate sliced trajectory. In our example, this translates as saying that the even slices of the total trajectory must satisfy the In-Out-In ensemble.
This is quite easy to implement:
End of explanation
"""


subtraj_sliced_member = SequentialEnsemble([
    AllInXEnsemble(vol),
    SlicedTrajectoryEnsemble(AllOutXEnsemble(vol), even_slice),
    AllInXEnsemble(vol) & LengthEnsemble(1)
])

"""
Explanation: Slicing the subtrajectory for a member of the SequentialEnsemble
Another approach we might desire is to slice the subtrajectory seen by one of the members of the SequentialEnsemble. Perhaps, for example, we want an ensemble which consists of a segment with all frames inside the state, then a segment where the even frames of that segement are outside the state, and then a segment with all frames inside the state.
Note carefully here that the slicing refers to the subtrajectory: it's the even frames counting from the beginning of the subtrajectory, regardless of how many frames are in the total trajectory.
In this case, we have a define a slightly different SequentialEnsemble:
End of explanation
"""


# TODO: We have yet to implement SlicedSequentialEnsemble. 
# This would be nice for completeness, but it really shouldn't be a priority.

"""
Explanation: Slicing the total trajectory for a member of the SequentialEnsemble
The last approach would be to use a slice based on the total trajectory, but only apply it to a member of the SequentialEnsemble. For example, we might want to have an ensemble that begins in a state, and after the first exit all even frames (counted from the start of the trajectory) are outside the state. This differs from the previous example because the count is based on the total trajectory, not the subtrajectory which is passed to the ensemble. To do this, we need to use a SlicedSequentialEnsemble, a subclass of SequentialEnsemble which requires initializing with a list containing one slice for each member ensemble (the special case None is turned into slice(None,None)).
End of explanation
"""


traj = {}
traj[0] = Trajectory([inV, outV, inV])
traj[1] = Trajectory([inV, inV, outV, outV, inV, inV])
traj[2] = Trajectory([inV, outV, outV, outV, inV])
traj[3] = Trajectory([inV, outV, inV, outV, inV, outV, inV, inV])
traj[4] = Trajectory([inV, outV, outV, inV, outV, inV, outV, inV, inV])
traj[5] = Trajectory([inV, outV, outV, inV, outV, inV, outV, inV])
traj[6] = Trajectory([inV, outV, inV, outV, inV, inV])
traj[7] = Trajectory([inV, outV, outV, outV, inV, outV])

"""
Explanation: Example trajectories for these ensembles
Let's take these abstract ideas and apply them to some real trajectories to highlight the differences between these ensembles.
End of explanation
"""


for i in range(len(traj)):
    print "ens(traj["+str(i)+"]) ==", ens(traj[i])

"""
Explanation: First, let's use the original ensemble ens: just from inspection, it should be pretty clear that traj[0] and traj[2] are accepted by that ensemble, and the others are not. Let's confirm that:
End of explanation
"""


for i in range(len(traj)):
    print "total_sliced_whole(traj["+str(i)+"]) ==", total_sliced_whole(traj[i])

"""
Explanation: What about when we slice the whole trajectory into its even components as with total_sliced_whole? Remembering that Python counts from 0, that means a trajectory like traj[4] becomes "in, out, out, out, in", which satifies the original ensemble. What happens with the other ensembles?
End of explanation
"""


for i in range(len(traj)):
    print "subtraj_sliced_member(traj["+str(i)+"]) ==", subtraj_sliced_member(traj[i])

"""
Explanation: Notice that traj[0] no longer works: it turned into "in, in", which does not satisfy the ensemble. Not surprisingly, 1, 2, and 4 satisfy the ensemble, and 3 and 6 don't (since all their "out" frames are odd).
But what about traj[5]? It seems like it should satisfy the ensemble: it starts and ends in the state, and all the even frames between are outside the state. But the problem is that the final frame is an odd-numbered frame, and therefore it gets removed by the slicing before the ensemble sees it. This is the same reason that traj[4] is accepted despite violating the LengthEnsemble(1) condition for the last ensemble of the sequence.
Let's also look more carefully at traj[7]: it may seem odd that this one is accepted, since it ends with a frame outside the volume. But again, the trajectory that the ensemble actually sees after the slicing doesn't include that frame, because it is an odd-numbered frame.
Next, we consider the case of subtraj_sliced_member, where we slice based on the subtrajectory instead of the total trajectory. Let's see what comes out of that:
End of explanation
"""


