# -*- coding: utf-8 -*-
"""
04b_cell2location_squidpy.py —— 去卷积（cell2location）与邻域富集（Squidpy）
对应教程第 09/10 章（Python 部分）。

用法：
  conda activate spatial
  python code/04b_cell2location_squidpy.py

说明：
  1) cell2location 用"参考单细胞图谱"估计每个 spot 的细胞类型丰度（去卷积），
     比 R 端 marker 打分更精细；
  2) squidpy 邻域富集验证"哪些细胞类型倾向于在空间上共定位"，
     为生态位定义提供证据。
  3) 完整跑 cell2location 需要下载/准备参考单细胞数据，本脚本给出核心流程；
     若只想快速体验，可直接运行第 4 节 squidpy 部分（用已有注释结果）。
"""

import scanpy as sc
import squidpy as sq
import numpy as np

# ---------------------------------------------------------------
# 1. 读入 Visium 数据（h5ad 或从 10x 目录读取）
# ---------------------------------------------------------------
# 如果 R 端已导出 h5ad：
# adata = sc.read_h5ad("data/processed/visium_annotated.h5ad")
# 或直接从 10x 目录读：
adata = sc.read_visium("data/raw/visium")

# ---------------------------------------------------------------
# 2. cell2location：参考单细胞模型 → spot 细胞丰度
#    参考数据（scRNA）通常来自同一组织的公开单细胞图谱。
# ---------------------------------------------------------------
# import cell2location
# from cell2location.models import RegressionModel, Cell2location
#
# # (a) 训练参考模型（在单细胞数据上）
# sc_ref = sc.read_h5ad("data/ref/breast_cancer_scRNA.h5ad")
# sc_ref.X = sc_ref.layers["counts"].copy()
# cell2location.models.RegressionModel.setup_anndata(sc_ref, batch_key="batch")
# ref_model = RegressionModel(sc_ref)
# ref_model.train(max_epochs=250)
# ref_model.export_posterior(sc_ref)
#
# # (b) 在 Visium 上估计细胞丰度
# cell2location.models.Cell2location.setup_anndata(adata, batch_key="sample")
# mod = Cell2location(adata, sc_ref, N_cells_per_location=30,
#                     detection_alpha=20)
# mod.train(max_epochs=30000)
# mod.export_posterior(adata)
#
# # (c) 结果：adata.obsm["q05_cell_abundance_w_sf"] = spot × 细胞类型
# adata.obsm["abundance"] = adata.obsm["q05_cell_abundance_w_sf"]
# adata.obs["dominant_celltype"] = adata.obsm["abundance"].idxmax(axis=1).values
# 注意：cell2location 训练较慢（数小时），建议 GPU 或先跑小数据。

# ---------------------------------------------------------------
# 3. （可选）把 R 端注释结果导入：celltype 列已在 obs 中
# ---------------------------------------------------------------
if "celltype" not in adata.obs.columns:
    adata.obs["celltype"] = adata.obs.get("dominant_celltype", "unknown")

# ---------------------------------------------------------------
# 4. Squidpy 邻域富集（生态位验证的核心证据）
# ---------------------------------------------------------------
# (a) 构建空间邻接图（k 近邻）
sq.gr.spatial_neighbors(adata, coord_type="generic", n_neighs=6)

# (b) 邻域富集：观察共现 vs 随机期望（z 分数，正=富集）
sq.gr.nhood_enrichment(adata, cluster_key="celltype")
sq.pl.nhood_enrichment(adata, cluster_key="celltype")

# (c) 共现曲线（随距离衰减）
sq.gr.co_occurrence(adata, cluster_key="celltype")
sq.pl.co_occurrence(adata, cluster_key="celltype", clusters="CD8 T")

# ---------------------------------------------------------------
# 5. 空间可视化
# ---------------------------------------------------------------
sc.pl.spatial(adata, color="celltype", spot_size=1.2, title="细胞类型（空间）")
sc.pl.spatial(adata, color="niche", spot_size=1.2, title="空间生态位")

print("去卷积与邻域富集分析完成！")
