sset0 = storage.samplesets[0]
numeric_labels = { s.ensemble : s.replica for s in sset0}
string_labels = { s.ensemble : str(s.replica) for s in sset0 }
numeric_to_string = { numeric_labels[e] : string_labels[e] for e in numeric_labels.keys()}

"""
Explanation: Defining labels
Everything is actually done in terms of ensembles. We can map the ensembles to any labels. In our case, we use the initial replica ID associated with the ensemble. We use this as both the numeric and string label.
End of explanation
"""


%%time
trace_1 = paths.trace_ensembles_for_replica(0, storage.steps)

plt.plot([numeric_labels[e] for e in trace_1])

"""
Explanation: Trace of ensemble visited by a replica
In the plot below, you'll see we set the labels on the axis as sset0[e].replica, so we use the replica ID associated with the ensemble in the first timestep.
End of explanation
"""


repx_net = paths.ReplicaNetwork(scheme, storage.steps)
print repx_net

flow = repx_net.flow(bottom=retis.minus_ensemble, top=retis.ensembles[-1])
print flow

flow_num = {numeric_labels[k] : flow[k] for k in flow.keys()}
print flow_num

sorted_vals = []
for k in sorted(flow_num.keys()):
    sorted_vals.append(flow_num[k])

plt.plot(sorted(flow_num.keys()), sorted_vals)

"""
Explanation: Replica flow
End of explanation
"""


repx_net.trips(bottom=retis.minus_ensemble, top=retis.ensembles[-1])

"""
Explanation: Trips
Now we calculate "up" trips, "down" trips, and round-trips.
End of explanation
"""


repx_net.transition_matrix()

"""
Explanation: Transition matrix
The transition matrix is the (unsymmetrized) matrix of the transition probabilities. By default, it automatically sets its order from the Cuthill-McKee reverse ordering algorithm.
End of explanation
"""


import numpy as np
perm = np.random.permutation(len(mstis.all_ensembles))
print perm

order = [mstis.all_ensembles[p] for p in perm]
repx_net.transition_matrix(index_order=order)

"""
Explanation: If you would like to set a different order, that can be done by providing a list of the ensembles in whatever order you choose:
End of explanation
"""


repx_net.mixing_matrix()

"""
Explanation: Mixing matrix
Same as the transition matrix $T$, but $\frac{1}{2}(T+T^T)$.
End of explanation
"""


repxG = paths.ReplicaNetworkGraph(repx_net)

# draw('graphviz') gives better results, but requires pygraphviz
repxG.draw('spring')

"""
Explanation: Making a pretty picture
End of explanation
"""


transitions = repx_net.transitions_from_traces(storage.steps)

for (k1, k2)  in transitions.keys():
    print numeric_labels[k1], numeric_labels[k2], transitions[(k1, k2)]

for (k1, k2) in repx_net.analysis['n_accepted'].keys():
        print numeric_labels[k1], numeric_labels[k2], repx_net.analysis['n_accepted'][(k1, k2)]

"""
Explanation: Blue is a minus interface, red is a normal interface. Multiple state outer interfaces (not in this example) would be green.
Alternate way of calculating transitions
There's another, perhaps better, way to calculate transitions. This does double count, but doesn't care if the how the transition happened (only that it did).
End of explanation
"""


