from prody import *
from numpy import *
from matplotlib.pyplot import *

"""
Explanation: Ensemble Analysis
This example compares experimental structural data analyzed using Principal Component Analysis (PCA) with the theoretical data predicted by Anisotropic Network Model (ANM):
First make the necessary imports:
End of explanation
"""


p38 = parsePDB('1p38', compressed=False) # MAP KINASE

"""
Explanation: Retrieve dataset
One way to retrieve data is to run an NCBI BLAST search against the PDB with the function blastPDB.
To do this, we first need to obtain a sequence and one way to do that is from the PDB:
End of explanation
"""


p38_sequence = p38['A'].getSequence()
p38_sequence

"""
Explanation: We just want one sequence, so we get the sequence of chain A from the pdb file
End of explanation
"""


import pickle

blast_record = pickle.load(open('blast_record3.pkl', 'rb'))

blast_record

"""
Explanation: Once we have this sequence, we can use it in the function blastPDB. We also provide a filename to save the output so we don't need to run it again. To reduce demand on the NCBI webserver, we have provided you this file so please do not run this command.
Instead, please load the data from the pickle file provided:
End of explanation
"""


hits = blast_record.getHits(percent_identity=90, percent_overlap=70)

pdbids = list(hits.keys())
len(pdbids)

pdbids[:10]

"""
Explanation: We can get hits from this record using certain parameters to filter them and extract a list of PDB IDs from them.
End of explanation
"""


from os import mkdir
from os.path import isdir

if not isdir('pdbs'):
    mkdir('pdbs')

pathPDBFolder('pdbs')

pdbs = parsePDB(pdbids)

"""
Explanation: Next, we will use the parsePDB function to import each one of the structures corresponding to these IDs. 
Before doing that, we will make a folder to put them in (if it doesn't already exist) and configure ProDy to use that folder with the function pathPDBFolder.
End of explanation
"""


pathPDBFolder('')

"""
Explanation: After parsing the structures from the PDB, we can use this function again to reset the default download folder back to our current directory:
End of explanation
"""


ref_structure = p38
ref_selection = ref_structure.select('resnum 5 to 31 36 to 114 122 to '
                                     '169 185 to 351 and calpha')

"""
Explanation: Set reference chain
Next, we make a selection to use as the reference for ensemble building:
End of explanation
"""


ref_chain = ref_selection['A']
repr(ref_chain)

"""
Explanation: We extract chain A by indexing to get a Chain object to make things easier.
End of explanation
"""


pdbs.insert(0, ref_chain)

ensemble = buildPDBEnsemble(pdbs, ref=ref_chain, title='p38')

ensemble

"""
Explanation: Ensemble Preparation
We will prepare a PDBEnsemble by mapping each structure against the reference chain and adding a coordinates set corresponding to the mapped atoms. We first make sure that our list of PDB structures (pdbs) includes the ref_chain.
End of explanation
"""


pca = PCA('p38 xray')           # Instantiate a PCA instance
pca.buildCovariance(ensemble)   # Build covariance for the ensemble
pca.calcModes()                 # Calculate modes (20 of the by default)

"""
Explanation: Ensemble Dynamics
Now we will examine the structural dynamics of this ensemble using two different methods
1. Principal Component Analysis (PCA)
PCA is a method that identifies the components which account for the greatest amount of variability in your dataset, i.e. ensemble.
End of explanation
"""


for mode in pca[:5]:    # Print % variance explained by top PCs
    var = calcFractVariance(mode)*100
    print('{0:s}  % variance = {1:.2f}'.format(repr(mode), var))

"""
Explanation: The components/modes of variation are sorted such that the first modes contribute the greatest fractional variance, which we can show as follows:
End of explanation
"""


anm = ANM('1p38')             # Instantiate a ANM instance
anm.buildHessian(ref_chain)   # Build Hessian for the reference chain
anm.calcModes()               # Calculate slowest non-trivial 20 modes

"""
Explanation: The first modes with the highest fractional variance are called the principal components (PCs).
2. Anisotropic Network Model (ANM) Normal Mode Analysis (NMA)
The ANM allows for the identification of the most impactful (slowest) modes in dynamics of a single protein, which we can compare to the principal components from PCA.
End of explanation
"""


for mode in pca[:3]:    # Print PCA mode collectivity
    coll = calcCollectivity(mode)
    print('{0:s}  collectivity = {1:.2f}'.format(repr(mode), coll))

for mode in anm[:3]:    # Print ANM mode collectivity
    coll = calcCollectivity(mode)
    print('{0:s}  collectivity = {1:.2f}'.format(repr(mode), coll))

"""
Explanation: Analysis of PCA and ANM modes
Collectivity of modes
One property that we can calculate and compare is the collectivity, which describes the extent to which a mode collectively recruits large portions of the structure. We see that most of the first modes from both calculations are highly collective.
End of explanation
"""


printOverlapTable(pca[:3], anm[:3]) # Top 3 PCs vs slowest 3 ANM modes

showOverlapTable(pca[:6], anm[:6]);
title('PCA - ANM Overlap Table');

"""
Explanation: PCA - ANM overlap
We can also look at how well the modes produced from each method correlate with each other using the overlap (correlation cosine).
End of explanation
"""


showOverlap(pca[0], anm);
showCumulOverlap(pca[0], anm, color='r');

"""
Explanation: The overlap table shows how each mode from the two methods overlap with each other. Some modes overlap very well with one other mode while others overlap with multiple modes to a lesser extent.
We can also look at the overlap between one mode and all others as follows. The cumulative overlap is the square root of the sum of squared overlaps.
End of explanation
"""


showScaledSqFlucts(pca[1], anm[1]);
legend();

showScaledSqFlucts(pca[3], anm[1]);
legend();

"""
Explanation: Square Fluctuations
In order to see where in the protein these important motions occur, we can visualize the square fluctuations of the principal components and/or slow ANM modes as follows. The function showScaledSqFlucts allows us to scale the square fluctuations from each set of modes to have the same overall size for easier comparison.
We can apply this to individual modes, such as overlapping mode 1 (index 0) from each method, or multiple modes such as the first 3.
End of explanation
"""


showCrossCorr(pca[0]);

showCrossCorr(anm[0]);

"""
Explanation: Cross Correlations
We can also see how correlated the motions for each residue are with each other residue. We see similar patterns for the two methods, especially when using a large number of modes.
End of explanation
"""


writePDB('p38_ref_chain.pdb', ref_chain)
saveEnsemble(ensemble)
saveModel(pca)
saveModel(anm)

"""
Explanation: Saving your work
We can save the outputs from each step for loading back into ProDy as follows:
End of explanation
"""


writeNMD('p38_pca.nmd',anm,ref_chain)

writeNMD('p38_anm.nmd',pca,ref_chain)

"""
Explanation: We can also prepare outputs in NMD format for visualising in VMD with the normal mode wizard:
End of explanation
"""


