'''
Main modules for this workflow
'''
import MDAnalysis as mda
from MDAnalysis.analysis import align, rms
from MDAnalysis.tests.datafiles import PSF, DCD #These are testing files from MDanalysis, when using your own MD files don't need to import this anymore

'''
Packages to visualize, plot data, and utilities
'''

import pandas as pd
import matplotlib as mpl
from matplotlib import pyplot as plt
import seaborn as sns
import py3Dmol
import progressbar # Not needed but useful to see analysis progress


%matplotlib inline
%config Completer.use_jedi = False

'''
MDA version
'''

mda.__version__

u=mda.Universe(PSF,DCD) #Loading the MD data topology=PSF, trajectory=DCD; 
print(u.segments,'with {} residues'.format(len(u.residues))) #Inspecting our system

bar=progressbar.ProgressBar(max_value=len(u.trajectory)) #loading a progressbar using the lenght of the simulation
RMSD_hmap=pd.DataFrame() # The Pairwise values will be saved on a pandas dataFrame
for i in range(len(u.trajectory)): #Depending on simulation 's lenght the calculations could be time-consuming. It could be advisable to compute the RMSD in a different step interval than 1
    for j in range(len(u.trajectory)):
        bb = u.select_atoms('backbone') # I'll compute only Backbone RMSD of the protein
        u.trajectory[i]
        A = bb.positions.copy() # coordinates of first frame
        u.trajectory[j]         # forward to second frame
        B = bb.positions.copy()  # coordinates of second frame
        rmsd=rms.rmsd(A, B, center=True, superposition=True) # Center and superposition is enabled to consider only protein movements and not any influence of Periodic Boundary Conditions
        RMSD_hmap.loc[i,j]=rmsd
    bar.update(i+1)

RMSD_hmap.head(5)

fig, ax = plt.subplots(figsize=(8,6))

n = 10
while len(RMSD_hmap)/n > 6:   #This generates new ticks list to label every 10 units for readability of the plot
    n += 10
    
ax=sns.heatmap(RMSD_hmap,square=True,xticklabels=n,yticklabels=n,vmin=0,cmap='RdBu_r',
               cbar_kws=dict(label='RMSD (Å)',shrink=1,orientation='vertical',spacing='uniform',pad=0.02))

plt.title('4AKE',size=26,weight='bold',color='k')
plt.ylabel('Frame',fontsize=22,fontweight='bold', rotation=90)
plt.xlabel('Frame',fontsize=22,fontweight='bold', rotation=0)

#ax.xaxis.tick_top()
plt.tick_params ('both',width=2,labelsize=18)
cax = plt.gcf().axes[-1]
cax.tick_params(labelsize=16)
ax.figure.axes[-1].yaxis.label.set_size(22)
plt.tight_layout()

plt.show()

calphas = u.select_atoms("name CA") #Commonly the RMSD is computed using the aC of the protein to evaluate protein flucturation respecting the initial coordinates
rmsfer = rms.RMSF(calphas,verbose=False).run()
RMSF_table=pd.DataFrame(rmsfer.rmsf,index=calphas.resnums,columns=['RMSF']) #Saving the RMSF value for every aC residue 

RMSF_table.head(5)

plt.rcParams['axes.linewidth'] = 1.5
fig, ax = plt.subplots(figsize=(8,6))


plt.plot(calphas.resnums, rmsfer.rmsf,linewidth=1.5,color='k')

plt.xlabel ('Residue αC',fontsize=16,fontweight='bold')
plt.ylabel ('RMSF (Å)',fontsize=16,fontweight='bold')

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

plt.tick_params ('both',width=2,labelsize=12)
plt.grid (axis='y',alpha=0.5)
plt.tight_layout()
plt.show()

u.add_TopologyAttr('tempfactors') #Initialization of B factors (tempfactors) by MDA in the Universe atoms
rmsf=[]
for atom in u.atoms:
    rmsf.append(RMSF_table.loc[atom.resid,'RMSF']) #A new list of RMSF must be created by ATOM. Our table has this values per aC. We can create a longer list easily

with mda.Writer("RMSf_factor.pdb", u) as PDB: #This is the functionality from MDA which allos us to save PDB files.
    for ts in u.trajectory[0:1]: #I'll save just the first frame in PDB as snapshot  
        u.atoms.tempfactors = rmsf #This adds the custom B factors to the PDB snapshot 
        PDB.write(u.atoms) 

view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',height=400,width=600)

view.addModel(open('RMSf_factor.pdb','r').read(),'pdb')

view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':min(RMSF_table['RMSF']),'max':max(RMSF_table['RMSF'])}}})

view.show()

fig, ax = plt.subplots(figsize=(6, 1))
fig.subplots_adjust(bottom=0.5)

cmap = mpl.cm.jet_r
norm = mpl.colors.Normalize(vmin=min(RMSF_table['RMSF']), vmax=max(RMSF_table['RMSF']))

cb1 = mpl.colorbar.ColorbarBase(ax, cmap=cmap,norm=norm,orientation='horizontal')
cb1.set_label('RMSF (Å)',fontsize=20,fontweight='bold')
cb1.ax.tick_params(labelsize=16) 
fig.show()