import sys
import logging

# Import the Protein class
from ssbio.core.protein import Protein

# Printing multiple outputs per cell
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"

"""
Explanation: Protein - Structure Mapping, Alignments, and Visualization
This notebook gives an example of how to map a single protein sequence to its structure, along with conducting sequence alignments and visualizing the mutations.
<div class="alert alert-info">

**Input:** Protein ID + amino acid sequence + mutated sequence(s)

</div>

<div class="alert alert-info">

**Output:** Representative protein structure, sequence alignments, and visualization of mutations

</div>

Imports
End of explanation
"""


# Create logger
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # SET YOUR LOGGING LEVEL HERE #

# Other logger stuff for Jupyter notebooks
handler = logging.StreamHandler(sys.stderr)
formatter = logging.Formatter('[%(asctime)s] [%(name)s] %(levelname)s: %(message)s', datefmt="%Y-%m-%d %H:%M")
handler.setFormatter(formatter)
logger.handlers = [handler]

"""
Explanation: Logging
Set the logging level in logger.setLevel(logging.&lt;LEVEL_HERE&gt;) to specify how verbose you want the pipeline to be. Debug is most verbose.

CRITICAL
Only really important messages shown


ERROR
Major errors


WARNING
Warnings that don't affect running of the pipeline


INFO (default)
Info such as the number of structures mapped per gene


DEBUG
Really detailed information that will print out a lot of stuff



<p><div class="alert alert-warning">**Warning:** `DEBUG` mode prints out a large amount of information, especially if you have a lot of genes. This may stall your notebook!</div></p>
End of explanation
"""


# SET FOLDERS AND DATA HERE
import tempfile
ROOT_DIR = tempfile.gettempdir()

PROTEIN_ID = 'SRR1753782_00918'
PROTEIN_SEQ = 'MSKQQIGVVGMAVMGRNLALNIESRGYTVSVFNRSREKTEEVIAENPGKKLVPYYTVKEFVESLETPRRILLMVKAGAGTDAAIDSLKPYLEKGDIIIDGGNTFFQDTIRRNRELSAEGFNFIGTGVSGGEEGALKGPSIMPGGQKDAYELVAPILTKIAAVAEDGEPCVTYIGADGAGHYVKMVHNGIEYGDMQLIAEAYSLLKGGLNLSNEELANTFTEWNNGELSSYLIDITKDIFTKKDEDGNYLVDVILDEAANKGTGKWTSQSALDLGEPLSLITESVFARYISSLKAQRVAASKVLSGPKAQPAGDKAEFIEKVRRALYLGKIVSYAQGFSQLRAASDEYHWDLNYGEIAKIFRAGCIIRAQFLQKITDAYAENADIANLLLAPYFKKIADEYQQALRDVVAYAVQNGIPVPTFSAAVAYYDSYRAAVLPANLIQAQRDYFGAHTYKRTDKEGIFHTEWLE'

# Create the Protein object
my_protein = Protein(ident=PROTEIN_ID, root_dir=ROOT_DIR, pdb_file_type='mmtf')

# Load the protein sequence
# This sets the loaded sequence as the representative one
my_protein.load_manual_sequence(seq=PROTEIN_SEQ, ident='WT', write_fasta_file=True, set_as_representative=True)

"""
Explanation: Initialization of the project
Set these three things:

ROOT_DIR
The directory where a folder named after your PROTEIN_ID will be created


PROTEIN_ID
Your protein ID


PROTEIN_SEQ
Your protein sequence



A directory will be created in ROOT_DIR with your PROTEIN_ID name. The folders are organized like so:
```
    ROOT_DIR
    └── PROTEIN_ID
        ├── sequences  # Protein sequence files, alignments, etc.
        └── structures  # Protein structure files, calculations, etc.
```
End of explanation
"""


# Mapping using BLAST
my_protein.blast_representative_sequence_to_pdb(seq_ident_cutoff=0.9, evalue=0.00001)
my_protein.df_pdb_blast.head()

"""
Explanation: Mapping sequence --> structure
Since the sequence has been provided, we just need to BLAST it to the PDB.
<p><div class="alert alert-info">**Note:** These methods do not download any 3D structure files.</div></p>

Methods
End of explanation
"""


# Download all mapped PDBs and gather the metadata
my_protein.download_all_pdbs()
my_protein.df_pdb_metadata.head(2)

# Set representative structures
my_protein.set_representative_structure()

"""
Explanation: Downloading and ranking structures
Methods
<div class="alert alert-warning">

**Warning:** 
Downloading all PDBs takes a while, since they are also parsed for metadata. You can skip this step and just set representative structures below if you want to minimize the number of PDBs downloaded.

</div>
End of explanation
"""


my_protein.__dict__

"""
Explanation: Loading and aligning new sequences
You can load additional sequences into this protein object and align them to the representative sequence.
End of explanation
"""


# Input your mutated sequence and load it
mutated_protein1_id = 'N17P_SNP'
mutated_protein1_seq = 'MSKQQIGVVGMAVMGRPLALNIESRGYTVSVFNRSREKTEEVIAENPGKKLVPYYTVKEFVESLETPRRILLMVKAGAGTDAAIDSLKPYLEKGDIIIDGGNTFFQDTIRRNRELSAEGFNFIGTGVSGGEEGALKGPSIMPGGQKDAYELVAPILTKIAAVAEDGEPCVTYIGADGAGHYVKMVHNGIEYGDMQLIAEAYSLLKGGLNLSNEELANTFTEWNNGELSSYLIDITKDIFTKKDEDGNYLVDVILDEAANKGTGKWTSQSALDLGEPLSLITESVFARYISSLKAQRVAASKVLSGPKAQPAGDKAEFIEKVRRALYLGKIVSYAQGFSQLRAASDEYHWDLNYGEIAKIFRAGCIIRAQFLQKITDAYAENADIANLLLAPYFKKIADEYQQALRDVVAYAVQNGIPVPTFSAAVAYYDSYRAAVLPANLIQAQRDYFGAHTYKRTDKEGIFHTEWLE'

my_protein.load_manual_sequence(ident=mutated_protein1_id, seq=mutated_protein1_seq)

# Input another mutated sequence and load it
mutated_protein2_id = 'Q4S_N17P_SNP'
mutated_protein2_seq = 'MSKSQIGVVGMAVMGRPLALNIESRGYTVSVFNRSREKTEEVIAENPGKKLVPYYTVKEFVESLETPRRILLMVKAGAGTDAAIDSLKPYLEKGDIIIDGGNTFFQDTIRRNRELSAEGFNFIGTGVSGGEEGALKGPSIMPGGQKDAYELVAPILTKIAAVAEDGEPCVTYIGADGAGHYVKMVHNGIEYGDMQLIAEAYSLLKGGLNLSNEELANTFTEWNNGELSSYLIDITKDIFTKKDEDGNYLVDVILDEAANKGTGKWTSQSALDLGEPLSLITESVFARYISSLKAQRVAASKVLSGPKAQPAGDKAEFIEKVRRALYLGKIVSYAQGFSQLRAASDEYHWDLNYGEIAKIFRAGCIIRAQFLQKITDAYAENADIANLLLAPYFKKIADEYQQALRDVVAYAVQNGIPVPTFSAAVAYYDSYRAAVLPANLIQAQRDYFGAHTYKRTDKEGIFHTEWLE'

my_protein.load_manual_sequence(ident=mutated_protein2_id, seq=mutated_protein2_seq)

# Conduct pairwise sequence alignments
my_protein.pairwise_align_sequences_to_representative()

# View IDs of all sequence alignments
[x.id for x in my_protein.sequence_alignments]

# View the stored information for one of the alignments
my_alignment = my_protein.sequence_alignments.get_by_id('WT_Q4S_N17P_SNP')
my_alignment.annotations
str(my_alignment[0].seq)
str(my_alignment[1].seq)

# Summarize all the mutations in all sequence alignments
s,f = my_protein.sequence_mutation_summary(alignment_type='seqalign')
print('Single mutations:')
s
print('---------------------')
print('Mutation fingerprints')
f

"""
Explanation: Methods
End of explanation
"""


import ssbio.databases.uniprot

this_examples_uniprot = 'P14062'
sites = ssbio.databases.uniprot.uniprot_sites(this_examples_uniprot)
my_protein.representative_sequence.features = sites
my_protein.representative_sequence.features

"""
Explanation: Some additional methods
Getting binding site/other information from UniProt
End of explanation
"""


# Returns a dictionary mapping sequence residue numbers to structure residue identifiers
# Will warn you if residues are not present in the structure
structure_sites = my_protein.map_seqprop_resnums_to_structprop_resnums(resnums=[1,3,45], 
                                                                       use_representatives=True)
structure_sites

"""
Explanation: Mapping sequence residue numbers to structure residue numbers
Methods
End of explanation
"""


# View just the structure
view = my_protein.representative_structure.view_structure(recolor=True)
view

view.add_spacefill(selection='( :A ) and not hydrogen and 17', label_type='res', color='orange')

# Map the mutations on the visualization (scale increased) - will show up on the above view
my_protein.add_mutations_to_nglview(view=view, alignment_type='seqalign', scale_range=(4,7), 
                                    use_representatives=True)

# Add sites as shown above in the table to the view
my_protein.add_features_to_nglview(view=view, use_representatives=True)

"""
Explanation: Viewing structures
The awesome package nglview is utilized as a backend for viewing structures within a Jupyter notebook. ssbio view functions will either return a NGLWidget object, which is the same as using nglview like the below example, or act upon the widget object itself.
```python
This is how NGLview usually works - it will load a structure file and return a NGLWidget "view" object.
import nglview
view = nglview.show_structure_file(my_protein.representative_structure.structure_path)
view
```
Methods
End of explanation
"""


import os.path as op
my_protein.save_json(op.join(my_protein.protein_dir, '{}.json'.format(my_protein.id)))

"""
Explanation: Saving
End of explanation
"""


