"""dingwall marker deep-dive: single rank_genes_groups call cross-referenced against canonical panels."""
from pathlib import Path
import warnings, json, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp
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])))
ROOT = Path(str(PANDA_ROOT))
OUT = ROOT / "discovery/pan_skin/marker"
OUT.mkdir(parents=True, exist_ok=True)

PANELS = {
    "eden-dermal-niche": ["S100a4", "Twist2", "Prrx1", "Pdgfra", "Fap", "Fn1"],
    "eccrine-secretory": ["Dcd", "Aqp5", "Muc7", "Cst6", "Krt7"],
    "eccrine-ductal":    ["Krt77", "Krt5", "Krt14", "Cldn6", "Grhl3"],
    "basal-multipotent": ["Krt5", "Krt14", "Trp63", "Itgb4", "Sox2"],
    "hair-placode":      ["Shh", "Sox9", "Lhx2", "Foxi3", "Wnt10a"],
    "melanocyte":        ["Dct", "Mlana", "Tyrp1", "Pmel", "Sox10"],
    "endothelial":       ["Pecam1", "Cdh5", "Kdr", "Flt1"],
    "spinous":           ["Krt10", "Krt1", "Dsp"],
    "basal-IFE":         ["Krt5", "Krt14", "Krt15", "Col17a1"],
    "immune":            ["Ptprc", "Cd68", "Cd3d", "Cd19"],
    "fibroblast":        ["Col1a1", "Dcn", "Pdgfra"],
}

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

print("[load] Dingwall raw + predictions", flush=True)
raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
pred_df = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv")
common = raw.obs_names.intersection(pd.Index(pred_df["cell_id"].astype(str)))
raw = raw[list(common)].copy()
pred_map = dict(zip(pred_df["cell_id"].astype(str), pred_df["pred_label"]))
raw.obs["pred_label"] = pd.Categorical([pred_map.get(c, "unknown") for c in raw.obs_names])
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"))
print(f"[align] {raw.n_obs} cells across {raw.obs['pred_label'].nunique()} classes", flush=True)

# subset to classes with >=30 cells for stable Wilcoxon
counts = raw.obs["pred_label"].value_counts()
keep_cls = counts[counts >= 30].index.tolist()
raw = raw[raw.obs["pred_label"].isin(keep_cls)].copy()
raw.obs["pred_label"] = raw.obs["pred_label"].astype(str).astype("category")
print(f"[filter] kept {raw.n_obs} cells × {len(keep_cls)} classes", flush=True)

sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw)

# single-call with groupby is much faster than per-class loop
print("[wilcoxon] single-call across all predicted classes...", flush=True)
sc.tl.rank_genes_groups(raw, groupby="pred_label", method="wilcoxon", n_genes=25, use_raw=False)
print("[wilcoxon] done", flush=True)

rows = []
for cls in raw.uns["rank_genes_groups"]["names"].dtype.names:
    mask = raw.obs["pred_label"] == cls
    if mask.sum() < 30: continue
    genes = list(raw.uns["rank_genes_groups"]["names"][cls][:20])
    pvals = [float(x) for x in raw.uns["rank_genes_groups"]["pvals_adj"][cls][:20]]
    logfc = [float(x) for x in raw.uns["rank_genes_groups"]["logfoldchanges"][cls][:20]]

    top_str = ",".join([f"{g}(LFC{lf:+.1f})" for g, lf in zip(genes[:10], logfc[:10])])
    panel_hits = {}
    for pname, plist in PANELS.items():
        hits = [g for g in plist if g in genes[:20]]
        panel_hits[pname] = f"{len(hits)}/{len(plist)}: {','.join(hits)}"
    gt = raw.obs["genotype"][mask]
    ncko = int((gt == "En1-cKO").sum()); nwt = int((gt == "WT").sum())
    frac_cko = ncko / max(1, ncko + nwt)
    best_panel = max(panel_hits.items(),
                     key=lambda x: int(x[1].split("/")[0]) / (int(x[1].split(":")[0].split("/")[1]) + 1e-6))

    rows.append({
        "predicted_class": cls,
        "n_cells": int(mask.sum()),
        "top_wilcoxon_markers": top_str,
        "min_p_adj_top5": min(pvals[:5], default=float("nan")),
        "best_canonical_panel_match": best_panel[0],
        "recovery": best_panel[1],
        "n_En1_cKO": ncko,
        "n_WT": nwt,
        "frac_En1_cKO": frac_cko,
    })

df = pd.DataFrame(rows).sort_values("n_cells", ascending=False)
df.to_csv(OUT / "90_dingwall_marker_deep_dive.csv", index=False)
print(f"[write] {OUT}/90_dingwall_marker_deep_dive.csv ({len(df)} classes)", flush=True)
print()
print(df[["predicted_class", "n_cells", "best_canonical_panel_match", "recovery", "frac_En1_cKO"]].to_string(index=False))