#!/usr/bin/env python
# coding: utf-8

# # Visualize antibody mix
# We will visualize a hypothetical polyclonal antibody mix same way we'd like to visualize the antibody mixes we deconvolve from deep mutational scanning experiments.
# 
# The hypothetical mix represents antibodies targeting three major neutralizing "epitopes" on the SARS-CoV-2 receptor-binding domain (RBD) using the classification scheme of [Barnes et al (2020)](https://www.nature.com/articles/s41586-020-2852-1).
# In particular, [Barnes et al (2020)](https://www.nature.com/articles/s41586-020-2852-1) divided anti-RBD antibodies that bind to the receptor-binding motif into three classes (see also [Greaney et al (2021)](https://www.nature.com/articles/s41467-021-24435-8)).
# For each class, we will use prior deep mutational scanning on a single well-studied monoclonal of that class to antibody to make plausible choices for how mutations affect that antibody class (of course, in reality as there are many somewhat distinct antibodies in each class).
# The antibodies used to represent each class are:
# 
#   - *LY-CoV016*: class 1, mutation estimates from [Starr et al (2021), Science](https://science.sciencemag.org/content/371/6531/850)
#   
#   - *LY-CoV555*: class 2, mutation estimates from [Starr et al (2021), Cell Reports Medicine](https://doi.org/10.1016/j.xcrm.2021.100255)
#   
#   - *REGN10987*: class 3, mutation estimates from [Starr et al (2021), Science](https://science.sciencemag.org/content/371/6531/850)
#   
# Read in the mutation-level escape values $\beta_{m,e}$ for each mutation against each antibody class for this simulated hypothetical mix:

# In[1]:


import pandas as pd

mut_escape_df = pd.read_csv('RBD_mut_escape_df.csv')

mut_escape_df


# Note that the data frame only includes 1932 of the $201 \times 19 = 3819$ possible amino-acid mutations to the RBD; this is because only about half of the mutations are functionally tolerated.
# 
# We also choose simulated activities $a_{\rm{wt},e}$ for each epitope $e$.
# We will let the activity of the polyclonal antibody mix be highest against the class 2 epitope, then next highest against the class 3 epitope, and lowest against the class 1 epitope ([experiments suggest](https://www.nature.com/articles/s41467-021-24435-8) this roughly corresponds to reality for SARS-CoV-2 polyclonal sera):

# In[2]:


activity_wt_df = pd.read_csv('RBD_activity_wt_df.csv')

activity_wt_df


# Now we will visualize the data using `Polyclonal` class provided by this package to model polyclonal antibody mixes:

# In[3]:


import polyclonal

poly_abs = polyclonal.Polyclonal(
                activity_wt_df=activity_wt_df,
                mut_escape_df=mut_escape_df)

print(f"Epitopes: {poly_abs.epitopes}")
print(f"Number of mutations: {len(poly_abs.mutations)}")
print(f"Number of sites: {len(poly_abs.sites)}")


# We can access the activity values, the mutation escape values, and site-level summaries of the mutation escape values:

# In[4]:


poly_abs.activity_wt_df


# In[5]:


poly_abs.mut_escape_df


# In[6]:


poly_abs.mut_escape_site_summary_df


# We can also **plot** the relevant values characterizing the polyclonal mix.
# 
# Here is the activity $a_{\rm{wt},e}$ against each epitope for the unmutated protein:

# In[7]:


# NBVAL_IGNORE_OUTPUT

poly_abs.activity_wt_barplot()


# Here is the mutation escape $\beta_{m,e}$ for each epitope at each site.
# For compact site-level plotting, the mutation escape at each site is summarized by a single number (e.g., as the mean or total of the $\beta_{m,e}$ for that site).
# Note that you can zoom on specific sites, and use the dropdown at the bottom of the plot to select different summary metrics of the escape at each site: 

# In[8]:


# NBVAL_IGNORE_OUTPUT

poly_abs.mut_escape_lineplot()


# If we want to actually interrogate the effect of each single mutation on escape, we can look at the actual $\beta_{m,e}$  in the form of a heatmap.
# Note that the heatmap is again zoomable, and you can mouse over specific mutations to get the value:

# In[9]:


# NBVAL_IGNORE_OUTPUT

poly_abs.mut_escape_heatmap()


# We can also project the site-level summary metrics of the mutation escape onto the protein structure.
# Here we do this using [PDB 6m0j](https://www.rcsb.org/structure/6M0J), which holds the SARS-CoV-2 RBD (chain `E`) in complex with ACE2 (chain `A`).
# Specific, the `Polyclonal` object has a method to make versions of the PDB in which the B-factor is re-assigned to one of the site-level summary metrics of escape (such as *mean* or *total positive*):

# In[10]:


poly_abs.mut_escape_pdb_b_factor(input_pdbfile='6M0J.pdb',
                                 chains='E',
                                 metric='mean',
                                 outfile='RBD_{metric}_{epitope}.pdb')


# These PDB files can then be colored in `pymol` by the escape metric.
# If you want the colors to match the same ones used above by the plotting from the `Polyclonal` object, get the colors and convert them to the RGB tuples used by `pymol`:

# In[11]:


import matplotlib.colors

for epitope, hex_color in poly_abs.epitope_colors.items():
    rgb = [round(val, 3) for val in matplotlib.colors.to_rgb(hex_color)]
    print(f"{epitope}: hex color is {hex_color}; RGB tuple is {rgb}")


# Then using these colors, we can use the `pymol spectrum` command to re-color by B-factor.
# Here is a Python script that can be run within `pymol` (via `run <script>`) to color the ACE2 within the PDB files so that sites in the RBD are colored according to the extent of mutant escape at that site files, with the key helices in ACE2 shown as gray cartoons:
# 
#     epitope_colors = {'class_1': [0.122, 0.467, 0.706],
#                       'class_2': [1.0, 0.498, 0.055],
#                       'class_3': [0.173, 0.627, 0.173],
#                       }
# 
#     for epitope, max_color in epitope_colors.items():
#         cmd.reinitialize()
#         infile = f"mean-{epitope}.pdb"
#         cmd.load(infile)
#         cmd.bg_color('white')
#         cmd.hide('everything')
#         cmd.show('cartoon', '(chain A) and resi 18-85')
#         cmd.color('gray', 'chain A')
#         cmd.set('cartoon_transparency', 0.4, 'chain A')
#         cmd.show('spheres', 'chain E')
#         cmd.set_color('max_color', epitope_colors[epitope])
#         cmd.spectrum('b', 'white max_color', 'chain E')
#         cmd.set_view([0.08, -0.9, -0.3,
#                       -0.9, -0.2, 0.3,
#                       -0.4, 0.3, -0.9,
#                       0.0, 0.0, -160.8,
#                       -32.5, 25.9, 21.2,
#                       87.8, 194.5, -20.0])
#         outfile = os.path.splitext(infile)[0] + '.png'
#         cmd.png(outfile, ray=1)
#         
# Here are the images that result if `pymol` is used to run the above script in the same directory to which we wrote the B-factor recolored PDB files:

# In[12]:


from IPython.display import display, Image

for epitope in poly_abs.epitopes:
    print(f"\n{epitope}:")
    png = f"RBD_mean_{epitope.replace(' ', '_')}.png"
    display(Image(png, width=400))