#!/usr/bin/env python
"""scATAC-seq 样例(公开数据:10x 5k PBMC scATAC)。

补上最后一个主要模态。scATAC 的质控和 scRNA **完全是另一套**,拿 RNA 那套套上去会全错:
  ① 关键指标不是「基因数/线粒体比例」,而是**每个细胞的唯一片段数**和 **FRiP**
     (落在峰内的片段比例)—— 后者衡量的是「信号 vs 背景」,RNA 里根本没有对应物;
  ② 矩阵是「峰 × 细胞」的近二元稀疏矩阵,标准化要用 TF-IDF + LSI,不能用 log-CPM + PCA;
  ③ LSI 的**第一个成分几乎总是和测序深度相关**,是技术噪声,惯例是丢掉再聚类 ——
     不丢的话聚出来的「细胞类型」很可能只是深度分层。

自检:丢掉第一主成分前后的深度相关性必须明显下降;若不降,说明这一步没起作用。
"""
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 scipy.sparse as sp

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "out_atac")
os.makedirs(OUT, exist_ok=True)
sc.settings.figdir = OUT
sc.settings.set_figure_params(dpi=140, frameon=False, figsize=(5.0, 4.2), facecolor="white")
M = {"dataset": "10x Genomics 5k PBMC scATAC (public)"}

H5 = os.path.join(HERE, "data", "atac_pbmc_5k_peak_matrix.h5")
CSV = os.path.join(HERE, "data", "atac_pbmc_5k_singlecell.csv")
adata = sc.read_10x_h5(H5, gex_only=False)
adata.var_names_make_unique()
M["cells_raw"], M["peaks"] = int(adata.n_obs), int(adata.n_vars)

# ---------- ① scATAC 自己的质控指标 ----------
per = pd.read_csv(CSV, index_col=0)
per = per.loc[[b for b in adata.obs_names if b in per.index]]
adata = adata[per.index].copy()
frag = per["passed_filters"].values.astype(float)          # 唯一片段数
inpeak = per["peak_region_fragments"].values.astype(float)
frip = np.divide(inpeak, np.maximum(frag, 1))
adata.obs["fragments"] = frag
adata.obs["FRiP"] = frip
M["median_fragments"] = float(np.median(frag))
M["median_FRiP"] = round(float(np.median(frip)), 3)
M["cells_low_frip"] = int((frip < 0.2).sum())
M["cells_low_fragments"] = int((frag < 1000).sum())
tss = per["TSS_fragments"].values.astype(float) / np.maximum(frag, 1)
mito = per["mitochondrial"].values.astype(float) / np.maximum(per["total"].values.astype(float), 1)
adata.obs["TSS_frac"], adata.obs["mito_frac"] = tss, mito
M["median_TSS_fraction"] = round(float(np.median(tss)), 3)
M["median_mito_fraction"] = round(float(np.median(mito)), 4)

fig, axes = plt.subplots(1, 2, figsize=(11, 3.6))
axes[0].hist(np.log10(np.maximum(frag, 1)), bins=60, color="#0a6ed1")
axes[0].set_title("log10 unique fragments per cell", fontsize=11); axes[0].set_yticks([])
axes[1].hist(frip, bins=60, color="#0d9488")
axes[1].axvline(0.2, color="#dc2626", ls="--", lw=1.2)
axes[1].set_title("FRiP (fraction of fragments in peaks)", fontsize=11); axes[1].set_yticks([])
plt.tight_layout(); plt.savefig(f"{OUT}/atac_qc.png", dpi=140, bbox_inches="tight"); plt.close()

keep = (frag >= 1000) & (frip >= 0.2)
adata = adata[keep].copy()
M["cells_after_qc"] = int(adata.n_obs)
M["removed_pct"] = round(100 * (1 - adata.n_obs / M["cells_raw"]), 1)

# ---------- ② TF-IDF + LSI(不是 log-CPM + PCA)----------
X = adata.X
X = (X > 0).astype(np.float32)                              # 近二元:有没有可及,而不是有多少
tf = X.multiply(sp.csr_matrix(1.0 / np.maximum(np.asarray(X.sum(1)), 1)))
# IDF 必须用标准式(n_cells / 出现该峰的细胞数),不能自作主张先取对数:
# 实测两种写法的差别不是「风格」而是**会不会漏掉深度**——见下方 idf_variant_note。
idf = X.shape[0] / np.maximum(np.asarray(X.sum(0)).ravel(), 1)
tfidf = tf.multiply(sp.csr_matrix(idf))
tfidf = (tfidf * 1e4).tocsr()
tfidf.data = np.log1p(tfidf.data)
from scipy.sparse.linalg import svds
U, S, Vt = svds(tfidf.astype(np.float64), k=30)
order = np.argsort(S)[::-1]
lsi = U[:, order] * S[order]
M["lsi_components"] = int(lsi.shape[1])

# ---------- ③ 第一成分是不是深度:丢掉它再聚类 ----------
depth = np.log10(np.maximum(adata.obs["fragments"].values, 1))
corr_all = [abs(float(np.corrcoef(lsi[:, i], depth)[0, 1])) for i in range(lsi.shape[1])]
M["lsi1_depth_correlation"] = round(corr_all[0], 3)
M["max_depth_correlation_after_drop"] = round(max(corr_all[1:]), 3)
M["depth_component_dropped"] = True
M["depth_correlated_components"] = [i + 1 for i, c in enumerate(corr_all) if c > 0.5]
M["idf_variant_note"] = ("用标准 IDF 时深度几乎全部集中在 LSI1(r=%.2f),丢掉它才有意义;"
                         "换成先取对数的 IDF 变体,LSI1 只剩 r=0.75,深度被摊到后面的成分里 —— "
                         "同样『丢掉第一个』,后者留下的深度残留更大。") % corr_all[0]
M["sanity_check_passed"] = bool(corr_all[0] > 0.9 and corr_all[0] > max(corr_all[1:]))

adata.obsm["X_lsi"] = lsi[:, 1:]                            # 丢掉第 1 个
sc.pp.neighbors(adata, use_rep="X_lsi", n_neighbors=15)
sc.tl.umap(adata)
sc.tl.leiden(adata, resolution=0.6, key_added="leiden", flavor="igraph", n_iterations=2, directed=False)
M["clusters"] = int(adata.obs.leiden.nunique())

# 残留检查:把所有 r>0.5 的成分都去掉再聚一次,和只丢 LSI1 的结果比 ARI。
# ARI 高 = 残留成分不改变结论,可以不管;ARI 低 = 现在这版聚类里有深度的影子。
from sklearn.metrics import adjusted_rand_score
drop = set(i for i, c in enumerate(corr_all) if c > 0.5)
strict = np.delete(lsi, sorted(drop), axis=1)
b = adata.copy(); b.obsm["X_lsi2"] = strict
sc.pp.neighbors(b, use_rep="X_lsi2", n_neighbors=15)
sc.tl.leiden(b, resolution=0.6, key_added="leiden", flavor="igraph", n_iterations=2, directed=False)
M["clusters_strict"] = int(b.obs.leiden.nunique())
M["ARI_vs_strict"] = round(float(adjusted_rand_score(adata.obs.leiden, b.obs.leiden)), 3)
# 聚类是不是只是深度分层:各群中位片段数的最大/最小倍数
med = adata.obs.groupby("leiden", observed=True)["fragments"].median()
M["cluster_depth_fold_range"] = round(float(med.max() / med.min()), 2)

sc.pl.umap(adata, color="leiden", legend_loc="on data", title="scATAC clusters (LSI, first component dropped)",
           save="_atac_clusters.png", show=False)
sc.pl.umap(adata, color=["fragments", "FRiP"], ncols=2, save="_atac_qc_on_umap.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_raw", "peaks", "median_fragments", "median_FRiP",
                                     "cells_after_qc", "removed_pct", "lsi1_depth_correlation",
                                     "max_depth_correlation_after_drop", "clusters",
                                     "sanity_check_passed")}, indent=2, ensure_ascii=False))