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

# In[1]:


import math

import ipywidgets as widgets
import numpy as np
from matplotlib import pyplot as plt

import wot


# # Notebook 5: Fate Matrices
# 
# In this notebook we show how to compute and visualize 'fates' of cells. 
# Consider a pair of time points $t_i < t_j$ and a collection of cell sets $C_1,\ldots,C_k$ spanning all cells at the later time point, $t_j$. (That is, each cell $y$ from time $t_j$ is in some $C_\ell$). 
# We show how to compute the probability that a cell $x$ from time $t_i$ will transition to a target destination $C_\ell$ at time $t_j$. 
# We call these the *fate probabilities* for cell $x$. The **fate matrix** $F_{t_i,t_j}$ is a matrix with a row containing the fate probabilities for each cell $x$ from time $t_i$.
# 
# ## Compute fates
# We begin by loading into memory
# * a `TransportMapModel` from a directory of pre-computed transport maps, and 
# * a dictionary of cell sets. 

# In[2]:


# Load transport map model and cell sets
tmap_model = wot.tmap.TransportMapModel.from_directory('tmaps/serum')
cell_sets = wot.io.read_sets('data/major_cell_sets.gmt', as_dict=True)


# Next, we initialize a list of `target_destinations` at time $t_j = 18$. 

# In[3]:


# create indicator vectors for each cell set
target_destinations = tmap_model.population_from_cell_sets(cell_sets, at_time=18)


# Finally, we compute the fate matrices for each earlier time point $t_i < t_j$ all at once with the  `fates` function. 

# In[4]:


fate_ds = tmap_model.fates(target_destinations)


# The resulting dataset `fate_ds` contains the fate matrices for each earlier time point $t_i < t_j$. We can then easily subset this dataset to access the fate matrix $F_{t_i,t_j}$ for any individual day $t_i$.
# 
# ## Visualizing fate matrices
# 
# In the following code blocks we provide two simple tools to visualize fate matrices. First, we plot a fate matrix with $k=3$ fates in barycentric coordinates. This is a handy coordinate system for viewing the three dimensional simplex in two dimensions. A probability vector $(p_1,p_2,p_3)$ is represented inside an equilateral triangle, where each vertex represents a pure fate like $(1,0,0)$ or $(0,1,0)$ or $(0,0,1)$.

# In[5]:


fate_dropdown1 = widgets.Dropdown(
    options=fate_ds.var.index,
    description='Fate 1:'
)
fate_dropdown2 = widgets.Dropdown(
    options=fate_ds.var.index,
    description='Fate 2:',
    value=fate_ds.var.index[1]
)
day_dropdown = widgets.Dropdown(
    options=fate_ds.obs['day'].unique(),
    description='Day',
    value=12
)


def update_fate_vis(name1, name2,day):
    figure = plt.figure(figsize=(10, 10))   

    fate1 = fate_ds[:,name1][fate_ds.obs['day']==day].X.flatten()
    fate2 = fate_ds[:,name2][fate_ds.obs['day']==day].X.flatten()

    Nrows = len(fate1)
    x = np.zeros(Nrows)
    y = np.zeros(Nrows)
    P = np.array([[1,0],[np.cos(2*math.pi/3),math.sin(2*math.pi/3)],[math.cos(4*math.pi/3),math.sin(4*math.pi/3)]])

    for i in range(0,Nrows):
        ff = np.array([fate1[i],fate2[i],1-(fate1[i]+fate2[i])])
        x[i] = (ff @ P)[0]
        y[i] = (ff @ P)[1]

    vx = P[:,0]
    vy = P[:,1]
    t1 = plt.Polygon(P, color=(0,0,0,0.1))
    plt.gca().add_patch(t1)
    
    plt.scatter(x,y)
    plt.scatter(vx,vy)
    #plt.scatter(0,0)
    plt.text(P[0,0]+.1, P[0,1], name1)
    plt.text(P[1,0]-.1, P[1,1]+.1, name2)
    plt.text(P[2,0]-.1, P[2,1]-.2, 'Other')
    plt.axis('equal')
    plt.axis('off')
    #plt.xticks([])
    #plt.yticks([])
    
    plt.title('{} vs. {} on day {}'.format(name1, name2,day))

widgets.interact(update_fate_vis, name1=fate_dropdown1, name2=fate_dropdown2,day=day_dropdown)


# ## Plotting log-odds for a pair of fates
# 
# Our second visualization tool displays the log-odds for a pair of fates (the odds are defined as the ratio of probabilities, and then we take a log). We plot the log-odds on the y-axis and use the x-axis to show time. In this way, we show the relative fate probabilities for each earlier time $t_i$ for a fixed pair of target destinations. 

# In[6]:


fate_dropdown1 = widgets.Dropdown(
    options=fate_ds.var.index,
    description='Fate 1:'
)
fate_dropdown2 = widgets.Dropdown(
    options=fate_ds.var.index,
    description='Fate 2:',
    value=fate_ds.var.index[1]
)

def update_fate_vis(name1, name2):
    figure = plt.figure(figsize=(10, 10))
    fate1 = fate_ds[:, name1].X
    fate2 = fate_ds[:, name2].X
    p = np.log(1e-9 + np.divide(fate1, fate2, out=np.zeros_like(fate1), where=fate2 != 0))
    plt.scatter(fate_ds.obs['day'], p, s=4, marker=',')
    plt.xlabel('Day')
    plt.ylabel('Log Odds')
    plt.title('{} vs. {}'.format(name1, name2))

widgets.interact(update_fate_vis, name1=fate_dropdown1, name2=fate_dropdown2)