"""melanocyte deep dive on aldrich: within-melanocyte cKO vs WT DE and pathway scoring."""
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"))

PATHWAYS = {
    "MITF_regulon":       ["Mitf", "Dct", "Tyr", "Pmel", "Mlana", "Tyrp1", "Slc24a5",
                           "Slc45a2", "Silv", "Sox10", "Pax3", "Kit", "Ednrb"],
    "Wnt_signaling":      ["Wnt3", "Wnt5a", "Wnt7a", "Wnt10b", "Ctnnb1", "Lef1", "Tcf4",
                           "Tcf7", "Axin2", "Dkk1", "Sfrp1", "Fzd7", "Lrp5", "Lrp6"],
    "BMP_TGF_signaling":  ["Bmp2", "Bmp4", "Bmp5", "Bmp7", "Bmpr1a", "Bmpr1b", "Bmpr2",
                           "Smad1", "Smad3", "Smad5", "Tgfb1", "Tgfb2", "Tgfbr1",
                           "Tgfbr2", "Id1", "Id2", "Id3"],
    "FGF_signaling":      ["Fgf1", "Fgf2", "Fgf7", "Fgf9", "Fgf10", "Fgf20", "Fgfr1",
                           "Fgfr2", "Fgfr3", "Fgfr4", "Etv1", "Etv4", "Etv5", "Spry2",
                           "Spry4", "Dusp6"],
    "Notch_signaling":    ["Notch1", "Notch2", "Notch3", "Jag1", "Jag2", "Dll1", "Dll3",
                           "Dll4", "Hes1", "Hes5", "Hey1", "Hey2", "Rbpj", "Maml1"],
    "Hedgehog":           ["Shh", "Ihh", "Ptch1", "Ptch2", "Smo", "Gli1", "Gli2", "Gli3",
                           "Hhip", "Sufu"],
    "EMT":                ["Zeb1", "Zeb2", "Snai1", "Snai2", "Twist1", "Twist2", "Vim",
                           "Cdh2", "Fn1", "Mmp2", "Mmp9", "Prrx1"],
    "Cell_cycle":         ["Ccnd1", "Ccnd2", "Ccne1", "Ccna2", "Ccnb1", "Cdk1", "Cdk2",
                           "Cdk4", "Cdk6", "Mki67", "Top2a", "Pcna", "Mcm2", "Mcm3"],
    "Neural_crest":       ["Sox10", "Sox9", "Sox2", "Pax3", "Foxd3", "Nes", "Ngfr",
                           "Ednrb", "Kit", "Tfap2a"],
    "Apoptosis":          ["Bax", "Bak1", "Bad", "Bid", "Bcl2", "Bcl2l1", "Casp3",
                           "Casp9", "Casp8", "Fas", "Fasl", "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"{name}"] = 0.0
            continue
        sc.tl.score_genes(sub, gene_list=present, score_name=f"{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

    mask = a.obs["pred_label"] == "melanocyte"
    print(f"[mel] melanocyte cells: {int(mask.sum())} "
          f"({int((mask & (a.obs['genotype']=='En1-cKO')).sum())} cKO, "
          f"{int((mask & (a.obs['genotype']=='WT')).sum())} WT)", flush=True)
    if mask.sum() < 30:
        print("[mel] too few melanocyte cells")
        return

    sub = a[mask].copy()
    sub = pathway_scoring(sub, PATHWAYS)

    print("\n[mel] pathway score cKO vs WT (positive = up in cKO):")
    rows = []
    for name in PATHWAYS.keys():
        s = sub.obs[name].astype(float).values
        g = sub.obs["genotype"].values
        c_scores = s[g == "En1-cKO"]
        w_scores = s[g == "WT"]
        stat, pval = mannwhitneyu(c_scores, w_scores, alternative="two-sided")
        delta = c_scores.mean() - w_scores.mean()
        rows.append({"pathway": name, "delta_cKO_minus_WT": round(delta, 4),
                     "MannU_p": pval,
                     "cKO_mean": round(c_scores.mean(), 4),
                     "WT_mean": round(w_scores.mean(), 4)})
        star = "***" if pval < 0.001 else "**" if pval < 0.01 else "*" if pval < 0.05 else ""
        print(f"  {name:20s} delta={delta:+.4f} p={pval:.2e} {star}")

    pd.DataFrame(rows).sort_values("MannU_p").to_csv(OUT / "56_melanocyte_pathways.csv", index=False)

    print("\n[mel] within-melanocyte cKO vs WT DE:")
    sc.tl.rank_genes_groups(sub, "genotype", method="wilcoxon", n_genes=50, use_raw=False)
    de = pd.DataFrame({
        "cKO_up_gene": sub.uns["rank_genes_groups"]["names"]["En1-cKO"][:20],
        "cKO_up_lfc":  sub.uns["rank_genes_groups"]["logfoldchanges"]["En1-cKO"][:20],
        "WT_up_gene":  sub.uns["rank_genes_groups"]["names"]["WT"][:20],
        "WT_up_lfc":   sub.uns["rank_genes_groups"]["logfoldchanges"]["WT"][:20],
    })
    print(de.to_string(index=False))
    de.to_csv(OUT / "56_melanocyte_wilcoxon_cko_vs_wt.csv", index=False)

    lines = ["# Melanocyte deep dive — Aldrich En1-cKO vs WT",
             "",
             f"PANDA-MLP-predicted melanocyte cells: **{int(mask.sum())} total** "
             f"({int((mask & (a.obs['genotype']=='En1-cKO')).sum())} cKO, "
             f"{int((mask & (a.obs['genotype']=='WT')).sum())} WT).",
             "",
             "The v3b+v3c cross-model replication showed ~2× melanocyte enrichment in cKO "
             "(log2fc +1.04, Fisher p ≈ 4×10⁻⁶). Here we probe what's happening WITHIN the "
             "melanocyte compartment.",
             "",
             "## Pathway score comparisons (MannU cKO vs WT)",
             ""]
    dfp = pd.DataFrame(rows).sort_values("MannU_p")
    lines.append(dfp.to_markdown(index=False))
    lines += ["",
              "## Within-melanocyte Wilcoxon DE (top 20 each direction)",
              ""]
    de_show = de.copy()
    de_show["cKO_up_lfc"] = de_show["cKO_up_lfc"].round(3)
    de_show["WT_up_lfc"]  = de_show["WT_up_lfc"].round(3)
    lines.append(de_show.to_markdown(index=False))
    (OUT / "56_melanocyte_deep_dive.md").write_text("\n".join(lines))
    print(f"[mel] wrote {OUT}/56_melanocyte_deep_dive.md")


if __name__ == "__main__":
    main()