import scanpy as sc
import numpy as np
import pandas as pd
import os
from scipy import sparse
import sys
import importlib
import VITAE
import seaborn

import random
import tensorflow as tf

## Preprocess datasets from "Molecular logic of cellular diversification in the mouse cerebral cortex"
## One can download the dataset and meta data through the URL:
## https://singlecell.broadinstitute.org/single_cell/study/SCP1290/molecular-logic-of-cellular-diversification-in-the-mammalian-cerebral-cortex
## transoform.h5ad file we used means data has been normalized to 10000 and done log(x+1) transformation.


###################################### Preprocess Di Bella's dataset #################################################
dd=sc.read("transform.h5ad")
metadata = pd.read_csv('metaData_scDevSC.txt', delimiter='\t', index_col = 0)
dd.obs['Day'] = metadata['orig_ident'][1:metadata.shape[0]]
dd.obs['Clusters'] = metadata['New_cellType'][1:metadata.shape[0]]
dd.obs['Clusters'] = pd.Categorical(dd.obs['Clusters'], categories = [
    'Apical progenitors', 'Intermediate progenitors', 'Migrating neurons',
    'Immature neurons', 'Cajal Retzius cells', 'CThPN', 'SCPN',
    'NP', 'Layer 6b', 'Layer 4', 'DL CPN', 'DL_CPN_1', 'DL_CPN_2', 'UL CPN',
    'Interneurons', 'Astrocytes', 'Oligodendrocytes', 'Microglia',
    'Cycling glial cells', 'Ependymocytes', 'Endothelial cells',
    'VLMC', 'Pericytes','Red blood cells', 'Doublet', 'Low quality cells'
    ], ordered = True)
dd.obs['S_Score'] = pd.to_numeric(metadata['S_Score'][1:metadata.shape[0]])
dd.obs['G2M_Score'] = pd.to_numeric(metadata['G2M_Score'][1:metadata.shape[0]])
dd = dd[dd.obs['Clusters'].isin(['Doublet', 'Low quality cells']) == False]
dd.obs.index=dd.obs.index.tolist()
dd.obs['Day']=dd.obs['Day'].tolist()
dd.obs['Clusters']=dd.obs['Clusters'].tolist()
dd.obs['S_Score']=dd.obs['S_Score'].tolist()
dd.obs['G2M_Score']=dd.obs['G2M_Score'].tolist()

sc.pp.highly_variable_genes(dd, flavor = "seurat")
sc.pp.scale(dd, max_value=10)

day18 = np.array([np.nan] * dd.shape[0]).astype(object)
day18[(dd.obs["Day"] == "E18").values] = "E18"
day18[(dd.obs["Day"] == "E18_S1").values] = "E18_S1"
day18[(dd.obs["Day"] == "E18_S3").values] = "E18_S3"
dd.obs["merge_day_18"] = day18

dayP1 = np.array([np.nan] * dd.shape[0]).astype(object)
dayP1[dd.obs["Day"] == "P1"] = "P1"
dayP1[dd.obs["Day"] == "P1_S1"] = "P1_S1"
dd.obs["merge_P1"] = dayP1

dd.obs["merge_day_18"] = dd.obs["merge_day_18"].astype("category")
dd.obs["merge_P1"] = dd.obs["merge_P1"].astype("category")

dd.obs["Source"] = "Di Bella"



###################################### Preprocess Yuzwa's and Ruan's merged dataset #################################################
mouse = load_data(path='data/',file_name="mouse_brain_merged")
sc.pp.normalize_total(mouse, target_sum=1e4)
sc.pp.log1p(mouse)
sc.pp.highly_variable_genes(mouse, min_mean=0.0125, max_mean=3, min_disp=0.5)
sc.pp.scale(mouse, max_value=10)
sc.tl.pca(mouse, n_comps = 64)
sc.pp.neighbors(mouse, n_neighbors = 20)
sc.tl.umap(mouse)

mouse_source = np.repeat("Yuzwa",mouse.shape[0])
mouse_source[mouse.obs["covariate_2"]==0] = "Ruan"
mouse.obs["Source"] = mouse_source
sc.pl.umap(mouse,color="Source",save = "_mouse_Source.pdf")

mouse_day = mouse.obs.index.values.copy()
mouse_day = [x[:3] for x in mouse_day]
mouse.obs["Day"] = mouse_day
mouse.obs.rename(columns={"grouping": "Clusters", "covariate_0": "S_Score",
                     "covariate_1":"G2M_Score"}, inplace = True)



###################################### Merge two datasets  #################################################

dd = dd.concatenate(mouse,join="inner")
## define highly_variable in merged dataset
dd.var["highly_variable"] = dd.var["highly_variable-0"].astype(bool) | dd.var["highly_variable-1"].astype(bool)

## union different name of cell types in two datasets
group_dict = {"Immature Neuron" : "Immature neurons",
             "NEC":"Apical progenitors",
             "RGC":"Apical progenitors",
              "Layer I":"Cajal Retzius cells",
             "OPC":"Oligodendrocytes",
             "Interneurons":"Interneurons",
             "Endothelial Cell":"Endothelial cells",
             "Microglia":"Microglia",
             "Pericyte":"Pericytes",
             "Intermediate progenitors":"IPC"}

c = dd.obs["Clusters"].values.copy()
c = [x if group_dict.get(x) == None else group_dict.get(x) for x in c]
dd.obs["tidy_clusters"] = c.copy()

## Add covariates for different source
a = np.zeros(dd.shape[0])
a[dd.obs["Source"] == "Ruan"] = 1
dd.obs["cov1"] = a

a = np.zeros(dd.shape[0])
a[dd.obs["Source"] == "Yuzwa"] = 1
dd.obs["cov2"] = a

a = np.zeros(dd.shape[0])
a[dd.obs["Source"] == "Di Bella"] = 1
dd.obs["cov3"] = a

a = np.array([np.nan] * dd.shape[0])
a[dd.obs["Day"] == "E18"] = 1
a[dd.obs["Day"] == "E18_S1"] = 2
a[dd.obs["Day"] == "E18_S3"] = 3
dd.obs["merge_18"] = a


a = np.array([np.nan] * dd.shape[0])
a[dd.obs["Day"] == "P1"] = 1
a[dd.obs["Day"] == "P1_S1"] = 2
dd.obs["merge_P1"] = a

dd.obs["merge_18"] = dd.obs["merge_18"].astype("category")
dd.obs["merge_P1"] = dd.obs["merge_P1"].astype("category")

a = np.array([np.nan] * dd.shape[0]).astype(object)
a[dd.obs["merge_P1"] == 1] = "P1"
a[dd.obs["merge_P1"] == 2] = "P1_S1"
dd.obs["merge_day_P1"] = a

a = dd.obs["tidy_clusters"].values.copy().astype(str)
a[(dd.obs["Day"].isin(["E14","E15","E16"])) & (a == "SCPN")] = "SCPN1"
dd.obs["Cluster2"] = a.copy()

###################################### Naive Merge  #################################################
## See what happened if we directly merge them ##
dd.obs["Source"] = dd.obs["Source"].astype(pd.CategoricalDtype(
                  ['Ruan', 'Yuzwa','Di Bella'], ordered=True))

sc.tl.pca(dd, n_comps = 64)
sc.pp.neighbors(dd, n_neighbors = 20)
sc.tl.umap(dd)
seaborn.set(rc={'figure.figsize':(6,6)},style = "white")
sc.pl.umap(dd,color="Source",save="_naive_Source.pdf",size = 5,palette= ["#F8766D","#00BFC4","#E5E8E8"],alpha = 0.5,na_color="#E5E8E8")

color_ggplot=["#F8766D","#E38900","#C49A00","#99A800","#53B400","#00BC56","#00C094","#00BFC4","#00B6EB","#06A4FF","#A58AFF","#DF70F8","#FB61D7","#FF66A8"]
seaborn.set(rc={'figure.figsize':(5,5)},style = "white")
sc.pl.umap(dd,color="Day",size = 2, alpha=0.4,palette=color_ggplot,save="_naive_merge_day.pdf")


###################################### Train VITAE  #################################################


seed = 400
tf.keras.backend.clear_session()
random.seed(seed)
np.random.seed(seed)
tf.random.set_seed(seed)

model = VITAE.VITAE(adata = dd,
    #        adata_layer_counts = 'counts',
            hidden_layers = [48, 24],
            latent_space_dim = 16,
            model_type = 'Gaussian',covariates = ['S_Score', 'G2M_Score',"cov1","cov2"],npc = 96)


model.pre_train(learning_rate = 1e-3,early_stopping_tolerance = 0.01, early_stopping_relative = True,L=1,phi=1,verbose=False,gamma = 1)
model.init_latent_space(cluster_label= 'Cluster2',ratio_prune=0.45)
model.train(beta = 1, learning_rate = 1e-3, early_stopping_tolerance = 0.01, early_stopping_relative = True,gamma=0,phi=1, verbose=False)
model.posterior_estimation(batch_size=32, L=10)
model.save_model(path_to_file="model.checkpoint")