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

# # Visualization of signal flow in EARM with two different calibrated parameter sets
# 
# ## *This notebook reproduces the Figure 3 in the paper
# 
# ### We start by importing the libraries required to do the analysis and the visualizations

# In[ ]:


from pyvipr.examples_models.lopez_embedded import model
import pyvipr.pysb_viz as viz
import numpy as np
from pysb.simulator import ScipyOdeSimulator
import os
import matplotlib.pyplot as plt
from scipy.spatial.distance import pdist, squareform
from numpy import genfromtxt
get_ipython().run_line_magic('matplotlib', 'inline')


# # Obtaining maximally different parameter sets
# 
# ### First, we load the calibrated parameters obtained with the Particle Swarm Optimization algorithm. Next, we standardize the parameter set values using the StandardScaler() class from sklearn. Then, we applied an euclidean distance and choose the two parameters that are the most dissimilar.

# In[ ]:


# Standardization of calibrated parameters by removing the mean and scaling to unit variance

pars = np.load('data/calibrated_6572pars.npy')
kpars_idxs = [i for i, j in enumerate(model.parameters) if not j in model.parameters_initial_conditions()]
kpars = pars[:, kpars_idxs]
kpars_means = np.mean(kpars, axis=0)
kpars_means = np.tile(kpars_means, (kpars.shape[0], 1))
kpars_stds = np.std(kpars, axis=0)
kpars_stds = np.tile(kpars_stds, (kpars.shape[0], 1))

kpars_standard = (kpars - kpars_means) / kpars_stds


# In[ ]:


# Calculate the pairwise distances
dis = squareform(pdist(kpars_standard, metric='euclidean'))


# In[ ]:


# Obtain the two most dissimilar parameters
ind = np.unravel_index(np.argmax(dis, axis=None), dis.shape)

print(ind)
pars4697 = pars[4697]
pars5023 = pars[5023]


# ### The parameter sets with the index 4697 and 5023 are the most dissimilar.

# ### We load the experimental data to confirm that the calibrated parameters actually fit the data, and create a new function called display that plots the experimental data alongside the simulated trajectories of the calibrated parameter set.

# In[ ]:


# load experimental data

data_path = os.path.join(os.path.abspath(''), 'data',
                         'EC-RP_IMS-RP_IC-RP_data_for_models.csv')

exp_data = genfromtxt(data_path, delimiter=',', names=True)
# Mean and variance of Td (delay time) and Ts (switching time) of MOMP, and
# yfinal (the last value of the IMS-RP trajectory)
momp_data = np.array([9810.0, 180.0, model.parameters['Smac_0'].value])
momp_var = np.array([7245000.0, 3600.0, 1e4])


# In[ ]:


# Function to display the experimendal data and the calibrated simulations to observe goodness of fit
def display_exp_data_sims(position, save_name):
    param_values = position
    traj = solver.run(param_values=param_values)

    # normalize trajectories
    bid_traj = traj.observables['mBid'] / model.parameters['Bid_0'].value
    cparp_traj = traj.observables['cPARP'] / model.parameters['PARP_0'].value
    aSmac_traj = traj.observables['aSmac'] / model.parameters['Smac_0'].value

    # create all plots for each observable
    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(9, 3), sharex=True, sharey=True)

    fig.add_subplot(111, frameon=False)
    plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)
    plt.xlabel("Time(s)")
    plt.ylabel("Population")
    
    # plot cleaved parp
    ax[0].plot(tspan, bid_traj, color='r', marker='^', label='tBID sim')
    ax[0].errorbar(exp_data['Time'], exp_data['norm_ICRP'],
                 yerr=exp_data['nrm_var_ICRP'] ** .5,
                 ecolor='black', color='black', elinewidth=0.5, capsize=0)
    ax[0].legend(loc=2)

    # plot cleaved parp
    ax[1].plot(tspan, cparp_traj, color='blue', marker='*', label='cPARP sim')
    ax[1].errorbar(exp_data['Time'], exp_data['norm_ECRP'],
                 yerr=exp_data['nrm_var_ECRP'] ** .5,
                 ecolor='black', color='black', elinewidth=0.5, capsize=0)
    ax[1].legend(loc=2)

    # plot activated SMAC
    ax[2].plot(tspan, aSmac_traj, color='g', label='aSMAC sim')
    ax[2].axvline(momp_data[0], -0.05, 1.05, color='black', linestyle=':',
                label='exp aSMAC')
    ax[2].legend(loc=2)
    plt.show()
#     fig.savefig('{}.png'.format(save_name), dpi=500, bbox_inches='tight')
#     plt.close()


# In[ ]:


# Function to display experimental data of tBid and reproduce Figure 3.A
def display_figure3_A(positions, save_name):
    traj = solver.run(param_values=positions)

    # normalize trajectories
    bid_traj0 = traj.observables[0]['mBid'] / model.parameters['Bid_0'].value
    cparp_traj0 = traj.observables[0]['cPARP'] / model.parameters['PARP_0'].value
    aSmac_traj0 = traj.observables[0]['aSmac'] / model.parameters['Smac_0'].value
    
    bid_traj1 = traj.observables[1]['mBid'] / model.parameters['Bid_0'].value
    cparp_traj1 = traj.observables[1]['cPARP'] / model.parameters['PARP_0'].value
    aSmac_traj1 = traj.observables[1]['aSmac'] / model.parameters['Smac_0'].value

    # create all plots for each observable
    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5), sharex=True, sharey=True)

    fig.add_subplot(111, frameon=False)
    plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)
    plt.xlabel("Time(s)", fontsize=fz, labelpad=15)
    plt.ylabel("Population", fontsize=fz, labelpad=15)
    
    # plot tBid t=202.20
    ax[0].scatter(tspan[::2], bid_traj0[::2], color='r', marker='*', label='tBID par 1 ', alpha=0.8, linewidth=lw)
    ax[0].scatter(tspan[::3], bid_traj1[::3], color='blue', marker='+', label='tBID par 2', linewidth=lw)

    ax[0].errorbar(exp_data['Time'], exp_data['norm_ICRP'],
                 yerr=exp_data['nrm_var_ICRP'] ** .5,
                 ecolor='black', color='black', elinewidth=0.5, capsize=0, alpha=0.5, fmt='none')
    ax[0].legend(loc=2)
    ax[0].annotate("", xy=(202.20, 0), xytext=(202.20, 0.2), arrowprops=dict(arrowstyle="->", lw=5))
    
    # plot tBid t=4040.40
    ax[1].scatter(tspan[::2], bid_traj0[::2], color='r', marker='*', label='tBID par 1 ', alpha=0.8, linewidth=lw)
    ax[1].scatter(tspan[::3], bid_traj1[::3], color='blue', marker='+', label='tBID par 2', linewidth=lw)

    ax[1].errorbar(exp_data['Time'], exp_data['norm_ICRP'],
                 yerr=exp_data['nrm_var_ICRP'] ** .5,
                 ecolor='black', color='black', elinewidth=0.5, capsize=0, alpha=0.5, fmt='none')
    ax[1].legend(loc=2)
    ax[1].annotate("", xy=(4040.40, 0), xytext=(4040.40, 0.2), arrowprops=dict(arrowstyle="->", lw=5))

    # plot tBid t=7474.75
    ax[2].scatter(tspan[::2], bid_traj0[::2], color='r', marker='*', label='tBID par 1 ', alpha=0.8, linewidth=lw)
    ax[2].scatter(tspan[::3], bid_traj1[::3], color='blue', marker='+', label='tBID par 2', linewidth=lw)

    ax[2].errorbar(exp_data['Time'], exp_data['norm_ICRP'],
                 yerr=exp_data['nrm_var_ICRP'] ** .5,
                 ecolor='black', color='black', elinewidth=0.5, capsize=0, alpha=0.5, fmt='none')
    ax[2].legend(loc=2)
    ax[2].annotate("", xy=(7474.75, 0.2), xytext=(7474.75, 0.4), arrowprops=dict(arrowstyle="->", lw=5))

    ax[0].tick_params(axis='both', which='major', labelsize=fz)
    ax[1].tick_params(axis='both', which='major', labelsize=fz)
    ax[2].tick_params(axis='both', which='major', labelsize=fz)
#     fig.savefig('{}.png'.format(save_name), dpi=500, bbox_inches='tight')
    plt.show()


# In[ ]:


### Linewith and fontsize values to use for all plots
lw = 2.5 # Linewidth
fz = 14 # Fontsize


# In[ ]:


tspan = np.linspace(0, 20000, 100)
solver = ScipyOdeSimulator(model, tspan, compiler='python')


# In[ ]:


display_figure3_A([pars4697, pars5023], 'data_calibration_par4697_par5023_202_20')


# # Parameter 1 visualization: Bid-Bax dominant reaction
# 
# ### First, we make sure that the calibrated parameter actually fit the experimental data. To do this, we run a simulation with the parameter set 4697 and then plot the experimental data and the simulation results.

# In[ ]:


sim4697 = ScipyOdeSimulator(model, tspan, compiler='python').run(param_values=[pars4697])


# In[ ]:


display_exp_data_sims(pars4697, 'data_calibration_par4697')


# ### Then, we can explore how the concentrations of the different complexes on which mBid is involved changes over time. From this we can see that the concentration doesn't tell much information about how the information is flowing through the apoptosis network.

# In[ ]:


plt.plot(tspan, sim4697.all['__s37'], label='mBid', linewidth=2.5)
plt.plot(tspan, sim4697.all['__s43'], label='mBcl2-mBid', linewidth=2.5)
plt.plot(tspan, sim4697.all['__s44'], label='mBclxL-mbid', linewidth=2.5)
plt.plot(tspan, sim4697.all['__s41'], label='cBax-mBid', linewidth=2.5)
plt.xlabel('Time(s)', fontsize=14)
plt.ylabel('# Molecules', fontsize=14)
plt.legend(fontsize=14)
plt.savefig('earm_trajectories_example.png', dpi=200)


# ### By visually inspecting the experimental data and the simulation, we see that the simulation results fit the experimental data. Then, we proceed to explore the dynamics of the model with that parameter by using the function sp_dyn_view from pyvipr.
# 
# ### We focus on the dynamics of mitochondrial Bid because it plays an important role in the regulation of Mitochondrial Outer Membrane Permeabilization. We click the mBid node in the graph to highlight all the species to which mBid can interact.

# In[ ]:


viz.sp_comm_dyn_view(sim4697, layout_name='klay', random_state=1)


# ### For this parameter set, we observed that most of mBid was used to transport cytosolic Bax to MOM while no activation of Bak occurred, indicating that the pores in MOM were primarily made by Bax and that the model with this parameter set is particularly sensitive to Bax inhibition
# 
# ### To verify our visualization-based analysis, we carried out an in-silico experiment. We did a Bax knockout and ran a simulation of EARM with the parameter set 4697.

# In[ ]:


pars_bax_ko= np.copy(pars4697)
pars_mcl1_ko= np.copy(pars4697)
pars_bax_ko[63] = pars_bax_ko[63] * 0 # Setting the initial condition of Bax to zero
pars_mcl1_ko[57] = pars_mcl1_ko[57] * 0 # Setting the initial condition of Mcl1 to zero
sim4697_kd_bax = ScipyOdeSimulator(model, tspan, compiler='python').run(param_values=[pars_bax_ko, pars_mcl1_ko])


# In[ ]:


plt.scatter(tspan[:57], sim4697.all['cPARP'][:57], label='WT', linewidth=lw, marker='*', color='green')
plt.plot(tspan[:57], sim4697_kd_bax.all[1]['cPARP'][:57], label='Mcl1 KO', linewidth=lw)
plt.plot(tspan[:57], sim4697_kd_bax.all[0]['cPARP'][:57], label='Bax KO', linewidth=lw)
plt.xlabel('Time', fontsize=fz)
plt.xticks(fontsize=fz)
plt.ylabel('Molecules', fontsize=fz)
plt.yticks(fontsize=fz)
plt.title('Parameter set 1: cPARP', fontsize=22)
plt.legend(fontsize=14)
plt.show()
# plt.savefig('par4697_bax_ko.png', dpi=500, bbox_inches='tight')


# ### We found that the knockout protected cells from apoptosis induction with TRAIL, confirming that Bax has an essential role in apoptosis.

# # Parameter 2 visualization: Bid-Mcl1 dominant reaction
# 
# ### First, we make sure that the calibrated parameter actually fit the experimental data. To do this, we run a simulation with the parameter set 5023 and then plot the experimental data and the simulation results

# In[ ]:


sim5023 = ScipyOdeSimulator(model, tspan, compiler='python').run(param_values=pars5023)


# In[ ]:


display_exp_data_sims(pars5023, 'data_calibration_par5023')


# ### By visually inspecting the experimental data and the simulation, we see that the simulation results fit the experimental data. Then, we proceed to explore the dynamics of the model with that parameter set by using the function sp_dyn_view from pyvipr.

# In[ ]:


viz.sp_comm_dyn_view(sim5023, layout_name='klay', random_state=1)


# ### For the parameter set 2 we observed that mBid activity was inhibited primarily by the anti-apoptotic protein Mcl1, indicating that it plays an important role in throttling apoptosis
# 
# ### To verify our visualization-based analysis, we carried out another in-silico experiment. We did an Mcl1 knockout and ran a2 simulation of EARM with the parameter set 2.  

# In[ ]:


pars_for_mcl1_kd = np.copy(pars5023)
pars_for_bax_kd = np.copy(pars5023)
pars_for_mcl1_kd[57] = pars_for_mcl1_kd[57] * 0 # Setting the initial condition of Mcl1 to zero
pars_for_bax_kd[63] = pars_for_bax_kd[63] * 0   # Setting the initial condition of Bax to zero
sim5023_kd_bcl2 = ScipyOdeSimulator(model, tspan, compiler='python').run(param_values=[pars_for_mcl1_kd, pars_for_bax_kd])


# In[ ]:


plt.scatter(tspan[:57], sim5023.all['cPARP'][:57], label='WT', linewidth=lw, marker='*', color='green')
plt.plot(tspan[:57], sim5023_kd_bcl2.all[0]['cPARP'][:57], label='Mcl1 KO', linewidth=lw)
plt.plot(tspan[:57], sim5023_kd_bcl2.all[1]['cPARP'][:57], label='Bax KO', linewidth=lw)
plt.xticks(fontsize=fz)
plt.yticks(fontsize=fz)

plt.xlabel('Time', fontsize=fz)
plt.ylabel('Molecules', fontsize=fz)
plt.title('Parameter set 2: cPARP', fontsize=22)
plt.legend(fontsize=14)
plt.savefig('par5023_mcl1_ko.png', dpi=500, bbox_inches='tight')


# ### We found that time-to-death was reduced by 22.6%, corroborating that Mcl1 was delaying the apoptosis execution by binding to mBid.  

# ### These results demonstrated that although these two parameter sets fit the data equally well, they executed the apoptosis signal in different ways; specifically, in this case the parameter sets determined whether Bax or Mcl1 played the key role in regulating apoptosis execution.

# In[ ]: