"""per-class pathway scoring cKO vs WT on aldrich, MannU per (class, pathway)."""
from __future__ import annotations
from pathlib import Path
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import anndata as ad
import scanpy as sc
from scipy.stats import mannwhitneyu

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
TARGET = Path(str(PANDA_ROOT / "data/processed/skin/adata_processed.h5ad"))
PROJ   = Path(str(PANDA_ROOT / "discovery/pan_skin/marker/50_aldrich_projections.h5ad"))
OUT    = Path(str(PANDA_ROOT / "discovery/pan_skin/marker"))

CLASSES_OF_INTEREST = [
    "basal-IFE", "spinous", "granular",
    "fibroblast-papillary", "fibroblast-reticular",
    "endothelial", "immune", "melanocyte",
]

PATHWAYS = {
    "MITF_regulon":         ["Mitf", "Dct", "Tyr", "Pmel", "Mlana", "Tyrp1", "Slc24a5",
                             "Slc45a2", "Sox10", "Pax3", "Kit", "Ednrb"],
    "Wnt_signaling":        ["Wnt3", "Wnt5a", "Wnt7a", "Wnt10b", "Ctnnb1", "Lef1", "Tcf4",
                             "Tcf7", "Axin2", "Dkk1", "Sfrp1", "Fzd7", "Lrp5"],
    "BMP_signaling":        ["Bmp2", "Bmp4", "Bmp5", "Bmp7", "Bmpr1a", "Bmpr1b", "Bmpr2",
                             "Smad1", "Smad5", "Id1", "Id2", "Id3"],
    "TGFB_signaling":       ["Tgfb1", "Tgfb2", "Tgfbr1", "Tgfbr2", "Smad3", "Smad7"],
    "FGF_signaling":        ["Fgf1", "Fgf2", "Fgf7", "Fgf9", "Fgf10", "Fgfr1", "Fgfr2",
                             "Etv1", "Etv4", "Etv5", "Spry2", "Dusp6"],
    "Notch_signaling":      ["Notch1", "Notch2", "Notch3", "Jag1", "Dll1", "Hes1", "Hes5",
                             "Hey1", "Hey2", "Rbpj"],
    "Hedgehog":             ["Shh", "Ptch1", "Smo", "Gli1", "Gli2", "Gli3"],
    "Eda_ectodysplasin":    ["Eda", "Edar", "Edaradd", "Nfkb1", "Nfkb2", "Rela"],
    "EMT":                  ["Zeb1", "Zeb2", "Snai1", "Snai2", "Twist1", "Twist2", "Vim",
                             "Cdh2", "Fn1", "Prrx1"],
    "Cell_cycle":           ["Ccnd1", "Ccne1", "Ccna2", "Ccnb1", "Cdk1", "Cdk2", "Cdk4",
                             "Mki67", "Top2a", "Pcna", "Mcm2", "Mcm3"],
    "KC_differentiation":   ["Krt1", "Krt10", "Ivl", "Lor", "Flg", "Flg2", "Klk5", "Klk7",
                             "Cdsn"],
    "Basal_keratinocyte":   ["Krt5", "Krt14", "Krt15", "Trp63", "Itga6", "Itgb1", "Itga3"],
    "Sweat_gland":          ["Foxi3", "Foxa1", "En1", "Krt8", "Krt18", "Krt19", "Muc5b",
                             "Aqp5", "Cutl1"],
    "Hair_placode":         ["Wnt10b", "Shh", "Lef1", "Foxi3", "Edar", "Bmp4", "Msx2"],
    "Neural_crest":         ["Sox10", "Sox9", "Sox2", "Pax3", "Foxd3", "Nes", "Tfap2a"],
    "Apoptosis":            ["Bax", "Bak1", "Bad", "Bcl2", "Casp3", "Casp9", "Trp53",
                             "Cdkn1a"],
}


def pathway_scoring(sub, pathway_dict):
    for name, genes in pathway_dict.items():
        present = [g for g in genes if g in sub.var_names]
        if not present:
            sub.obs[f"pw_{name}"] = 0.0
            continue
        sc.tl.score_genes(sub, gene_list=present, score_name=f"pw_{name}",
                          random_state=0, use_raw=False)
    return sub


def main():
    a = ad.read_h5ad(TARGET)
    p = ad.read_h5ad(PROJ)
    a.obs["pred_label"] = p.obs["pred_bbse_label"].values
    print(f"[pw] classes in target: {a.obs['pred_label'].value_counts().to_dict()}", flush=True)

    rows = []
    for cls in CLASSES_OF_INTEREST:
        mask = a.obs["pred_label"] == cls
        n_c = int((mask & (a.obs["genotype"]=="En1-cKO")).sum())
        n_w = int((mask & (a.obs["genotype"]=="WT")).sum())
        if n_c < 15 or n_w < 15:
            print(f"[pw] {cls}: skip (n_cKO={n_c}, n_WT={n_w})")
            continue
        sub = a[mask].copy()
        sub = pathway_scoring(sub, PATHWAYS)
        for pw in PATHWAYS.keys():
            s = sub.obs[f"pw_{pw}"].astype(float).values
            g = sub.obs["genotype"].values
            cvals = s[g=="En1-cKO"]; wvals = s[g=="WT"]
            try:
                _, pval = mannwhitneyu(cvals, wvals, alternative="two-sided")
            except Exception:
                pval = 1.0
            delta = cvals.mean() - wvals.mean()
            rows.append({"class": cls, "pathway": pw,
                         "n_cKO": n_c, "n_WT": n_w,
                         "delta_cKO_minus_WT": round(delta, 4),
                         "MannU_p": pval})
        print(f"[pw] {cls}: {n_c} cKO, {n_w} WT — scored")

    df = pd.DataFrame(rows)
    df.to_csv(OUT / "57_pathway_class_by_pathway.csv", index=False)

    pivot_delta = df.pivot(index="pathway", columns="class", values="delta_cKO_minus_WT")
    pivot_p     = df.pivot(index="pathway", columns="class", values="MannU_p")
    def stars(p): return "***" if p<0.001 else "**" if p<0.01 else "*" if p<0.05 else ""
    disp = pivot_delta.copy().astype(object)
    for pw in disp.index:
        for c in disp.columns:
            d = pivot_delta.loc[pw, c]; p = pivot_p.loc[pw, c]
            if pd.isna(d): disp.loc[pw, c] = ""
            else: disp.loc[pw, c] = f"{d:+.3f}{stars(p)}"

    lines = ["# Pathway score contrasts by class — Aldrich En1-cKO vs WT",
             "",
             "Delta = cKO mean − WT mean of `sc.tl.score_genes` pathway score.",
             "Sig: * p<0.05, ** p<0.01, *** p<0.001 (MannU two-sided).\n",
             disp.to_markdown()]
    (OUT / "57_pathway_class_by_pathway.md").write_text("\n".join(lines))
    print(f"[pw] wrote {OUT}/57_pathway_class_by_pathway.md")

    df_sig = df[df["MannU_p"] < 0.01].sort_values("MannU_p")
    print("\n[pw] Strongest cKO/WT pathway shifts (p<0.01):")
    print(df_sig[["class","pathway","delta_cKO_minus_WT","MannU_p"]].to_string(index=False))
    df_sig.to_csv(OUT / "57_pathway_top_hits.csv", index=False)


if __name__ == "__main__":
    main()