#CellTypistに必要
!pip install celltypist

#decouplerに必要
!pip install decoupler
!pip install omnipath

#umapに必要
!pip3 install igraph
!pip3 install leidenalg
!pip install python-igraph louvain

#umapに必要
!pip3 install igraph
!pip3 install leidenalg
!pip install python-igraph louvain

# ScanpyとAnnData
!pip install scanpy anndata

from google.colab import drive
drive.mount('/content/drive')

import scanpy as sc
import anndata as ad
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import scipy.io

import celltypist
from celltypist import models

import decoupler as dc


# Google Drive上のデータパスを指定
data_dir = '/content/drive/MyDrive/???' #???はデータを置いたフォルダのpathを指定してください

barcodes_path = f'{data_dir}/barcodes.tsv'
features_path = f'{data_dir}/features.tsv'
matrix_path = f'{data_dir}/matrix.mtx'

# バーコード、特徴量、マトリックスの読み込み
barcodes = pd.read_csv(barcodes_path, header=None)
features = pd.read_csv(features_path, header=None, sep='\\t')
matrix = scipy.io.mmread(matrix_path)

# AnnDataオブジェクトの作成
adata = sc.AnnData(X=matrix.T.tocsr())
adata.obs_names = barcodes[0].values #細胞名　obs.indexとおなじ
adata.var_names = features[1].values #遺伝子名 var.indexとおなじ


# 遺伝子の重複を取り除く
new_data = adata.to_df()
new_data = new_data.groupby(new_data.columns, axis=1).sum() 
adata =  sc.AnnData(X=new_data.values)
adata.obs_names = new_data.index
adata.var_names = new_data.columns

# Quality Control

def QC(adata):
    # ミトコンドリア遺伝子の発現割合を計算してobsに追加
    percent_mt = adata[:,mt_genes].X.sum(axis=1) / adata.X.sum(axis=1) * 100
    # 各細胞におけるユニークな遺伝子の数（非ゼロの遺伝子数）
    nFeature_RNA = (adata.X > 0).sum(axis=1)
    # 各細胞における全てのRNA分子の総数
    nCount_RNA = adata.X.sum(axis=1)

    df_violin_QC =\\
    pd.DataFrame({"nFeature_RNA":nFeature_RNA, "nCount_RNA":nCount_RNA, "percent_mt":percent_mt},\\
                index = adata.obs_names)
    return df_violin_QC
    
df_violin_QC = QC(adata)

adata_QC = adata[(200< df_violin_QC['nFeature_RNA'])\\
                 & (df_violin_QC['nFeature_RNA'] < 5000)\\
                 & (df_violin_QC['percent_mt'] < 15)]

violin_adata = sc.AnnData(df_violin_QC)

# プロット1: バイオリンプロット
sc.pl.violin(violin_adata, ["nFeature_RNA", "nCount_RNA", "percent_mt"], jitter=0.4, multi_panel=True)}}}

# プロット2: 散布図
corr_nCount_percent_mt = np.corrcoef(df_violin_QC['nCount_RNA'], df_violin_QC['percent_mt'])[0,1]
corr_nCount_nFeature_RNA = np.corrcoef(df_violin_QC['nCount_RNA'], df_violin_QC['nFeature_RNA'])[0, 1]

sc.pl.scatter(violin_adata, x='nCount_RNA', y='percent_mt', title=f'nCount_RNA vs percent_mt (corr: {corr_nCount_percent_mt:.2f})')
sc.pl.scatter(violin_adata, x='nCount_RNA', y='nFeature_RNA', title=f'nCount_RNA vs nFeature_RNA (corr: {corr_nCount_nFeature_RNA:.2f})')


# 正規化
sc.pp.normalize_total(adata_QC, target_sum=1e4)
sc.pp.log1p(adata_QC) # celltypistはlog1pで変換されていることが前提になっている。

# 高可変性遺伝子の検出
sc.pp.highly_variable_genes(adata_QC, flavor='seurat_v3', n_top_genes=2000)

# 高変動遺伝子のマスクを作成
highly_variable_mask = adata_QC.var['highly_variable']
# 高変動遺伝子のみをスケーリング
adata_highly_variable = adata_QC[:, highly_variable_mask].copy()
sc.pp.scale(adata_highly_variable)

# 上位10個の高変動遺伝子を取得
top10 = adata_QC.var[adata_QC.var['highly_variable']].sort_values(by='highly_variable_rank').index[:10]

# プロットの作成
plt.figure(figsize=(10, 6))
ax = sc.pl.highly_variable_genes(adata_QC, show=False)

# 上位10個の遺伝子にラベルを付ける
for gene in top10:
    idx = adata_QC.var.index.get_loc(gene)
    ax.annotate(gene, (adata_QC.var['means'][idx], adata_QC.var['variances'][idx]),
                fontsize=7, color='red', ha='right')

plt.show()

# CellTypistはスケーリング・次元削減前のデータを用いるので、ここで保存する。
# スケーリング前のデータを保存
adata_QC.layers['before_DimRed'] = adata_QC.X.copy()

# スケーリングされた高変動遺伝子のデータを元のデータに戻す
adata_QC[:, highly_variable_mask].X = adata_highly_variable.X

# PCAの実行
sc.tl.pca(adata_QC,svd_solver='arpack')

# UMAPの実行
sc.pp.neighbors(adata_QC, n_pcs=20)  # UMAPの前に隣接グラフを作成する必要があります
sc.tl.umap(adata_QC)

# 隣接グラフの構築
sc.pp.neighbors(adata_QC, n_pcs=10)

# クラスタリングの実行
sc.tl.leiden(adata_QC, resolution=0.5)

# UMAPプロットの作成
sc.pl.umap(adata_QC, color='leiden', legend_loc='on data', title='UMAP with Clusters', frameon=False)

#decouplerはクラスタリング後のデータを用いるので、ここで保存する。
adata_clustered = adata_QC.copy()


adata_decoupler = adata_clustered.copy()

markers = dc.get_resource('PanglaoDB')

# 細胞マーカーをロードする
markers = dc.get_resource('PanglaoDB')
# markers = markers[markers['human'] & markers['canonical_marker'] & (markers['human_sensitivity'] > 0.5)]
markers = markers[~markers.duplicated(['cell_type', 'genesymbol'])]

# ORAを実行する
dc.run_ora(mat=adata_decoupler, net=markers, source='cell_type', target='genesymbol', min_n=3, verbose=True,use_raw=False)

# ORAのスコアを保存
acts = dc.get_acts(adata_decoupler, obsm_key='ora_estimate')
acts_v = acts.X.ravel()
max_e = np.nanmax(acts_v[np.isfinite(acts_v)])
acts.X[~np.isfinite(acts.X)] = max_e

#ORAスコアを元にクラスターごとに細胞タイプのランキングを作る
df_decoupler = dc.rank_sources_groups(acts, groupby='leiden', reference='rest', method='t-test_overestim_var')

# ランキング1位の細胞タイプを取り出す
annotation_dict = df_decoupler.groupby('group').head(1).set_index('group')['names'].to_dict()

# アノテーション
adata_decoupler.obs['cell_type'] = [annotation_dict[clust] for clust in adata_decoupler.obs['leiden']]

#　UMAPで可視化
sc.pl.umap(adata_decoupler, color='cell_type')


adata_CellTypist= adata_QC.copy()
adata_CellTypist.X = adata_QC.layers['before_DimRed'].copy()

# modelをアップデート
models.download_models(force_update = True)

# モデルのインスタンス生成。免疫細胞を低解像度にアノテーションするモデルを設定
model = models.Model.load(model = 'Immune_All_Low.pkl')

# アノテーションの実行
predictions = celltypist.annotate(adata_CellTypist, model = 'Immune_All_Low.pkl', majority_voting = True)

# アノテーション結果の追加

adata_CellTypist = predictions.to_adata()

# アノテーション結果の保存
adata.write('annotated_data.h5ad')

# スケーリングされた高変動遺伝子のデータを元のデータに戻す
adata_CellTypist[:, highly_variable_mask].X = adata_highly_variable.X

# PCAの実行
sc.tl.pca(adata_CellTypist,svd_solver='arpack')

# UMAPの実行
sc.pp.neighbors(adata_CellTypist, n_pcs=20)  # UMAPの前に隣接グラフを作成する必要があります
sc.tl.umap(adata_CellTypist)

# 隣接グラフの構築
sc.pp.neighbors(adata_CellTypist, n_pcs=10)

# クラスタリングの実行
sc.tl.leiden(adata_CellTypist, resolution=0.5)
# UMAPプロットの作成
# sc.pl.umap(adata_CellTypist, color='predicted_labels', legend_loc='on data', title='UMAP with Annotations', frameon=False)

sc.pl.umap(adata_CellTypist, color='predicted_labels', title='UMAP with Annotations', frameon=False)


# バーコード、特徴量、マトリックスの読み込み
barcodes = pd.read_csv(barcodes_path, header=None)
features = pd.read_csv(features_path, header=None, sep='\\t')
matrix = scipy.io.mmread(matrix_path)

# AnnDataオブジェクトの作成
adata = sc.AnnData(X=matrix.T.tocsr())
adata.obs_names = barcodes[0].values #細胞名　obs.indexとおなじ
adata.var_names = features[1].values #遺伝子名 var.indexとおなじ


# 遺伝子の重複を取り除く
new_data = adata.to_df()
new_data = new_data.groupby(new_data.columns, axis=1).sum() 
adata =  sc.AnnData(X=new_data.values)
adata.obs_names = new_data.index
adata.var_names = new_data.columns

def QC(adata):
		mt_genes = adata.var_names.str.startswith('MT-')
    # ミトコンドリア遺伝子の発現割合を計算してobsに追加
    percent_mt = adata[:,mt_genes].X.sum(axis=1) / adata.X.sum(axis=1) * 100
    # 各細胞におけるユニークな遺伝子の数（非ゼロの遺伝子数）
    nFeature_RNA = (adata.X > 0).sum(axis=1)
    # 各細胞における全てのRNA分子の総数
    nCount_RNA = adata.X.sum(axis=1)

    df_violin_QC =\\
    pd.DataFrame({"nFeature_RNA":nFeature_RNA, "nCount_RNA":nCount_RNA, "percent_mt":percent_mt},\\
                index = adata.obs_names)
    return df_violin_QC
    
df_violin_QC = QC(adata)

adata_QC = adata[(200< df_violin_QC['nFeature_RNA'])\\
                 & (df_violin_QC['nFeature_RNA'] < 5000)\\
                 & (df_violin_QC['percent_mt'] < 15)]

# 正規化
sc.pp.normalize_total(adata_QC, target_sum=1e4)
sc.pp.log1p(adata_QC) # celltypistはlog1pで変換されていることが前提になっている。

# 高可変性遺伝子の検出
sc.pp.highly_variable_genes(adata_QC, flavor='seurat_v3', n_top_genes=2000)

# 高変動遺伝子のマスクを作成
highly_variable_mask = adata_QC.var['highly_variable']
# 高変動遺伝子のみをスケーリング
adata_highly_variable = adata_QC[:, highly_variable_mask].copy()
sc.pp.scale(adata_highly_variable)

# CellTypistはスケーリング・次元削減前のデータを用いるので、ここで保存する。
# スケーリング前のデータを保存
adata_QC.layers['before_DimRed'] = adata_QC.X.copy()

# スケーリングされた高変動遺伝子のデータを元のデータに戻す
adata_QC[:, highly_variable_mask].X = adata_highly_variable.X

# PCAの実行
sc.tl.pca(adata_QC,svd_solver='arpack')

# UMAPの実行
sc.pp.neighbors(adata_QC, n_pcs=20)  # UMAPの前に隣接グラフを作成する必要があります
sc.tl.umap(adata_QC)

# 隣接グラフの構築
sc.pp.neighbors(adata_QC, n_pcs=10)

# クラスタリングの実行
sc.tl.leiden(adata_QC, resolution=0.5)

# UMAPプロットの作成
sc.pl.umap(adata_QC, color='leiden', legend_loc='on data', title='UMAP with Clusters', frameon=False)

#decouplerはクラスタリング後のデータを用いるので、ここで保存する。
adata_clustered = adata_QC.copy()

adata_decoupler = adata_clustered.copy()

markers = dc.get_resource('PanglaoDB')

# 細胞マーカーをロードする
markers = dc.get_resource('PanglaoDB')
# markers = markers[markers['human'] & markers['canonical_marker'] & (markers['human_sensitivity'] > 0.5)]
markers = markers[~markers.duplicated(['cell_type', 'genesymbol'])]

# ORAを実行する
dc.run_ora(mat=adata_decoupler, net=markers, source='cell_type', target='genesymbol', min_n=3, verbose=True,use_raw=False)

# ORAのスコアを保存
acts = dc.get_acts(adata_decoupler, obsm_key='ora_estimate')
acts_v = acts.X.ravel()
max_e = np.nanmax(acts_v[np.isfinite(acts_v)])
acts.X[~np.isfinite(acts.X)] = max_e

#ORAスコアを元にクラスターごとに細胞タイプのランキングを作る
df_decoupler = dc.rank_sources_groups(acts, groupby='leiden', reference='rest', method='t-test_overestim_var')

# ランキング1位の細胞タイプを取り出す
annotation_dict = df_decoupler.groupby('group').head(1).set_index('group')['names'].to_dict()

# アノテーション
adata_decoupler.obs['cell_type'] = [annotation_dict[clust] for clust in adata_decoupler.obs['leiden']]

#　UMAPで可視化
sc.pl.umap(adata_decoupler, color='cell_type')

adata_CellTypist= adata_QC.copy()
adata_CellTypist.X = adata_QC.layers['before_DimRed'].copy()

# modelをアップデート
models.download_models(force_update = True)

# モデルのインスタンス生成。免疫細胞を低解像度にアノテーションするモデルを設定
model = models.Model.load(model = 'Immune_All_Low.pkl')

# アノテーションの実行
predictions = celltypist.annotate(adata_CellTypist, model = 'Immune_All_Low.pkl', majority_voting = True)

# アノテーション結果の追加

adata_CellTypist = predictions.to_adata()

# アノテーション結果の保存
adata.write('annotated_data.h5ad')

# スケーリングされた高変動遺伝子のデータを元のデータに戻す
adata_CellTypist[:, highly_variable_mask].X = adata_highly_variable.X

# PCAの実行
sc.tl.pca(adata_CellTypist,svd_solver='arpack')

# UMAPの実行
sc.pp.neighbors(adata_CellTypist, n_pcs=20)  # UMAPの前に隣接グラフを作成する必要があります
sc.tl.umap(adata_CellTypist)

# 隣接グラフの構築
sc.pp.neighbors(adata_CellTypist, n_pcs=10)

# クラスタリングの実行
sc.tl.leiden(adata_CellTypist, resolution=0.5)
# UMAPプロットの作成
# sc.pl.umap(adata_CellTypist, color='predicted_labels', legend_loc='on data', title='UMAP with Annotations', frameon=False)

sc.pl.umap(adata_CellTypist, color='predicted_labels', title='UMAP with Annotations', frameon=False)

# 100以上の出現回数を持つ細胞タイプを取得
cell_type_counts = adata_CellTypist.obs['predicted_labels'].value_counts()
frequent_cell_types = cell_type_counts[cell_type_counts >= 70].index.tolist()
rare_cell_types = cell_type_counts[(cell_type_counts < 70) & (cell_type_counts >= 5)].index.tolist()

# 頻出する細胞タイプのインデックスを取得
selected_indices = adata_CellTypist.obs['predicted_labels'].isin(frequent_cell_types)

# 新しいAnnDataオブジェクトを作成（選択された細胞タイプ）
adata_CellTypist_selected = sc.AnnData(
    X=adata_CellTypist.X[selected_indices],
    obs=adata_CellTypist.obs.loc[selected_indices],
    var=adata_CellTypist.var,
    uns=adata_CellTypist.uns.copy(),
    obsm={'X_umap': adata_CellTypist.obsm['X_umap'][selected_indices]},
    varm=adata_CellTypist.varm,
    layers={key: layer[selected_indices] for key, layer in adata_CellTypist.layers.items()},
    obsp={key: connectivities[selected_indices][:, selected_indices] for key, connectivities in adata_CellTypist.obsp.items()}
)

# predicted_labelsのカテゴリを選択された細胞タイプに限定
adata_CellTypist_selected.obs['predicted_labels'] = pd.Categorical(
    adata_CellTypist_selected.obs['predicted_labels'],
    categories=frequent_cell_types
)

# 稀な細胞タイプのインデックスを取得

rare_indices = adata_CellTypist.obs['predicted_labels'].isin(rare_cell_types)

# 新しいAnnDataオブジェクトを作成（稀な細胞タイプ）
adata_CellTypist_rare = sc.AnnData(
    X=adata_CellTypist.X[rare_indices],
    obs=adata_CellTypist.obs.loc[rare_indices],
    var=adata_CellTypist.var,
    uns=adata_CellTypist.uns.copy(),
    obsm={'X_umap': adata_CellTypist.obsm['X_umap'][rare_indices]},
    varm=adata_CellTypist.varm,
    layers={key: layer[rare_indices] for key, layer in adata_CellTypist.layers.items()},
    obsp={key: connectivities[rare_indices][:, rare_indices] for key, connectivities in adata_CellTypist.obsp.items()}
)

# predicted_labelsのカテゴリを稀な細胞タイプに限定
rare_cell_types = list(set(adata_CellTypist_rare.obs['predicted_labels'].unique()) - set(frequent_cell_types))
adata_CellTypist_rare.obs['predicted_labels'] = pd.Categorical(
    adata_CellTypist_rare.obs['predicted_labels'],
    categories=rare_cell_types
)
# UMAPプロット
sc.pl.umap(adata_CellTypist_selected, 
           color='predicted_labels',
           title = "common cells", 
           frameon=False)

# UMAPプロット
sc.pl.umap(adata_CellTypist_rare, 
           color='predicted_labels', 
           title = "Rare Cells",
           frameon=False)
           
           
#少ない細胞数のグループを、多い細胞数のグループの上に重ねて表示する。
  # プロットの設定
fig, ax = plt.subplots(figsize=(12, 10))

# adata_CellTypist_selectedをグレーで表示（キャプションなし）
sc.pl.umap(adata_CellTypist_selected, 
           ax=ax, 
           color='predicted_labels', 
           palette=['lightgrey'] * len(adata_CellTypist_selected.obs['predicted_labels'].cat.categories),
           legend_loc=None,  # キャプションを非表示
           show=False, 
           title='')

# adata_CellTypist_rareをcolor='predicted_labels'で表示（キャプション付き）
sc.pl.umap(adata_CellTypist_rare, 
           ax=ax, 
           color='predicted_labels', 
        #    legend_loc='on data',  # キャプションをプロット内に表示
           legend_fontsize=8,
           legend_fontoutline=2,
           size=100,
           show=False, 
           title='')
plt.tight_layout()
plt.show()