import sys
sys.path.append("../nmpath/")
from tools_for_notebook0 import *
%matplotlib inline
from mappers import rectilinear_mapper
from ensembles import Ensemble, DiscreteEnsemble, PathEnsemble, DiscretePathEnsemble

"""
Explanation: Testing the non-Markovian Path Analysis Package
End of explanation
"""


plot_traj([],[])

"""
Explanation: 2D Toy model
End of explanation
"""


#Generating MC trajectories
mc_traj1_2d = mc_simulation2D(100000)
mc_traj2_2d = mc_simulation2D(10000)

"""
Explanation: MC simulation
End of explanation
"""


# Empty ensemble with no trajectories
my_ensemble = Ensemble()

"""
Explanation: 1 - Ensemble class (analysis of continuos trajectories)
Stores an esemble (list) of trajectories (np.arrays). The ensemble could have any number of trajectories including no trajectories at all.
Creating an Ensemble
End of explanation
"""


# from a single trajectory
my_ensemble = Ensemble([mc_traj1_2d],verbose=True)

"""
Explanation: From a single trajectory:
End of explanation
"""


# We have to set list_of_trajs = True

my_list_of_trajs = [mc_traj1_2d, mc_traj2_2d]

my_ensemble = Ensemble(my_list_of_trajs, verbose=True)

"""
Explanation: From a list of trajectories:
End of explanation
"""


for traj in my_ensemble:
    print(len(traj))

"""
Explanation: Ensembles are iterable objects
End of explanation
"""


my_ensemble = Ensemble(verbose=True)

my_ensemble.add_trajectory(mc_traj1_2d)

my_ensemble.add_trajectory(mc_traj2_2d)

"""
Explanation: Adding trajectories to the Ensemble
New trajectories can be added to the ensemble as long as there is consistency in the number of variables.
End of explanation
"""


print(my_ensemble)

"""
Explanation: "Printing" the ensemble
End of explanation
"""


stateA = [[0,pi],[0,pi]]
stateB = [[5*pi,6*pi],[5*pi,6*pi]]

my_ensemble.empirical_mfpts(stateA, stateB)

"""
Explanation: Defining states and computing MFPTs
The states are considered intervals in the is the class is Ensemble
End of explanation
"""


seq1 = mc_simulation2D(20000)
seq2 = mc_simulation2D(20000)

my_e1 = Ensemble([seq1])
my_e2 = Ensemble([seq2])

ensemble1 = my_e1 + my_e2

"""
Explanation: Sum of ensembles (ensemble + ensemble)
End of explanation
"""


e1 = Ensemble([[1.,2.,3.,4.]],verbose=True)
e2 = Ensemble([[2,3,4,5]])
e3 = Ensemble([[2,1,1,4]])

my_ensembles = [e1, e2, e3]

ensemble_tot = Ensemble([])

for ens in my_ensembles:
    ensemble_tot += ens

#ensemble_tot.mfpts([1,1],[4,4])

"""
Explanation: Another simple example
End of explanation
"""


n_states = N**2
bin_bounds = [[i*pi for i in range(7)],[i*pi for i in range(7)]]
C1 = my_ensemble._count_matrix(n_states, map_function=rectilinear_mapper(bin_bounds))
print(C1)

K1 = my_ensemble._mle_transition_matrix(n_states, map_function=rectilinear_mapper(bin_bounds))
print(K1)

"""
Explanation: Computing the count matrix and transition matrix
End of explanation
"""


#p_ensemble = PathEnsemble()

"""
Explanation: 2 - PathEnsemble class
Creating a path ensemble object
End of explanation
"""


p_ensemble = PathEnsemble.from_ensemble(my_ensemble, stateA, stateB)

print(p_ensemble)

"""
Explanation: From ensemble
End of explanation
"""


p_ensemble.empirical_mfpts(stateA, stateB)

"""
Explanation: MFPTs
End of explanation
"""


print(p_ensemble._count_matrix(n_states, mapping_function2D))

#clusters = p_ensemble.cluster(distance_metric = 'RMSD', n_cluster=10, method = 'K-means')

"""
Explanation: Count matrix
End of explanation
"""


d_ens = DiscreteEnsemble.from_ensemble(my_ensemble, mapping_function2D)
print(d_ens)

"""
Explanation: 3 - DiscreteEnsemble class
We can generate a discrete ensemble from the same mapping function and we should obtain exaclty the same result:
End of explanation
"""


C2 = d_ens._count_matrix(n_states)
print(C2)

K2= d_ens._mle_transition_matrix(n_states)
print(K2)

"""
Explanation: Count matrix and transition matrix
End of explanation
"""


stateA = [0]
stateB = [N*N-1]

d_ens.empirical_mfpts(stateA, stateB)

"""
Explanation: Defining states and computing MFPTs
The states are now considered sets, defining the states as follow we should obtain the same results
End of explanation
"""


d_ens2 = DiscreteEnsemble.from_transition_matrix(K2, sim_length = 100000)

#d_ens2.mfpts(stateA,stateB)

"""
Explanation: Generating a Discrete Ensemble from the transition matrix
End of explanation
"""


dpathEnsemble = DiscretePathEnsemble.from_ensemble(my_ensemble, stateA, stateB, mapping_function2D)
print(dpathEnsemble)

#MFPT from the transition matrix
dpathEnsemble.nm_mfpt(ini_probs = None, n_states = N*N)

"""
Explanation: 4 - DiscretePathEnsemble class
Creating the DPE
From Ensemble
End of explanation
"""


n_paths = 200

dpathEnsemble2 = DiscretePathEnsemble.from_transition_matrix\
                (K2, stateA = stateA, stateB = stateB, n_paths = n_paths,ini_pops = [1])
    
print(dpathEnsemble2)

"""
Explanation: From the transition matrix
End of explanation
"""


FSs, weights, count = dpathEnsemble2.weighted_fundamental_sequences(K2)
size = len(FSs)

paths = dpathEnsemble2.trajectories
print(count)

"""
Explanation: Fundamental sequence
End of explanation
"""


discrete = [True for i in range(size)]

plot_traj([[paths[i],[]] for i in range(size)] , discrete, \
          line_width=0.2, std=0.5, color='k', title = '{} paths A->B'.format(n_paths))


"""
Explanation: Plotting paths A -> B
End of explanation
"""


plot_traj([[FSs[i],[]] for i in range(size)] ,discrete, \
          line_width=0.5, std=0.2, color='k', title = '{} FSs A->B'.format(n_paths))


lw = [weights[i]*100 for i in range(size)]

#np.random.seed(12)
plot_traj([[FSs[i],[]] for i in range(size)] ,discrete=[True for i in range(size)],\
          line_width = lw,std = 0.002, alpha=0.25)

"""
Explanation: Plotting Fundamental Sequences A -> B
End of explanation
"""


