# Imports
import os
import warnings
import numpy as np
import pandas as pd
import scipy as sp
import seaborn as sns
from matplotlib import pyplot as plt
from matplotlib import ticker as mticker

# Plotting colors
tol_muted_adjusted = [
    "#000000",
    "#CC6677", 
    "#1f78b4", 
    "#DDCC77", 
    "#117733", 
    "#882255", 
    "#88CCEE",
    "#44AA99", 
    "#999933", 
    "#AA4499", 
    "#EE7733",
    "#CC3311",
    "#DDDDDD",
]

# Seaborn style settings
sns.set(rc={
    "figure.dpi":300, 
    "savefig.dpi":300,
    "svg.fonttype":"none",
})
sns.set_style("ticks")
sns.set_palette(tol_muted_adjusted)

# Suppress warnings
warnings.simplefilter("ignore")

# this cell is tagged as `parameters` for papermill parameterization
validation_titers_path = None
functional_scores_path = None
out_dir = None
saved_image_path = None

# Parameters
validation_titers_path = "data/single_mutant_functional_validations.csv"
functional_scores_path = (
    "results/filtered_func_effect_CSVs/293T_filtered_func_effects.csv"
)
out_dir = "results/validation_plots/"
saved_image_path = "results/validation_plots/functional_validation_correlation.svg"

# # Uncomment for running interactive
# validation_titers_path = "../data/single_mutant_functional_validations.csv"
# functional_scores_path = "../results/filtered_func_effect_CSVs/293T_filtered_func_effects.csv"
# out_dir = "../results/validation_plots/"
# saved_image_path="../results/validation_plots/functional_validation_correlation.svg"

# Read titer data
validation_titers = pd.read_csv(validation_titers_path, na_filter=None)

# Read icXX model predictions
functional_scores_df = pd.read_csv(functional_scores_path)

# Create new column to match single mutant viruses
functional_scores_df["aa_substitutions"] = (
    functional_scores_df["wildtype"] + 
    functional_scores_df["site"].astype(str) + 
    functional_scores_df["mutant"]
)
functional_scores_df = (
    functional_scores_df.drop(columns=[
        "wildtype",
        "site",
        "mutant",
    ])
)

# Merge model predictions with measured titers
validation_vs_prediction = (
    validation_titers.merge(
        functional_scores_df,
        how="left",
        on=["aa_substitutions"],
        validate="many_to_one",
    )
    .rename(columns={"aa_substitutions" : "Amino acid substitutions"})
    .fillna(0)
)

# Rename dictionary
rename = {
    "Unmutated" : "unmutated (n=12)",
    "N89D" : "N89D (n=1)",
    "N119S" : "N119S (n=4)",
    "K125L" : "K125L (n=8)",
    "K126L" : "K126L (n=6)",
    "Y129L" : "Y129L (n=4)",
    "S135L" : "S135L (n=1)",
    "S138N" : "S138N (n=1)",
    "N148R" : "N148R (n=1)",
    "Q149H" : "Q149H (n=6)",
    "R257S" : "R257S (n=3)",
    "D401I" : "D401I (n=3)",
    "E404L" : "E404L (n=3)",
}

# Rename virus name
validation_vs_prediction["Amino acid substitutions"] = validation_vs_prediction["Amino acid substitutions"].map(rename)

# Calculate log10 of fold change for accuate r value calculation
validation_vs_prediction["log10 fold change"] = np.log10(validation_vs_prediction["Fold change relative to unmutated median"])

r, p = sp.stats.pearsonr(
    x=validation_vs_prediction["effect"], 
    y=validation_vs_prediction["log10 fold change"]
)
print(f"R={r}")
print(f"R^2={r**2}")

# Plot predicted vs measured
fig, ax = plt.subplots(figsize=(2.5, 2))
corr_chart = sns.scatterplot(
    data=validation_vs_prediction,
    x="effect", 
    y="Fold change relative to unmutated median",
    hue="Amino acid substitutions",
    ax=ax,
    alpha=0.8,
    edgecolor="black",
    s=30,
)
# Change all spines
for axis in ["top", "bottom", "left", "right"]:
    corr_chart.spines[axis].set_linewidth(1)
corr_chart.tick_params(axis="both", length=4, width=1)
sns.despine()
corr_chart.set_xlabel(
    "effect on cell entry measured\nby DMS (arbitrary units)", 
    # weight="bold",
    fontsize=8,
)
corr_chart.set_ylabel(
    "virus titer fold change\ncompared to unmutated GPC", 
    # weight="bold",
    fontsize=8,
)

corr_chart.set_yscale("log")
corr_chart.set_xlim(-2.25, 1.25)
corr_chart.set_ylim(0.005, 2.5)
corr_chart.set_xticks([-2, -1, 0, 1])
corr_chart.set_xticklabels(corr_chart.get_xticks(), size=8)
corr_chart.yaxis.set_minor_locator(mticker.NullLocator())  # no minor ticks
corr_chart.set_yticks([0.01, 0.1, 1])
corr_chart.set_yticklabels(["0.01", "0.1", "1"], size=8, horizontalalignment="right")
sns.move_legend(
    corr_chart, 
    "upper left", 
    bbox_to_anchor=(1, 1),
    fontsize=8,
    markerscale=1,
    handletextpad=0.1,
    title="amino acid\nsubstitutions",
    title_fontproperties = {
        "size" : 8, 
        "weight" : "bold",
    },
    frameon=False,
    borderaxespad=0.1,
)
corr_chart.text(
    -2, 
    1, 
    f"r={r:.2f}", 
    horizontalalignment="left",  
    weight="bold",
    fontsize=8,
)

# Add edges to legend markers to match scatter plot
for ha in corr_chart.legend_.legendHandles:
    ha.set_edgecolor("black")
    ha.set_linewidths(0.5)

# Set square ratio
corr_chart.set_box_aspect(1)

# Make output dir if doesn't exist
if not os.path.exists(out_dir):
    os.mkdir(out_dir)

# Save fig
plt.savefig(saved_image_path)

R=0.7606759679353398
R^2=0.578627928194366

Single mutant pseudovirus titers versus functional scores¶