# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: pytorch_new
#     language: python
#     name: pytorch_new
# ---

# +
# from utils import * 

import os
os.chdir("../../scVI/")
os.getcwd()

import pickle
import numpy as np
import pandas as pd
from copy import deepcopy
from statsmodels.stats.multitest import multipletests
import gseapy as gp

save_path = '../CSF/Notebooks/'
# -

import matplotlib
matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
import matplotlib.pyplot as plt


# %matplotlib inline

celllabels = np.load(save_path + 'meta/celllabels.npy')
celltypes, labels = np.unique(celllabels,return_inverse=True)

# +
# latent_u = np.load(save_path + 'UMAP/all_dataset.umap.npy')
# with open(save_path + 'dataset/all_datasets.pkl', 'rb') as f:
#     gene_dataset = pickle.load(f)

# logX = np.log10(1+gene_dataset.X.todense())
# scaling_factor = logX.mean(axis=1)
# norm_X = logX - scaling_factor.reshape(len(scaling_factor), 1)

isMS = np.load(save_path+'meta/isMS.npy')
isCSF = np.load(save_path+'meta/isCSF.npy')

# +
# with open(save_path + 'DE/control_tissue.DEres.pkl', 'rb') as f:
#     tissueDE,tissueDEclust = pickle.load(f)
# -

# clean = [1,2,3,4,5,6,7,8,10,11,13,14,17,18,19,20,21]
# temp=[x in clean for x in tissueDEclust]
# DEres = [tissueDE[i] for i,x in enumerate(temp) if x is True]
# DEclust = [tissueDEclust[i] for i,x in enumerate(temp) if x is True]
# genenames = pd.read_csv('../CSF/RFiles/genenames.csv',header=None)


# +
# celltype = []
# combinedDEres = []
# for i,x in enumerate(celltypes[DEclust]):
#     temp  = pd.read_csv('../CSF/DE/wilcoxon/MannWhitneyU.norm.tissue_control.%s.csv'%x)
#     fdr_wil = multipletests(temp['pvalue'],method='fdr_bh')
#     temp.index = list(genenames[0])
#     temp = pd.concat([temp,DEres[i]],axis=1,sort=True)
#     temp = temp.sort_values(by='bayes1',ascending=False)
#     combined = deepcopy(temp[['stat','pvalue','bayes1','bayes2','mean1','mean2','nonz1','nonz2','clusters']])
#     combined['scVI_logFC'] = np.log(temp['scale1']/temp['scale2'])
#     combined['norm_logFC'] = np.log(temp['norm_mean1']/temp['norm_mean2'])
#     temp  = pd.read_csv('../CSF/EdgeR/tissue_control.batch_corrected.%s.edgeR.csv'%x)
#     fdr_edgeR = multipletests(temp['PValue'],method='fdr_bh')
#     temp.index = list(genenames[0])
#     temp = temp[['logFC','logCPM','F','PValue']]
#     combined = pd.concat([combined,temp],axis=1,sort=True)
#     combined['fdr_wil'] = fdr_wil[1]
#     combined['fdr_edgeR'] = fdr_edgeR[1]
#     celltype.append(x)
#     combinedDEres.append(combined)

    
# with open(save_path + 'DE/tissue_control.combined.DEres.pkl', 'wb') as f:
#     pickle.dump((combinedDEres,celltype),f)

# +

with open(save_path + 'DE/tissue_control.combined.DEres.pkl', 'rb') as f:
    combinedDEres,celltype = pickle.load(f)

# +
# from pandas import ExcelWriter
# import xlsxwriter


# writer = pd.ExcelWriter(save_path + 'DE/tissue_control.xlsx', engine='xlsxwriter')
# for i, x in enumerate(celltype):
#     combinedDEres[i].to_excel(writer, sheet_name=str(x))
# writer.close()

# -

celltype

# # DE measure correlation between clusters

combinedDEres[0].columns

# +
plt.figure(figsize=(15,15))
corr = []
for k, stat in enumerate(['stat','pvalue','bayes1','bayes2','scVI_logFC','norm_logFC','logFC','F','PValue','fdr_wil','fdr_edgeR']):
    df = pd.concat([x[stat] for x in combinedDEres],axis=1)
    df.columns = celltype
    df = df[['B2','CD4','Tregs','CD8n','CD8a','Tdg','NK1','NK2','Gran','ncMono','pDC','mDC2']]
    ax = plt.subplot(4,4,(k+1))
    res = df.corr(method='spearman')
    corr.append(res)
    ax.matshow(res,vmin=0,vmax=1)
    plt.title(stat,fontsize=30)
    plt.xticks(color='w')
    plt.yticks(range(len(df.columns)), df.columns)
    plt.tight_layout()

plt.savefig(save_path + 'figures/SupFigure10/CSF_blood.DEcor.pdf')

# -

plt.matshow(res,vmin=0,vmax=1)
plt.colorbar()
plt.savefig(save_path + 'figures/SupFigure10/CSF_blood.colorbar.pdf')

# +
# from pandas import ExcelWriter
# import xlsxwriter


# writer = pd.ExcelWriter(save_path + 'DE/MSinCSF.xlsx', engine='xlsxwriter')
# for i, x in enumerate(celltype):
#     combinedDEres[i].to_excel(writer, sheet_name=str(x))
# writer.close()

# -

# # GSEA

gmt = ['c7.all.v6.2.symbols.gmt',
       'Hemato.geneset.gmt',
       'cell_cycle_Tirosh.gmt',
       'chenling.gmt',
       'netPath.gmt',
       'codeset.gmt',
       'signatures_NY_private.gmt']
gmt = ['../CSF/signatures/unmodified/'+x for x in gmt]


def DE_GSEA(filtered,gmt,celltypes,filename):
    gsea_res = []
    for x in filtered:
        if len(x)>0:
            res1 = gp.enrichr(gene_list=list(x.index),
                     gene_sets=['KEGG_2016'],
                     cutoff=0.5                    )
            res2 = gp.enrichr(gene_list=list(x.index),
                             gene_sets=gmt,
                             cutoff=0.5,
                             background=10266)
            res = pd.concat([res1.results,res2.results])
            res = res.sort_values(by='P-value')
            res['cluster'] = celltypes[list(x['clusters'])[0]]
            gsea_res.append(res.loc[res['Adjusted P-value']<0.05])
    gsea_res = pd.concat(gsea_res)
    gsea_res = gsea_res[['Adjusted P-value','Genes','Term','Overlap','cluster']]
#     gsea_res.to_csv(filename)
    return(gsea_res)



# # upregulated in CSF

# +
filtered = [x.loc[(x['fdr_wil']<0.2)&(x['fdr_edgeR']<0.2)& 
                  (x['bayes1']>0.3) &
                  (x['logFC']<0) & (x['norm_logFC']>0)] for x in combinedDEres]

DEgenes = pd.concat(filtered)
geneid, nocc = np.unique(DEgenes.index,return_counts=True)
shared = pd.DataFrame([geneid,nocc],index=['genename','occ']).T
shared = shared.sort_values(by='occ',ascending=False)

DEgenes['clusters'] = np.asarray(celltypes)[np.asarray(DEgenes['clusters']).astype(int)]
DEgenes = DEgenes.sort_index()
# DEgenes.to_csv(save_path+'DE/control_tissue.upregulated.relaxed.csv')
# -

DEgenes.loc[['FGF9','MT1E']]

DEgenes.loc[['CD1E','CCNC','CXCL16','CXCR5','CD48']]

up = DEgenes

# +
# # %%capture
# filtered_GSEA = DE_GSEA(filtered,gmt,celltypes,save_path+'DE/control_tissue.upregulated.relaxed.GSEA.csv')

# +
filtered = [x.loc[(x['fdr_wil']<0.05)&(x['fdr_edgeR']<0.05)& 
                  (x['bayes1']>0.5) &
                  (x['logFC']<0) & (x['norm_logFC']>0)] for x in combinedDEres]

DEgenes = pd.concat(filtered)
geneid, nocc = np.unique(DEgenes.index,return_counts=True)
shared = pd.DataFrame([geneid,nocc],index=['genename','occ']).T
shared = shared.sort_values(by='occ',ascending=False)

DEgenes['clusters'] = np.asarray(celltypes)[np.asarray(DEgenes['clusters']).astype(int)]
DEgenes = DEgenes.sort_index()
# DEgenes.to_csv(save_path+'DE/control_tissue.upregulated.csv')

# -



# # downregulated in MS

# +

filtered = [x.loc[(x['fdr_wil']<0.2)&(x['fdr_edgeR']<0.2)& 
                  (x['bayes2']>0.3) &
                  (x['logFC']>0) & (x['norm_logFC']<0)] for x in combinedDEres]


DEgenes = pd.concat(filtered)
geneid, nocc = np.unique(DEgenes.index,return_counts=True)
shared = pd.DataFrame([geneid,nocc],index=['genename','occ']).T
shared = shared.sort_values(by='occ',ascending=False)
DEgenes = DEgenes.sort_index()
DEgenes['clusters'] = np.asarray(celltypes)[np.asarray(DEgenes['clusters']).astype(int)]

# DEgenes.to_csv(save_path+'DE/control_tissue.downregulated.relaxed.csv')
# -

down = DEgenes

# %%capture
filtered_GSEA = DE_GSEA(filtered,gmt,celltypes,save_path+'DE/control_tissue.downregulated.relaxed.GSEA.csv')

filtered = [x.loc[(x['fdr_wil']<0.05)&(x['fdr_edgeR']<0.05)& 
                  (x['bayes2']>0.5) &
                  (x['logFC']>0) & (x['norm_logFC']<0)] for x in combinedDEres]
DEgenes = pd.concat(filtered)
geneid, nocc = np.unique(DEgenes.index,return_counts=True)
shared = pd.DataFrame([geneid,nocc],index=['genename','occ']).T
shared = shared.sort_values(by='occ',ascending=False)
DEgenes = DEgenes.sort_index()
DEgenes['clusters'] = np.asarray(celltypes)[np.asarray(DEgenes['clusters']).astype(int)]
# DEgenes.to_csv(save_path+'DE/control_tissue.downregulated.csv')


filtered_GSEA

# %%capture
filtered_GSEA = DE_GSEA(filtered,gmt,celltypes,save_path+'DE/control_tissue.downregulated.GSEA.csv')

allgenes = pd.concat(combinedDEres)

# +
genes, occ = np.unique(allgenes.index,return_counts=True)

np.sum(occ==1)/10266
# -

np.sum(occ==2)/10266

allgenes.loc[['CCNC']]



StarcoderdataJupyterScriptsDedupFiltered