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

# # Reference site numbering
# So far all the preceding examples have involved analysis of proteins with sequential integer site numbers (e.g., 1, 2, 3, ...).
# This sequential integer numbering can also start at sites other than one (e.g., 331, 332, 333, ...).
# 
# However, often proteins are numbered in relation to alignment to sequential integer numbering of some reference homolog.
# For instance, SARS-CoV-2 spike sequences are usually numbered in relation to the Wuhan-Hu-1 reference.
# But if the protein has accumulated insertions or deletions relative to the reference, there may be missing sites (gaps) or insertions (typically numbered as `214`, `214a`, `214b`, etc).
# 
# This notebook shows how to perform an analysis and analyze the results with non-sequential reference site numbering.
# The advantage of doing this is that the results can directly be visualized/analyzed using the conventional site numbering scheme.
# 
# It does this by analyzing deep mutational scanning of the SARS-CoV-2 Omicron BA.1 spike, which has several indels relative to the Wuhan-Hu-1 numbering reference.

# ## Fit and visualize a model with non-integer site numbering
# 
# First, import Python modules:

# In[1]:


import polyclonal

import pandas as pd


# Now read the data to fit:

# In[2]:


# read data w `na_filter=None` so empty aa_substitutions read as such rather than NA
data_to_fit = pd.read_csv(
    "Lib-2_2022-06-22_thaw-1_LyCoV-1404_1_prob_escape.csv",
    na_filter=None,
)

data_to_fit.head()


# Notice how these data have sites numbered in two different schemes:
# 
#  1. `aa_substitutions_sequential`: sequential integer numbering (1, 2, 3, ...) of the protein used in the experiment. You could just analyze the mutations this way, but then the results will not be in standard reference numbering scheme.
#  
#  2. `aa_substitutions_reference`: the referencey based site numbering, which skips sites with indels and has some non-numeric sites where there are insertions (eg, `214a`), as for example in the variant shown below:

# In[3]:


data_to_fit.query("aa_substitutions_reference.str.contains('a')").head(n=1)


# Here we will use the reference based amino-acid substitutions, so create a column with the name used by `Polyclonal` (the "aa_substitutions" column) that uses this numbering scheme:

# In[4]:


data_to_fit["aa_substitutions"] = data_to_fit["aa_substitutions_reference"]


# Importantly, in order to use reference based numbering that is not sequential integer, you also have to provide a list of the sites in order.
# The reason is that otherwise it's not possible for `Polyclonal` to figure out for instance if sites are just missing from the data are are actually deletions.
# 
# Here we read a data frame that maps sequential to reference site numbering, and then use that to extract the list of reference sites.
# Note this plot also contains the protein regions, which we will use later:

# In[5]:


site_numbering_map = pd.read_csv("BA.1_site_numbering_map.csv")
display(site_numbering_map.head())

print("Note how some reference sites differ from sequential ones due to indels:")
display(
    site_numbering_map[
        site_numbering_map["sequential_site"].astype(str)
        != site_numbering_map["reference_site"]
    ].head()
)

sites = site_numbering_map["reference_site"].tolist()


# Now initialize and fit the `Polyclonal` model, but pass the `sites` argument so we can use these non-sequential-integer reference sites.
# (If you don't pass the `sites` argument, sites are assumed to be sequential integer):

# In[6]:


model = polyclonal.Polyclonal(
    # `polyclonal` expects the concentration column to be named "concentration"
    data_to_fit=data_to_fit.rename(columns={"antibody_concentration": "concentration"}),
    n_epitopes=1,
    alphabet=polyclonal.AAS_WITHSTOP_WITHGAP,
    sites=sites,
)


# Note how the model has its `sequential_integer_sites` attribute set to `False`:

# In[7]:


assert set(model.sites) == set(sites)

assert model.sequential_integer_sites is False


# Now fit the model:

# In[8]:


# NBVAL_IGNORE_OUTPUT

_ = model.fit(
    logfreq=200,
    reg_escape_weight=0.1,
    fix_hill_coefficient=True,
    fix_non_neutralized_frac=True,
)


# Look at output.
# First we look at the `mut_escape_df`.
# The site entries are str:

# In[9]:


# NBVAL_IGNORE_OUTPUT

assert all(model.mut_escape_df["site"].astype(str) == model.mut_escape_df["site"])
assert set(model.mut_escape_df["site"]).issubset(sites)

model.mut_escape_df.head().round(2)


# The same is true for `mut_escape_site_summary_df`:

# In[10]:


# NBVAL_IGNORE_OUTPUT

assert all(
    model.mut_escape_site_summary_df()["site"].astype(str)
    == model.mut_escape_site_summary_df()["site"]
)
assert set(model.mut_escape_site_summary_df()["site"]).issubset(sites)

model.mut_escape_site_summary_df().head().round(2)


# Now plot the escape values.
# Note how we merge in the data frame with sequential site numbers and protein regions and use it to color the zoom bar and add sequential sites to the lineplot and heatmap tooltips:

# In[11]:


# NBVAL_IGNORE_OUTPUT

model.mut_escape_plot(
    df_to_merge=site_numbering_map.rename(columns={"reference_site": "site"}),
    addtl_tooltip_stats=["sequential_site"],
    site_zoom_bar_color_col="region",
    addtl_slider_stats={"times_seen": 2},
)


# ## Confirm same results for reference and sequential integer number
# To demonstrate how the results are the same regardless of which numbering scheme is used for the fitting, we also fit a model with the sequentially numbered variants.
# This section of the notebook can almost be considered a test rather than an example.

# In[12]:


# NBVAL_IGNORE_OUTPUT

model_sequential = polyclonal.Polyclonal(
    data_to_fit=(
        data_to_fit.drop(columns="aa_substitutions").rename(
            columns={
                "antibody_concentration": "concentration",
                "aa_substitutions_sequential": "aa_substitutions",
            }
        )
    ),
    n_epitopes=1,
    alphabet=polyclonal.AAS_WITHSTOP_WITHGAP,
)

assert model_sequential.sequential_integer_sites is True

_ = model_sequential.fit(
    logfreq=200,
    reg_escape_weight=0.1,
    fix_hill_coefficient=True,
    fix_non_neutralized_frac=True,
)


# Now make sure the fitting gives nearly the same result regardless of which site numbering is used.
# First check the activity values:

# In[13]:


pd.testing.assert_frame_equal(
    model_sequential.curve_specs_df,
    model.curve_specs_df,
    atol=0.02,
)


# Now compare the mutation-escape values.
# In order to do this, we have to re-number the sequential values to reference numbering:

# In[14]:


min_times_seen = 50

mut_escape = model.mut_escape_df.drop(columns="mutation")

mut_escape_sequential = model_sequential.mut_escape_df.assign(
    site=lambda x: x["site"].map(
        site_numbering_map.set_index("sequential_site")["reference_site"].to_dict()
    )
).drop(columns="mutation")

# have to use fairly big atol to test this, so also do correlations
pd.testing.assert_frame_equal(
    mut_escape,
    mut_escape_sequential,
    atol=1.5,
)
assert 0.99 < mut_escape["escape"].corr(mut_escape_sequential["escape"])


# ## Test of averaging
# Just as a toy test, average the reference-site model with itself and make sure it still displays reference-based site numbers:

# In[15]:


avg_model = polyclonal.PolyclonalAverage(
    pd.DataFrame({"number": [1, 2], "model": [model, model]}),
)


# In[16]:


# NBVAL_IGNORE_OUTPUT

avg_model.mut_escape_plot()


# You cannot average models that don't have the same site-numbering setup in terms of sequential or non-sequential:

# In[17]:


# NBVAL_RAISES_EXCEPTION

# this will give an error
polyclonal.PolyclonalAverage(
    pd.DataFrame({"number": [1, 2], "model": [model_sequential, model]}),
)