#!/usr/bin/env python
"""RNA velocity 样例(公开数据:小鼠胰腺内分泌发育,scVelo 内置)。

选这个模块是因为它在论文里被过度解读得最厉害:一张箭头图看起来像是「证明了分化方向」,
但 velocity 的方向来自一个**模型假设**(未剪接 RNA 先于剪接 RNA),假设不成立时箭头照样会画出来。

所以这份样品的重点不是「我们也会画箭头」,而是把该报的三个数一起报出来:
  ① 未剪接/剪接的比例 —— 太低就根本不该做 velocity(信息量不够);
  ② velocity confidence —— 每个细胞的方向有多可信,低置信区域的箭头不能拿来讲故事;
  ③ 和已知生物学对不对得上 —— 胰腺内分泌发育的终点是 Alpha/Beta/Delta/Epsilon 四种内分泌细胞,
     若箭头指回祖细胞,就是模型没拟合好,而不是发现了「去分化」。
"""
import json, os, warnings
warnings.filterwarnings("ignore")
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
import scvelo as scv

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out_velocity")
os.makedirs(OUT, exist_ok=True)
scv.settings.figdir = OUT
scv.settings.set_figure_params("scvelo", dpi=140, dpi_save=140, facecolor="white", frameon=False)
M = {"dataset": "scVelo pancreas endocrinogenesis (Bastidas-Ponce et al. 2019, public)"}

adata = scv.datasets.pancreas()
M["cells"], M["genes"] = int(adata.n_obs), int(adata.n_vars)
M["cell_types"] = adata.obs["clusters"].value_counts().to_dict()

# ---------- ① 未剪接/剪接比例:够不够做 velocity ----------
prop = {k: float(v) for k, v in zip(["spliced", "unspliced"],
        [adata.layers["spliced"].sum(), adata.layers["unspliced"].sum()])}
tot = sum(prop.values())
M["layer_proportions"] = {k: round(100 * v / tot, 1) for k, v in prop.items()}
M["unspliced_pct"] = M["layer_proportions"]["unspliced"]
M["enough_unspliced"] = bool(M["unspliced_pct"] >= 10)

# scvelo 0.3.4 里 filter_and_normalize 会把多余的关键字透传给 normalize_per_cell 而报错,
# 且已不再导出 filter_genes_dispersion。改用 scvelo 过滤 + scanpy 选高变基因,不依赖那些版本细节。
scv.pp.filter_genes(adata, min_shared_counts=20)
scv.pp.normalize_per_cell(adata)
sc.pp.log1p(adata)
sc.pp.highly_variable_genes(adata, n_top_genes=2000, subset=True)
scv.pp.moments(adata, n_pcs=30, n_neighbors=30)
# stochastic 模式在 scvelo 0.3.4 + numpy 2.x 下会在最小二乘里崩(数组形状假设过时);
# deterministic 走的是另一条代码路径。**记下来是因为这本身是有用的信息**:
# velocity 这套工具链对版本很敏感,交付前必须实跑而不是照抄教程。
scv.tl.velocity(adata, mode="deterministic")
scv.tl.velocity_graph(adata, n_jobs=4)

# ---------- ② 方向的置信度 ----------
scv.tl.velocity_confidence(adata)
conf = adata.obs["velocity_confidence"].values
M["velocity_confidence"] = {"median": round(float(np.median(conf)), 3),
                            "low_confidence_cells_pct": round(100 * float((conf < 0.5).mean()), 1)}

scv.tl.velocity_pseudotime(adata)
M["pseudotime_range"] = [round(float(adata.obs["velocity_pseudotime"].min()), 3),
                         round(float(adata.obs["velocity_pseudotime"].max()), 3)]

# ---------- ③ 和已知生物学对不对得上 ----------
# 胰腺内分泌发育:Ductal → Ngn3 low → Ngn3 high → Pre-endocrine → Alpha/Beta/Delta/Epsilon
ORDER = ["Ductal", "Ngn3 low EP", "Ngn3 high EP", "Pre-endocrine", "Beta", "Alpha", "Delta", "Epsilon"]
pt = adata.obs.groupby("clusters")["velocity_pseudotime"].median().to_dict()
known = [c for c in ORDER if c in pt]
M["median_pseudotime_by_type"] = {k: round(float(pt[k]), 3) for k in known}
early = [c for c in ["Ductal", "Ngn3 low EP"] if c in pt]
late = [c for c in ["Beta", "Alpha", "Delta", "Epsilon"] if c in pt]
if early and late:
    M["early_before_late"] = bool(max(pt[c] for c in early) < min(pt[c] for c in late))
    M["sanity_check_passed"] = M["early_before_late"] and M["enough_unspliced"]

scv.pl.velocity_embedding_stream(adata, basis="umap", color="clusters",
                                 title="RNA velocity — pancreas endocrinogenesis",
                                 save="velocity_stream.png", show=False)
scv.pl.scatter(adata, color="velocity_confidence", cmap="coolwarm",
               title="Velocity confidence per cell", save="velocity_confidence.png", show=False)
scv.pl.proportions(adata, save="proportions.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: M[k] for k in ("cells", "layer_proportions", "enough_unspliced",
                                     "velocity_confidence", "early_before_late",
                                     "sanity_check_passed")}, indent=2, ensure_ascii=False))
print("各细胞类型的中位 velocity 拟时序:")
for k, v in M["median_pseudotime_by_type"].items():
    print(f"   {k:<16} {v}")