# import libraries
from Bio.PDB import *
import nglview as nv

"""
Explanation: Biophysics. Examples of code using Bio.PDB
Requirements
biophyton
jupyter
nglview
End of explanation
"""


pdb_file = '1ubq.pdb'

pdb_parser = PDBParser()

st = pdb_parser.get_structure('1UBQ', '1ubq.pdb')

import nglview as nv

view = nv.show_biopython(st)
view

"""
Explanation: 1.1 Loading structure from PDB file
End of explanation
"""


st.header

"""
Explanation: 1.2 Get Headers
End of explanation
"""


pdbl = PDBList()
cif_fn = pdbl.retrieve_pdb_file("1UBQ")
cif_fn

cif_parser = MMCIFParser()
cif_st = cif_parser.get_structure('1UBQ', cif_fn)

view = nv.show_biopython(st)
view

"""
Explanation: 1.3 Alternatively . Download from PDB (in mmCIF format, this is the default)
pdbl.retrieve_pdb_file gives as output the name of the downloaded file.
End of explanation
"""


mmcif_dict = MMCIF2Dict.MMCIF2Dict(cif_fn)
mmcif_dict

"""
Explanation: 1.4 Get Extended data from mmCIF file (very long output, structure defined at PDB mmCIF format)
End of explanation
"""


len(st)

"""
Explanation: Composition

3.1 Number of models
All elements are just lists of elements of the following level: structure is a list of models, model is a list of chains, etc. Here we just check for the length of the list of models.
End of explanation
"""


for chain in st[0]:
    print(chain.id, len(chain))

"""
Explanation: 3.2 Chains and length of the first model
We iterate along the list of chains contained in the first model (st[0]), len(chain) gives the length of the array, i.e. the number of residues that contains.
End of explanation
"""


for num_res in range(10):
    # Note that residues start in 1 in 1ubq, may not be the case
    # in other proteins
    res = st[0]['A'][num_res + 1]
    print(res.get_resname(), res.get_parent().id, res.id[1])

"""
Explanation: 3.3 List of first 10 residues of chain 'A'
Chains are lists of residues using the residue number as index. st[0] corresponds to the first model, st[0]['A'] points to the chain with id 'A' within st[0].
End of explanation
"""


for atom in st[0]['A'][1].get_atoms():
    print(atom.id)

vars(st[0]['A'][1]['N'])

"""
Explanation: 3.4 List of atoms of first residue and complete information of first atom
get_atoms() gives the list of atoms of any element. Can be used with structures, models, chains or residues. Here st[0]['A'][1] stands for the residue labelled as 1 in chain A of the first model (st[0]).
Vars is a standart python function to give the contents of a given dictionary/object. The atom dictionary is described in Bio.PDB.Atom. 
Note that the information in Atom corresponds to one line in the PDB format, in this case N atom on residue A 1
ATOM      1  N   MET A   1      27.340  24.430   2.614  1.00  9.67           N
End of explanation
"""


residues = ['ARG', 'ASP']
atoms = ['CA']
select = []
for atom in st.get_atoms():
    if atom.get_parent().get_resname() in residues and atom.id in atoms:
        select.append(atom)

for atom in select:
    print(atom.get_parent().get_resname(), atom.get_parent().id[1],
          atom.get_name(), atom.get_coord())

"""
Explanation: List of atoms and coordinates for a given list of residues and atoms

First loop goes through all structure (note that here get_atoms() is applied to st, could be also applied to other selections), selecting atoms and residues included in residues and atoms lists. 
Second loop goes through the selection and prints atom id and coordinates (as lists). It can be done in a single loop, but in this way select can be re-used for other purposes
End of explanation
"""


residue_num = 24
chain_id = 'A'
for atom in st[0][chain_id][residue_num]:
     print(atom.get_parent().get_resname(), atom.get_parent().id[1],
          atom.get_name(), atom.get_coord())

"""
Explanation: list atoms and coordinates for a given residue

Note here that here we select a specific residue not a residue type as above.
End of explanation
"""


# chain_A = st[0]['A'] 
residue_1 = 10
residue_2 = 20

import numpy as np

# simple functions to get atom and residue ids properly printed

def residue_id(res): #residue as ASP A32 
    #res.id[1] is integer, we should use str to get the corresponding string
    return res.get_resname() + ' ' + res.get_parent().id + str(res.id[1])

def atom_id(atom): #atom as ASP A32.N, re-uses residue_id
    return residue_id(atom.get_parent()) + '.' + atom.id

for at1 in chain_A[residue_1]:
    for at2 in chain_A[residue_2]:
        dist = at2 - at1     # Direct procedure with (-) to compute distances
        
        vector = at2.coord - at1.coord  # Or using numpy 
        distance = np.sqrt(np.sum(vector ** 2))
        
        print(atom_id(at1), atom_id(at2), dist, distance)

"""
Explanation: Print all distances between the atoms of two residues

We introduce here two functions to make nicer prints for residues and atoms. This is a quite standard way but other formats are possible. If model should be indicated, this is typically written as /model (GLY A10.N/0). 
Distances are evaluated in two alternative ways. Bio.PDB provides a shorcut for distance just substracting atoms. Second option uses numpy functions and calculates the length of a vector going from one atom to the other. Should not be any difference.
End of explanation
"""


center = np.array([10, 10, 10])

for at1 in chain_A[residue_1].get_atoms():
    vect = at1.coord - center
    distance = np.sqrt(np.sum(vect ** 2))
    print(atom_id(at1), distance)


"""
Explanation: Print all distances from atoms in the residue to a given point

Version of a same script using a central point of reference.
End of explanation
"""


MAXDIST = 10

select = []

#Select only CA atoms for short
for at in st.get_atoms():
    if at.id == 'CA':
        select.append(at)

# Preparing search
nbsearch = NeighborSearch(select)

ncontact = 0

for at1, at2 in nbsearch.search_all(MAXDIST): # All pairs with a distance less than MAXDIST
    ncontact += 1
    print("Contact: ", ncontact, atom_id(at1), atom_id(at2), at2 - at1)
    

"""
Explanation: Find contacts below some distance cut-off, using NeighborSearch

This can be solved iterating over all combination of pairs of residues and selecting those that are closer than MAXDIST. 
NeighborSearch uses a different internal structure that makes this search much faster, especially in large structures.
End of explanation
"""


st_6axg = pdb_parser.get_structure('6AXG', '6axg.pdb')

view = nv.show_biopython(st_6axg)
view

"""
Explanation: 9: Select some chain and discard others (e.g. to get a Biological assembly from an assymetrical unit)
This is the "Biopython" way of selecting parts of the structure. Can be used also to select residues, or atoms, but it can become complicated. Most experienced users will just edit the PDB file and remove the chains, residues or atoms manually. 
Here we use as example 6axg, a complex of two proteins (the biological assembly). Version downloaded from PDB (the asymmetric unit) has 6 copies of the complex. Chains A and C correspond to one of such copies.
End of explanation
"""


chains_ok = ['A', 'C']
chain_orig = [ch.id for ch in st_6axg[0]] 

for chain_id in chain_orig:
    if chain_id not in chains_ok:
        st_6axg[0].detach_child(chain_id)

view = nv.show_biopython(st_6axg)
view


"""
Explanation: selecting chains A and C (one biological assembly)
End of explanation
"""


pdbio = PDBIO()

output_pdb_path = '6axg_AC.pdb'

pdbio.set_structure(st_6axg)
pdbio.save(output_pdb_path)

"""
Explanation: Save it on a PDB file
End of explanation
"""


