#!/usr/bin/env python
"""公开数据的真实空间转录组分析样例(10x Visium,人淋巴结)。

空间是单价最高的一档服务($550–1,400/张切片),所以样品必须是真跑的:
质控 → 过滤 → 标准化 → 聚类 → 空间域可视化 → 空间自相关(Moran's I)找空间变异基因。
所有数字与图都由本脚本算出,可复算。
"""
import json, os, warnings
warnings.filterwarnings("ignore")
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import scanpy as sc
import numpy as np

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out_spatial")
os.makedirs(OUT, exist_ok=True)
sc.settings.figdir = OUT
sc.settings.set_figure_params(dpi=140, frameon=False, facecolor="white")
M = {}

adata = sc.datasets.visium_sge(sample_id="V1_Human_Lymph_Node")
adata.var_names_make_unique()
M["spots_raw"], M["genes_raw"] = int(adata.n_obs), int(adata.n_vars)

adata.var["mt"] = adata.var_names.str.startswith("MT-")
sc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], inplace=True, percent_top=None, log1p=False)
M["median_genes_per_spot"] = float(np.median(adata.obs["n_genes_by_counts"]))
M["median_umi_per_spot"] = float(np.median(adata.obs["total_counts"]))
M["median_pct_mito"] = round(float(np.median(adata.obs["pct_counts_mt"])), 2)

fig, axes = plt.subplots(1, 3, figsize=(12, 3.4))
for ax, key, title in zip(axes, ["total_counts", "n_genes_by_counts", "pct_counts_mt"],
                          ["UMIs per spot", "Genes per spot", "Mitochondrial %"]):
    ax.hist(adata.obs[key].values, bins=60, color="#0a6ed1")
    ax.set_title(title, fontsize=11); ax.set_yticks([])
plt.tight_layout(); plt.savefig(f"{OUT}/spatial_qc.png", dpi=140, bbox_inches="tight"); plt.close()

sc.pp.filter_cells(adata, min_counts=500)
sc.pp.filter_genes(adata, min_cells=10)
M["spots_after_qc"], M["genes_after_qc"] = int(adata.n_obs), int(adata.n_vars)
M["spots_removed_pct"] = round(100 * (1 - M["spots_after_qc"] / M["spots_raw"]), 1)

sc.pp.normalize_total(adata, inplace=True)
sc.pp.log1p(adata)
sc.pp.highly_variable_genes(adata, flavor="seurat", n_top_genes=2000)
M["hvg"] = int(adata.var.highly_variable.sum())
sc.pp.pca(adata, n_comps=50)
sc.pp.neighbors(adata)
sc.tl.umap(adata)
sc.tl.leiden(adata, resolution=1.0, key_added="domain", flavor="igraph", n_iterations=2, directed=False)
M["domains"] = int(adata.obs["domain"].nunique())
M["domain_sizes"] = {k: int(v) for k, v in adata.obs["domain"].value_counts().items()}

sc.pl.spatial(adata, color="domain", size=1.4, title="Spatial domains (Leiden)",
              save="_domains.png", show=False)
sc.pl.spatial(adata, color="total_counts", size=1.4, title="UMIs per spot",
              save="_counts.png", show=False)

# 空间自相关:Moran's I 找真正随位置变化的基因
sc.pp.neighbors(adata, key_added="spatial_neigh", use_rep="X_pca")
hvg = adata.var_names[adata.var.highly_variable][:1500]
mi = sc.metrics.morans_i(adata[:, hvg])
order = np.argsort(mi)[::-1]
top_sv = [(str(hvg[i]), round(float(mi[i]), 3)) for i in order[:15]]
M["top_spatially_variable"] = top_sv

sc.tl.rank_genes_groups(adata, "domain", method="wilcoxon")
M["domain_markers"] = {g: [str(x) for x in adata.uns["rank_genes_groups"]["names"][g][:6]]
                       for g in adata.obs["domain"].cat.categories}

sc.pl.spatial(adata, color=[t[0] for t in top_sv[:4]], size=1.4, ncols=4,
              save="_svg.png", show=False)

with open(f"{OUT}/metrics.json", "w") as f:
    json.dump(M, f, indent=2, ensure_ascii=False)
print(json.dumps({k: v for k, v in M.items() if k not in ("domain_markers", "domain_sizes")}, indent=2))
print("输出:", OUT)