"""within-class wilcoxon DE WT vs Kit_W41 + module scoring; checks dahlin paper claims (Myc/ISR)."""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys
warnings.filterwarnings("ignore")
import numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp
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])))
DAHLIN = Path(str(PANDA_ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract"))
PRED = Path(str(PANDA_ROOT / "discovery/hematopoiesis/marker/66_dahlin_predictions.csv"))
OUT = Path(str(PANDA_ROOT / "discovery/hematopoiesis/marker"))


def load_dahlin_all():
    parts = []
    for f in sorted(DAHLIN.glob("*.txt.gz")):
        gsm = f.name.split("_")[0]
        sample = f.name.split("_")[1].split(".")[0]
        df = pd.read_csv(f, sep="\t", compression="gzip", index_col=0)
        X = sp.csr_matrix(df.values.T.astype(np.float32))
        obs = pd.DataFrame(index=[f"{sample}_{bc}" for bc in df.columns.astype(str)])
        obs["sample"] = sample
        var = pd.DataFrame(index=df.index.astype(str))
        var["ensmusg"] = var.index.values
        a = ad.AnnData(X=X, obs=obs, var=var)
        parts.append(a)
    return ad.concat(parts, join="outer", label="_batch")


def main():
    print("[dahlin-mech] loading Dahlin …", flush=True)
    a = load_dahlin_all()

    import mygene
    mg = mygene.MyGeneInfo()
    ids = a.var_names.astype(str).tolist()
    print(f"[dahlin-mech] mapping {len(ids)} genes …", flush=True)
    res = mg.querymany(ids, scopes="ensembl.gene", fields="symbol", species="mouse",
                       verbose=False)
    id2sym = {r["query"]: r["symbol"] for r in res if "symbol" in r}
    syms = pd.Series(a.var_names.astype(str)).map(id2sym).values
    keep = pd.notna(syms)
    a = a[:, keep].copy(); a.var_names = syms[keep]
    a.var_names_make_unique()
    print(f"[dahlin-mech] converted, shape now {a.shape}", flush=True)

    preds = pd.read_csv(PRED, index_col=0)
    a.obs = a.obs.join(preds[["pred_label", "pred_conf", "genotype"]], how="left")
    print(f"[dahlin-mech] joined preds; genotype counts {a.obs['genotype'].value_counts().to_dict()}",
          flush=True)

    a_de = a.copy()
    sc.pp.normalize_total(a_de, target_sum=1e4)
    sc.pp.log1p(a_de)

    rows = []
    lg_class_min = 300
    class_counts = a.obs["pred_label"].value_counts()
    for cls in class_counts[class_counts >= lg_class_min].index:
        mask = a.obs["pred_label"] == cls
        sub = a_de[mask].copy()
        if len(sub.obs["genotype"].unique()) < 2:
            continue
        gc = sub.obs["genotype"].value_counts()
        if gc.min() < 30:
            continue
        try:
            sc.tl.rank_genes_groups(sub, "genotype", method="wilcoxon", n_genes=40,
                                    use_raw=False)
            for grp in ["Kit_W41", "WT"]:
                if grp not in sub.uns["rank_genes_groups"]["names"].dtype.names:
                    continue
                names = list(sub.uns["rank_genes_groups"]["names"][grp][:15])
                lfcs = list(sub.uns["rank_genes_groups"]["logfoldchanges"][grp][:15])
                padjs = list(sub.uns["rank_genes_groups"]["pvals_adj"][grp][:15])
                for g, l, p in zip(names, lfcs, padjs):
                    rows.append({"class": cls, "up_in": grp, "gene": g,
                                 "logfc": round(float(l), 3),
                                 "padj": float(p)})
        except Exception as exc:
            print(f"[dahlin-mech] DE failed for {cls}: {exc}")
    pd.DataFrame(rows).to_csv(OUT / "67_dahlin_within_class_de.csv", index=False)

    MODULES = {
        "MYC_targets":       ["Myc", "Mycn", "Nme1", "Odc1", "Ncl", "Ppat", "Cad",
                              "Ldha", "Nop58"],
        "Integrated_stress": ["Atf4", "Ddit3", "Chac1", "Trib3", "Asns", "Gdf15",
                              "Sesn2", "Slc7a11"],
        "Apoptosis_pro":     ["Bax", "Bak1", "Bad", "Bid", "Bbc3", "Pmaip1", "Casp3",
                              "Casp9", "Trp53"],
        "Apoptosis_anti":    ["Bcl2", "Bcl2l1", "Mcl1", "Bcl2l2"],
        "Erythroid_dev":     ["Klf1", "Gata1", "Gata2", "Epor", "Hba-a1", "Hbb-b1"],
        "Cell_cycle":        ["Ccnd1", "Ccnd2", "Ccne1", "Ccna2", "Ccnb1", "Cdk1",
                              "Mki67", "Top2a", "Pcna", "Mcm2", "Mcm3", "Mcm4"],
        "Kit_signaling":     ["Kit", "Kitl", "Sos1", "Rasgrf1", "Grb2", "Pik3ca"],
    }
    for name, genes in MODULES.items():
        present = [g for g in genes if g in a_de.var_names]
        if not present: continue
        sc.tl.score_genes(a_de, gene_list=present, score_name=f"pw_{name}",
                          random_state=0, use_raw=False)

    print("\n[dahlin-mech] Kit_W41 vs WT module score contrast per class:", flush=True)
    mod_rows = []
    for cls in class_counts[class_counts >= lg_class_min].index:
        mask = a_de.obs["pred_label"] == cls
        for name in MODULES.keys():
            col = f"pw_{name}"
            if col not in a_de.obs: continue
            s = a_de.obs.loc[mask, col].astype(float).values
            g = a_de.obs.loc[mask, "genotype"].astype(str).values
            cv = s[g == "Kit_W41"]; wv = s[g == "WT"]
            if len(cv) < 20 or len(wv) < 20: continue
            try:
                _, p = mannwhitneyu(cv, wv, alternative="two-sided")
            except Exception: p = 1.0
            delta = cv.mean() - wv.mean()
            star = "***" if p < 1e-6 else "**" if p < 1e-3 else "*" if p < 0.05 else ""
            print(f"  {cls:15s} {name:20s} delta={delta:+.4f} p={p:.2e} {star}")
            mod_rows.append({"class": cls, "module": name,
                             "delta_Kit_minus_WT": round(delta, 4),
                             "MannU_p": p})
    pd.DataFrame(mod_rows).to_csv(OUT / "67_dahlin_module_scores.csv", index=False)

    print(f"\n[dahlin-mech] complete. Outputs in {OUT}/")


if __name__ == "__main__":
    main()