"""veres marker deep-dive.

per-predicted-class wilcoxon on raw veres counts, cross-referenced with canonical
adult-beta / SC-alpha / EP panels. output: discovery/pancreas/marker/91_veres_marker_deep_dive.csv
"""
from pathlib import Path
import warnings, 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/pancreas/marker"
OUT.mkdir(parents=True, exist_ok=True)

# human panels (Veres is human hPSC)
PANELS = {
    "adult-beta":                  ["INS", "MAFA", "UCN3", "NKX6-1", "MNX1", "NEUROD1", "PDX1"],
    "adult-alpha":                 ["GCG", "ARX", "IRX2", "IRX1", "MAFB", "TTR"],
    "alpha (embryonic prototype)": ["GCG", "ARX", "IRX2", "MAFB"],
    "beta (embryonic prototype)":  ["INS", "NKX6-1", "MNX1", "NEUROD1", "PDX1"],
    "delta":                       ["SST", "HHEX", "LEPR"],
    "gamma":                       ["PPY", "PYY", "SLC38A4"],
    "epsilon":                     ["GHRL"],
    "endocrine-progenitor-early":  ["NEUROG3", "CBFA2T3", "BTBD17"],
    "endocrine-progenitor-Fev":    ["FEV", "INSM1"],
    "endocrine-progenitor-primed": ["PAX4", "ARX"],
    "acinar":                      ["PRSS1", "PRSS2", "CEL", "CTRB1"],
    "ductal":                      ["KRT19", "SOX9", "MUC1"],
    "endothelial":                 ["PECAM1", "CDH5", "KDR"],
    "immune":                      ["PTPRC", "CD68"],
    "mesenchymal":                 ["COL1A1", "COL3A1", "DCN"],
}

# Load Veres via existing loader logic
def load_veres():
    SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract"
    parts = []
    for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")):
        counts_file = str(meta_file).replace("cell_metadata", "processed_counts")
        if not Path(counts_file).exists(): continue
        meta = pd.read_csv(meta_file, sep="\t", compression="gzip")
        counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0)
        obs = meta.set_index("library.barcode")
        obs = obs.loc[obs.index.intersection(counts.index)]
        counts_al = counts.loc[obs.index]
        X = sp.csr_matrix(counts_al.values.astype(np.float32))
        a = ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=counts_al.columns))
        a.var_names_make_unique()
        parts.append(a)
    return ad.concat(parts, join="outer")

print("[load] Veres + predictions", flush=True)
raw = load_veres()
pred_df = pd.read_csv(ROOT / "discovery/pancreas/marker/veres_predictions.csv")
# strip the "veres_" prefix from prediction cell_ids so they align with raw.obs_names
pred_df["cell_id"] = pred_df["cell_id"].astype(str).str.replace(r"^veres_", "", regex=True)
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])
print(f"[align] {raw.n_obs} cells across {raw.obs['pred_label'].nunique()} classes", flush=True)

counts_s = raw.obs["pred_label"].value_counts()
keep_cls = counts_s[counts_s >= 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] {raw.n_obs} cells × {len(keep_cls)} classes", flush=True)

sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw)
print("[wilcoxon] running...", flush=True)
sc.tl.rank_genes_groups(raw, groupby="pred_label", method="wilcoxon", n_genes=25, use_raw=False)

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)}"
    best_panel = max(panel_hits.items(),
                     key=lambda x: int(x[1].split("/")[0]) / (int(x[1].split(":")[0].split("/")[1]) + 1e-6))

    # Stage enrichment
    stage_col = raw.obs.get("Stage", raw.obs.get("stage", pd.Series([""]*raw.n_obs, index=raw.obs.index)))
    stage_vals = pd.to_numeric(stage_col[mask], errors="coerce")
    top_stage = int(stage_vals.mode().iloc[0]) if len(stage_vals.dropna()) else -1
    stage6_frac = float((stage_vals == 6).sum() / max(1, mask.sum()))

    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],
        "top_stage": top_stage,
        "stage6_frac": stage6_frac,
    })

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