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

# # Trajectory-wise analyses
# 
# In this tutorial, we will use trajectory wise analysis to expore regulatory mechanisms of hematopoiesis.

# In[1]:


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# import Scribe as sb
import sys
import os

# import scanpy as sc
import dynamo as dyn
import seaborn as sns

# filter warnings for cleaner tutorials
import warnings
warnings.filterwarnings('ignore')

dyn.dynamo_logger.main_silence()


# In[2]:


adata_labeling = dyn.sample_data.hematopoiesis()


# take a glance at what is in `adata` object. All observations, embedding layers and other data in `adata` are computed within `dynamo`. Please refer to other dynamo tutorials regarding how to obtain these values from metadata and raw new/total and (or) raw spliced/unspliced gene expression values.

# In[3]:


adata_labeling


# In[4]:


pca_dim = [0, 1]
dyn.pl.streamline_plot(
    adata_labeling,
    color="cell_type",
    use_smoothed=False,
    cmap="bwr",
    basis="pca",
    sym_c=True,
    x=pca_dim[0],
    y=pca_dim[1],
    frontier=True,
    sort="abs",
    alpha=0.2,
    pointsize=0.1,
    despline=True,
    despline_sides=["right", "top"],
    deaxis=False,
    save_show_or_return="return",
)


# In[5]:


HSC_cells = dyn.tl.select_cell(adata_labeling, "cell_type", "HSC")
Meg_cells = dyn.tl.select_cell(adata_labeling, "cell_type", "Meg")
Ery_cells = dyn.tl.select_cell(adata_labeling, "cell_type", "Ery")
Bas_cells = dyn.tl.select_cell(adata_labeling, "cell_type", "Bas")
Mon_cells = dyn.tl.select_cell(adata_labeling, "cell_type", "Mon")
Neu_cells = dyn.tl.select_cell(adata_labeling, "cell_type", "Neu")


# In[6]:


# dyn.pd.fate(adata_labeling, init_cells=adata_labeling.obs_names[Mon_cells[0]], basis='pca', direction='backward', interpolation_num=500)
dyn.pd.fate(
    adata_labeling,
    init_cells=adata_labeling.obs_names[Meg_cells[0]],
    basis="pca",
    direction="backward",
    interpolation_num=500,
)


# In[7]:


adata_labeling.uns["fate_pca"]


# In[8]:


adata_labeling.uns["fate_pca"]["prediction"][0]


# In[9]:


dyn.pd.fate(
    adata_labeling,
    # init_cells=adata_labeling.obs_names[Meg_cells[0]],
    init_cells=adata_labeling.obs_names[Meg_cells[2]],
    basis="pca",
    direction="backward",
    interpolation_num=500,
)


# In[10]:


adata_labeling.uns["fate_pca"].keys(), adata_labeling.uns["fate_pca"]


# In[11]:


adata_labeling.uns["fate_pca"]["prediction"][0].shape, len(
    adata_labeling.uns["fate_pca"]["t"][0],
)


# In[12]:


from anndata import AnnData
from scipy.sparse import csr_matrix
from dynamo.vectorfield.utils import vector_transformation


path_id = 0
dyn.tools.utils.nearest_neighbors(
    adata_labeling.uns["fate_pca"]["prediction"][path_id].T[0, :], adata_labeling.obsm["X_pca"], 5
)[0][1]

dyn.tools.utils.nearest_neighbors(
    adata_labeling.uns["fate_pca"]["prediction"][path_id].T[140, :], adata_labeling.obsm["X_pca"], 5
)


vec_dict, vecfld = dyn.vf.utils.vecfld_from_adata(adata_labeling, basis="pca")


vector_field_class = dyn.vf.SvcVectorField()
vector_field_class.from_adata(adata_labeling, basis="pca")

X_data = adata_labeling.uns["fate_pca"]["prediction"][path_id].T[:140, :]
vel_norm = vector_field_class.func(X_data)
acc_norm, acc_mat = vector_field_class.compute_acceleration(X=X_data)
curv_norm, curv_mat = vector_field_class.compute_curvature(X=X_data)
div = vector_field_class.compute_divergence(X=X_data)

Jac_func = vector_field_class.get_Jacobian()
Js = Jac_func(X_data)

X_data.shape, Js.shape, adata_labeling.uns["fate_pca"]["prediction"][path_id].T.shape
adata_labeling.uns["PCs"].shape, adata_labeling.uns["pca_mean"].shape
Jacobian = dyn.vf.utils.subset_jacobian_transformation(Js, adata_labeling.uns["PCs"], adata_labeling.uns["PCs"])


# In[13]:


# project the expression state back to higher dimension
exprs = dyn.prediction.utils.pca_to_expr(X_data, adata_labeling.uns["PCs"], mean=adata_labeling.uns["pca_mean"])
exprs.shape


# In[14]:


trajectory_adata = AnnData(
    X=exprs, layers={"M_t": csr_matrix(exprs)}, var=adata_labeling[:, adata_labeling.var.use_for_pca].var
)

trajectory_adata.layers["velocity"] = csr_matrix(vector_transformation(vel_norm, adata_labeling.uns["PCs"]))
trajectory_adata.layers["acceleration"] = vector_transformation(acc_mat, adata_labeling.uns["PCs"])
trajectory_adata.layers["curvature"] = vector_transformation(curv_mat, adata_labeling.uns["PCs"])

adata_labeling.uns["jacobian_pca"].keys(), X_data.shape[0]
trajectory_adata.uns["jacobian_pca"] = {
    "cell_idx": np.arange(X_data.shape[0]),
    "effectors": trajectory_adata.var_names,
    "jacobian": Js,
    "jacobian_gene": Jacobian,
    "regulators": trajectory_adata.var_names,
}


# In[15]:


trajectory_adata.obs["integral_time"] = adata_labeling.uns["fate_pca"]["t"][path_id][: X_data.shape[0]]


# In[16]:


dyn.pp.recipe_monocle(trajectory_adata)


# In[17]:


trajectory_adata.var_names


# In[18]:


dyn.tl.reduceDimension(trajectory_adata)
dyn.pl.umap(trajectory_adata, color="SPI1")


# In[19]:


genes = ["SPI1", "GATA1"]
integral_time = trajectory_adata.obs.integral_time[::-1]
expression = trajectory_adata[:, genes].layers["M_t"].A[::-1, :]
velocity = trajectory_adata[:, genes].layers["velocity"].A[::-1, :]
acceleration = trajectory_adata[:, genes].layers["acceleration"].A[::-1, :]

fig1, f1_axes = plt.subplots(ncols=3, nrows=1, constrained_layout=True, figsize=(12, 4))
sns.scatterplot(integral_time, acceleration[:, 0], s=2, ax=f1_axes[0], ec=None, label="SPI1")
sns.scatterplot(integral_time, acceleration[:, 1], s=2, ax=f1_axes[0], ec=None, label="GATA1")
f1_axes[0].set_title("acceleration")
f1_axes[0].set_xlabel("integration time")
sns.scatterplot(integral_time, velocity[:, 0], s=2, ax=f1_axes[1], ec=None, label="SPI1")
sns.scatterplot(integral_time, velocity[:, 1], s=2, ax=f1_axes[1], ec=None, label="GATA1")
f1_axes[1].set_title("velocity")
f1_axes[1].set_xlabel("integration time")
sns.scatterplot(integral_time, expression[:, 0], s=2, ax=f1_axes[2], ec=None, label="SPI1")
sns.scatterplot(integral_time, expression[:, 1], s=2, ax=f1_axes[2], ec=None, label="GATA1")
f1_axes[2].set_title("expression")
f1_axes[2].set_xlabel("integration time")


# In[20]:


dyn.configuration.set_pub_style()
regulators, effectors = list(adata_labeling.uns["jacobian_pca"]["regulators"]), list(
    adata_labeling.uns["jacobian_pca"]["effectors"]
)
spi1_ind, gata1_ind = regulators.index("SPI1"), effectors.index("GATA1")
fig1, f1_axes = plt.subplots(ncols=4, nrows=1, constrained_layout=True, figsize=(8, 2))
expression = trajectory_adata[:, ["SPI1", "GATA1"]].layers["M_t"].A[::-1, :]

sns.scatterplot(
    expression[:, 0],
    Jacobian[gata1_ind, spi1_ind, ::-1],
    hue=integral_time * 3,
    palette="viridis",
    ax=f1_axes[0],
    ec=None,
    s=2,
)
f1_axes[0].set_title(r"$\partial f_{GATA1}/\partial x_{SPI1}$")
f1_axes[0].set_xlabel(r"SPI1 ($M_t$)")
sns.scatterplot(
    expression[:, 1],
    Jacobian[spi1_ind, gata1_ind, ::-1],
    hue=integral_time * 3,
    palette="viridis",
    ax=f1_axes[1],
    ec=None,
    s=2,
)
f1_axes[1].set_title(r"$\partial f_{SPI1}/\partial x_{GATA1}$")
f1_axes[1].set_xlabel(r"GATA1 ($M_t$)")
sns.scatterplot(
    expression[:, 0],
    Jacobian[spi1_ind, spi1_ind, ::-1],
    hue=integral_time * 3,
    palette="viridis",
    ax=f1_axes[2],
    ec=None,
    s=2,
)
f1_axes[2].set_title(r"$\partial f_{SPI1}/\partial x_{SPI1}$")
f1_axes[2].set_xlabel(r"SPI1 ($M_t$)")
sns.scatterplot(
    expression[:, 1],
    Jacobian[gata1_ind, gata1_ind, ::-1],
    hue=integral_time * 3,
    palette="viridis",
    ax=f1_axes[3],
    ec=None,
    s=2,
)
f1_axes[3].set_title(r"$\partial f_{GATA1}/\partial x_{GATA1}$")
f1_axes[3].set_xlabel(r"GATA1 ($M_t$)")


# In[21]:


dyn.configuration.set_pub_style(scaler=2)
dyn.pl.response(
    adata_labeling,
    np.array([["SPI1", "SPI1"], ["SPI1", "GATA1"], ["GATA1", "SPI1"], ["GATA1", "GATA1"]]),
    ykey="alpha",
    log=False,
    drop_zero_cells=False,
    grid_num=20,
)