scripts/analysis/100_primary_eden_discovery.py

"""primary EDEN discovery on Dingwall: PANDA-v3 dermal-fibro subset, Leiden res=1.5, wilcoxon markers + Fisher cKO depletion + module scoring."""
from pathlib import Path
import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, torch.nn.functional as F
from scipy.stats import fisher_exact, mannwhitneyu
warnings.filterwarnings("ignore"); sc.settings.verbosity = 0
import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
sys.path.insert(0, str(PANDA_ROOT))
from panda import PANDAEncoder

ROOT = Path(str(PANDA_ROOT))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

CKO_GSMS = {"GSM6833482", "GSM6833483"}  # CORRECTED: 480/481 are rttaControl (WT), not cKO
WT_GSMS  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}  # CORRECTED: 4 Cre-neg controls per GEO metadata

# secondary EDEN definition per Dingwall 2024
SECONDARY_EDEN_PANEL = ["S100a4", "Tnc", "Pdgfra"]

# En1-responsive eccrine program (from restored 57_pathway_analysis.py, En1 removed)
SWEAT_GLAND_PANEL_ENSMINUSEN1 = ["Foxi3", "Foxa1", "Krt8", "Krt18", "Krt19", "Muc5b", "Aqp5"]
# Eda pathway
EDA_PATHWAY_PANEL = ["Eda", "Edar", "Edaradd", "Nfkb1", "Nfkb2", "Rela"]


def load_dingwall_with_v3_predictions():
    print("[eden] loading Dingwall raw + v3 predictions", flush=True)
    raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
    pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv")
    pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
    raw.obs["pred_label"] = np.array([pred_map.get(c, "unknown") for c in raw.obs_names.astype(str)])
    raw.obs["genotype"] = np.where(raw.obs["sample"].astype(str).isin(list(CKO_GSMS)), "En1-cKO",
                          np.where(raw.obs["sample"].astype(str).isin(list(WT_GSMS)), "WT", "other"))
    labeled = raw.obs["genotype"].isin(["WT", "En1-cKO"]).values
    raw = raw[labeled].copy()

    dermal_mask = np.isin(raw.obs["pred_label"], ["fibroblast-papillary", "fibroblast-reticular"])
    dermal = raw[dermal_mask].copy()
    print(f"[eden] {dermal.n_obs} dermal-fibroblast cells for sub-clustering", flush=True)
    return dermal


def subcluster_dermal(dermal, resolution=1.5):
    print(f"[eden] preprocessing + PCA (Leiden resolution={resolution})", flush=True)
    sc.pp.normalize_total(dermal, target_sum=1e4); sc.pp.log1p(dermal)
    sc.pp.highly_variable_genes(dermal, n_top_genes=3000, flavor="seurat_v3",
                                inplace=True, batch_key=None)
    dermal_hvg = dermal[:, dermal.var["highly_variable"]].copy() if "highly_variable" in dermal.var else dermal
    sc.pp.scale(dermal_hvg, max_value=10)
    sc.tl.pca(dermal_hvg, n_comps=30, random_state=0)
    sc.pp.neighbors(dermal_hvg, n_neighbors=20, use_rep="X_pca")
    sc.tl.leiden(dermal_hvg, resolution=resolution, random_state=0)
    dermal.obs["leiden"] = dermal_hvg.obs["leiden"].astype(str)
    print(f"[eden] {dermal.obs['leiden'].nunique()} sub-clusters found", flush=True)
    return dermal


def score_modules(dermal):
    for name, genes in [("secondary_eden", SECONDARY_EDEN_PANEL),
                         ("sweat_gland", SWEAT_GLAND_PANEL_ENSMINUSEN1),
                         ("eda_pathway", EDA_PATHWAY_PANEL)]:
        present = [g for g in genes if g in dermal.var_names]
        if not present:
            dermal.obs[f"score_{name}"] = 0.0
            continue
        sc.tl.score_genes(dermal, gene_list=present, score_name=f"score_{name}",
                          random_state=0, use_raw=False)
    return dermal


def per_subcluster_analysis(dermal):
    n_wt_tot = int((dermal.obs["genotype"] == "WT").sum())
    n_cko_tot = int((dermal.obs["genotype"] == "En1-cKO").sum())
    baseline_cko_frac = n_cko_tot / (n_wt_tot + n_cko_tot)
    print(f"[eden] baseline: WT={n_wt_tot} cKO={n_cko_tot} (baseline cKO frac = {baseline_cko_frac:.3f})", flush=True)

    sc.tl.rank_genes_groups(dermal, "leiden", method="wilcoxon", n_genes=30, use_raw=False)

    rows = []
    for cls in sorted(dermal.obs["leiden"].unique(), key=int):
        sub = dermal[dermal.obs["leiden"] == cls]
        n_wt = int((sub.obs["genotype"] == "WT").sum())
        n_cko = int((sub.obs["genotype"] == "En1-cKO").sum())
        if n_wt + n_cko < 20:
            continue
        cko_frac = n_cko / (n_wt + n_cko)

        # fisher 2x2: (n_wt_in, n_wt_out) vs (n_cko_in, n_cko_out) — cluster depletion in cKO
        n_wt_elsewhere = n_wt_tot - n_wt
        n_cko_elsewhere = n_cko_tot - n_cko
        odds, p_fisher = fisher_exact([[n_wt, n_wt_elsewhere], [n_cko, n_cko_elsewhere]],
                                      alternative="two-sided")
        depletion_direction = "cKO-depleted" if cko_frac < baseline_cko_frac else "cKO-enriched"

        s2eden_mean_wt = float(sub[sub.obs["genotype"] == "WT"].obs["score_secondary_eden"].mean()) if n_wt > 0 else 0.0
        s2eden_mean_cko = float(sub[sub.obs["genotype"] == "En1-cKO"].obs["score_secondary_eden"].mean()) if n_cko > 0 else 0.0
        sg_mean_wt = float(sub[sub.obs["genotype"] == "WT"].obs["score_sweat_gland"].mean()) if n_wt > 0 else 0.0
        sg_mean_cko = float(sub[sub.obs["genotype"] == "En1-cKO"].obs["score_sweat_gland"].mean()) if n_cko > 0 else 0.0
        eda_mean_wt = float(sub[sub.obs["genotype"] == "WT"].obs["score_eda_pathway"].mean()) if n_wt > 0 else 0.0
        eda_mean_cko = float(sub[sub.obs["genotype"] == "En1-cKO"].obs["score_eda_pathway"].mean()) if n_cko > 0 else 0.0

        genes_list = list(dermal.uns["rank_genes_groups"]["names"][cls][:10])
        lfc_list = list(dermal.uns["rank_genes_groups"]["logfoldchanges"][cls][:10])
        top_markers = ", ".join([f"{g}({lfc:+.1f})" for g, lfc in zip(genes_list, lfc_list)])

        rows.append({
            "leiden_cluster": cls,
            "n_cells": n_wt + n_cko,
            "n_WT": n_wt, "n_cKO": n_cko,
            "cko_frac": cko_frac,
            "baseline_cko_frac": baseline_cko_frac,
            "depletion_direction": depletion_direction,
            "fisher_p_two_sided": p_fisher,
            "odds_ratio": odds,
            "score_secondary_eden_WT_mean": s2eden_mean_wt,
            "score_secondary_eden_cKO_mean": s2eden_mean_cko,
            "score_sweat_gland_WT_mean": sg_mean_wt,
            "score_sweat_gland_cKO_mean": sg_mean_cko,
            "score_eda_pathway_WT_mean": eda_mean_wt,
            "score_eda_pathway_cKO_mean": eda_mean_cko,
            "top_wilcoxon_markers": top_markers,
        })
    return pd.DataFrame(rows), baseline_cko_frac


def call_primary_and_secondary(df, baseline_cko_frac):
    # secondary EDEN: highest score_secondary_eden_WT_mean AND cKO-depleted (Fisher p<0.05)
    df_wt_ordered = df.sort_values("score_secondary_eden_WT_mean", ascending=False)
    secondary_candidates = df_wt_ordered[
        (df_wt_ordered["depletion_direction"] == "cKO-depleted") &
        (df_wt_ordered["fisher_p_two_sided"] < 0.05)
    ]
    secondary = secondary_candidates.iloc[0]["leiden_cluster"] if len(secondary_candidates) > 0 else None

    # primary EDEN: cKO-depleted + LOW secondary_eden (S100a4-neg) + HIGH Eda_pathway (En1-responsive)
    df_ranked = df.copy()
    df_ranked["depletion_score"] = -np.log10(df_ranked["fisher_p_two_sided"].clip(lower=1e-300)) * \
                                   (df_ranked["cko_frac"] < baseline_cko_frac).astype(int)
    primary_score = df_ranked["depletion_score"] * \
                    (1.0 / (df_ranked["score_secondary_eden_WT_mean"].abs() + 0.01)) * \
                    (df_ranked["score_eda_pathway_WT_mean"] + 0.1)
    df_ranked["primary_eden_composite_score"] = primary_score
    df_ranked = df_ranked.sort_values("primary_eden_composite_score", ascending=False)
    primary_candidates = df_ranked[
        (df_ranked["depletion_direction"] == "cKO-depleted") &
        (df_ranked["fisher_p_two_sided"] < 0.05) &
        (df_ranked["leiden_cluster"] != secondary)
    ].head(3)
    return secondary, primary_candidates, df_ranked


def main():
    dermal = load_dingwall_with_v3_predictions()
    dermal = subcluster_dermal(dermal, resolution=1.5)
    dermal = score_modules(dermal)
    df, baseline_cko = per_subcluster_analysis(dermal)
    secondary, primary_cands, df_ranked = call_primary_and_secondary(df, baseline_cko)

    out = ROOT / "discovery/pan_skin/marker"
    out.mkdir(parents=True, exist_ok=True)
    df_ranked.to_csv(out / "100_primary_eden_discovery.csv", index=False)

    summary = {
        "target": "Dingwall_GSE220977",
        "hypothesis": "Primary EDEN precedes Secondary EDEN (S100a4+/Tnc+ cluster 20/Derm10) in dermal lineage",
        "method": "PANDA-v3 predicts dermal-fibroblast compartment; Leiden sub-clustering "
                  "(resolution=1.5) resolves substructure; Wilcoxon markers + Fisher-exact "
                  "cKO enrichment + module scoring (Secondary_EDEN, Sweat_gland, Eda_pathway) "
                  "rank sub-clusters for Primary EDEN candidacy",
        "baseline_cko_frac": float(baseline_cko),
        "n_subclusters": int(len(df)),
        "secondary_eden_call": {
            "leiden_cluster": str(secondary),
            "criteria": "highest S100a4+Tnc+Pdgfra score AND Fisher cKO-depleted p<0.05",
            "row": df[df["leiden_cluster"] == secondary].iloc[0].to_dict() if secondary else None,
        },
        "primary_eden_candidates_top3": primary_cands[[
            "leiden_cluster", "n_cells", "n_WT", "n_cKO", "cko_frac",
            "fisher_p_two_sided", "score_secondary_eden_WT_mean",
            "score_sweat_gland_WT_mean", "score_eda_pathway_WT_mean",
            "top_wilcoxon_markers", "primary_eden_composite_score",
        ]].to_dict("records") if len(primary_cands) > 0 else [],
    }
    (out / "100_primary_eden_summary.json").write_text(json.dumps(summary, indent=2, default=str))

    print(f"\n[eden] wrote {out}/100_primary_eden_*", flush=True)
    print(f"\n=== SECONDARY EDEN CALL ===", flush=True)
    print(f"  leiden cluster: {secondary}", flush=True)
    if secondary:
        row = df[df["leiden_cluster"] == secondary].iloc[0]
        print(f"  n={row['n_cells']} (WT {row['n_WT']} / cKO {row['n_cKO']}), "
              f"cko_frac={row['cko_frac']:.3f} vs baseline {baseline_cko:.3f}", flush=True)
        print(f"  Fisher p={row['fisher_p_two_sided']:.2e}, "
              f"score_secondary_eden WT={row['score_secondary_eden_WT_mean']:.3f}", flush=True)

    print(f"\n=== PRIMARY EDEN CANDIDATES (top 3) ===", flush=True)
    for _, row in primary_cands.iterrows():
        print(f"  leiden {row['leiden_cluster']} n={row['n_cells']} (WT {row['n_WT']} / cKO {row['n_cKO']}), "
              f"cko_frac={row['cko_frac']:.3f}, Fisher p={row['fisher_p_two_sided']:.2e}", flush=True)
        print(f"    S2EDEN_WT={row['score_secondary_eden_WT_mean']:.3f}, "
              f"Sweat_WT={row['score_sweat_gland_WT_mean']:.3f}, "
              f"Eda_WT={row['score_eda_pathway_WT_mean']:.3f}", flush=True)
        print(f"    top markers: {row['top_wilcoxon_markers']}", flush=True)


if __name__ == "__main__":
    main()