from msmbuilder.dataset import dataset
import mdtraj as md
import numpy as np
from glob import glob
from mdtraj.utils import timing
import itertools
from msmbuilder.featurizer import AtomPairsFeaturizer
from msmbuilder.decomposition import tICA
from msmbuilder.cluster import MiniBatchKMeans
from msmbuilder.cluster import MiniBatchKMedoids
%matplotlib inline
from matplotlib import pyplot as plt

traj_list = sorted(glob("/Users/je714/wt_data/*/05*nc"))

top = md.load_prmtop("/Users/je714/wt_data/run8/WT-ff14SB_clean.prmtop")

ef_hand = top.select("resid 64 to 74")
cat_cal = top.select("resid 419")

atom_sel = np.append(ef_hand, cat_cal)

atom_sel

xyz = dataset("/Users/je714/wt_data/*/05*nc", fmt='mdtraj', 
               topology="/Users/je714/wt_data/run8/WT-ff14SB_clean.prmtop",
               atom_indices=atom_sel)

"""
Explanation: MSM of EF-hand
Build an MSM on the EF-hand domain of cTnC and the catalytic Ca<sup>2+</sup> atom.
End of explanation
"""


pairs = np.asarray(list(itertools.product(list(range(139)), [138])))

print(pairs.shape)

pairs_noIJ = pairs[pairs[:,0] != pairs[:,1]] # Drop out the elements that are i=j (whose distance is 0)

pairs_noIJ

traj1 = xyz[0]

dist_featurizer = AtomPairsFeaturizer(pair_indices=pairs)

with timing("Featurizing into distances..."):
    dists = dist_featurizer.fit_transform(xyz)

tica_model = tICA(n_components=10, lag_time=1)
with timing("tica..."):
    tica_trajs_dists = tica_model.fit_transform(dists)

tica_model = tICA

tica_trajs_dists[0].shape

def tica_plotter(tica_trajs):
    
    txx = np.concatenate(tica_trajs)

    plt.figure(figsize=(14, 4))
    plt.subplot(1, 2, 1)
    plt.hexbin(txx[:, 0], txx[:, 1], bins='log', mincnt=1)
    plt.xlabel('tIC 1')
    plt.ylabel('tIC 2')
    cb = plt.colorbar()
    cb.set_label('log10(N)')

    plt.subplot(1, 2, 2)
    
    time = np.linspace(start=0.0, stop=txx.shape[0]*0.02, num = txx.shape[0]) / 1000
    plt.plot(time, txx[:, 0])
    plt.xlabel("Aggregated time ($\mu$s)")
    plt.ylabel("tIC 1")
    
    plt.tight_layout()

tica_plotter(tica_trajs_dists)

clusterer = MiniBatchKMeans(n_clusters=400)
with timing("Clustering tica trajs from dist featurization..."):
    clustered_trajs = clusterer.fit_transform(tica_trajs_dists)

def cluster_plotter(trajectory, clusterer_object):
    txx = np.concatenate(trajectory)
    plt.hexbin(txx[:,0], txx[:,1], bins='log', mincnt=1)
    plt.scatter(clusterer_object.cluster_centers_[:,0],
                clusterer_object.cluster_centers_[:,1], 
                s=100, c='w')

cluster_plotter(tica_trajs_dists, clusterer)

from msmbuilder.msm import MarkovStateModel
from msmbuilder.utils import dump
msm = MarkovStateModel(lag_time=50)
msm.fit(clustered_trajs)
print("The MSM has %s states.\n" % msm.n_states_)
print(msm.left_eigenvectors_.shape)

plt.hexbin(txx[:,0], txx[:,1], bins='log', mincnt=1, cmap="Greys")
plt.scatter(clusterer.cluster_centers_[:,0],
            clusterer.cluster_centers_[:,1],
            s=1e4 * msm.populations_, # size by population
            c=msm.left_eigenvectors_[:,1], # color by eigenvector
            cmap="RdBu") 
plt.colorbar(label='First dynamical eigenvector')
plt.xlabel('tIC 1')
plt.ylabel('tIC 2')
plt.tight_layout()

from msmbuilder.lumping import PCCAPlus
pcca = PCCAPlus.from_msm(msm, n_macrostates=5)
macro_trajs = pcca.transform(clustered_trajs)

plt.hexbin(txx[:,0], txx[:,1], bins='log', mincnt=1, cmap="Greys")
plt.scatter(clusterer.cluster_centers_[:,0],
            clusterer.cluster_centers_[:,1],
            s=100,
            c=pcca.microstate_mapping_,
      )
plt.xlabel('tIC 1')
plt.ylabel('tIC 2')

"""
Explanation: We featurize the trajectories into the pairwise distance between all the atoms in the EF hand and the catalytic Ca<sup>2+</sup>. Since we only loaded the relevant atoms, msmbuilder renumbers the indices from 0.
End of explanation
"""


from sklearn.pipeline import Pipeline
from sklearn.cross_validation import KFold
from sklearn.grid_search import RandomizedSearchCV
from scipy.stats.distributions import randint

model = Pipeline([
    ('featurizer', AtomPairsFeaturizer(pair_indices=pairs)),
    ('tica', tICA(n_components=4)),
    ('clusterer', MiniBatchKMeans()),
    ('msm', MarkovStateModel(n_timescales=3, ergodic_cutoff='on', reversible_type='mle', verbose=False))
])
search = RandomizedSearchCV(model, n_iter=50, cv=3, refit=False, param_distributions={
        'tica__lag_time':randint(1,100),    
        'clusterer__n_clusters':randint(1,1000),
        'msm__lag_time':randint(1,100)
})

def plot_scores(score_matrix):
    from matplotlib import pyplot as pp
    scores = np.array([[np.mean(e.cv_validation_scores),
                        np.std(e.cv_validation_scores),
                        e.parameters['tica__lag_time'],
                        e.parameters['clusterer__n_clusters'],
                        e.parameters['msm__lag_time']]
                        for e in score_matrix.grid_scores_])

    mean_score_value = scores[:,0]
    std_score_value = scores[:,1]
    lags_tica = scores[:,2]
    cluster_number = scores[:,3]
    lags_msm = scores[:,4]

    pp.figure(figsize=(14,4))
    pp.grid(False)
    interval = 2*std_score_value
    # subplot1
    pp.subplot(1,3,1, axisbg='white')
    pp.errorbar(x=lags_tica, y=mean_score_value, yerr=interval, color = 'b', fmt='o')
    pp.plot(pp.xlim(), [search.best_score_]*2, 'k-.', label = 'best')
    pp.xlabel("tICA lag time (Frame)")
    pp.ylabel("Score")
    # subplot2
    pp.subplot(1,3,2, axisbg='white')
    pp.errorbar(x=cluster_number, y=mean_score_value, yerr=interval,  color = 'r', fmt='o')
    pp.plot(pp.xlim(), [search.best_score_]*2, 'k-.', label = 'best')
    pp.xlabel("Number of clusters")
    pp.ylabel("Score")
    # subplot3
    pp.subplot(1,3,3, axisbg='white')
    pp.errorbar(x=lags_msm, y=mean_score_value, yerr=interval, color = 'g', fmt='o')
    pp.plot(pp.xlim(), [search.best_score_]*2, 'k-.', label = 'best')
    pp.xlabel("MSM lag time (Frame)")
    pp.ylabel("Score")

    pp.tight_layout()

# with timing("Cross validation..."):
#     xyz_scores = search.fit(xyz)
# for key, value in xyz_scores.best_params_.items():
#     print("%s\t%s" % (key, value))
# print("Score = %.2f" % xyz_scores.best_score_)

# plot_scores(xyz_scores)

model = Pipeline([
    ('tica', tICA()),
    ('clusterer', MiniBatchKMedoids()),
    ('msm', MarkovStateModel(ergodic_cutoff='on', reversible_type='mle', verbose=False))
])
search = RandomizedSearchCV(model, n_iter=100, cv=5, refit=False, param_distributions={
        'tica__lag_time':randint(1,200),
        'tica__n_components':randint(1,15),
        'clusterer__n_clusters':randint(1,1000),
        'msm__n_timescales':randint(1,15),
        'msm__lag_time':randint(1,100)
})

print(search)

with timing("Cross validation on trajectories already fit to distance feature..."):
    feat_scores = search.fit(dists)

feat_scores.best_params_

plot_scores(feat_scores)

"""
Explanation: Cross validation
End of explanation
"""


opt_tica = tICA(lag_time=118, n_components=13)
opt_cluster = MiniBatchKMedoids(n_clusters=346)
opt_msm = MarkovStateModel(lag_time=20, n_timescales=14)

opt_tica_trajs = opt_tica.fit_transform(dists)

tica_plotter(opt_tica_trajs)

clustered_opt_tica_trajs = opt_cluster.fit_transform(opt_tica_trajs)

cluster_plotter(opt_tica_trajs, opt_cluster)

opt_msm.fit(clustered_opt_tica_trajs)

lag_times = [x*2 for x in range(1,55,2)]
print(lag_times)

opt_msm.timescales_[0:3]

first_three_timescales = []
for lag_time in lag_times:
    msm = MarkovStateModel(lag_time=lag_time, n_timescales=14)
    msm.fit(clustered_opt_tica_trajs)
    first_three_timescales.append(msm.timescales_[0:3])

first_three_timescales

first_timescale = []
for comb in first_three_timescales:
    first_timescale.append(comb[0])

plt.scatter(x=lag_times, y=first_timescale)

"""
Explanation: feat_scores.grid_scores_
Use best scores from cross-validation to build an MSM
End of explanation
"""


