"""rank dingwall cell types by En1-cKO vs WT DEG count (|log2FC|>1, padj<0.05)."""
from __future__ import annotations

import json
import warnings
from pathlib import Path

import anndata as ad
import numpy as np
import pandas as pd
import scanpy as sc

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)

RAW  = ROOT / "data/raw/GSE220977_combined.h5ad"
PRED = ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv"
CKO_GSMS = {"GSM6833482", "GSM6833483"}
WT_GSMS  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}
MIN_PER_GROUP = 15
LFC_THRESH = 1.0
PADJ_THRESH = 0.05


def main():
    print("[load]", flush=True)
    a = ad.read_h5ad(RAW)
    pred = pd.read_csv(PRED)
    pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
    a.obs["pred_label"] = pd.Categorical(
        [pred_map.get(c, "unknown") for c in a.obs_names])
    samp = a.obs["sample"].astype(str)
    a.obs["group"] = np.where(samp.isin(list(CKO_GSMS)), "En1-cKO",
                     np.where(samp.isin(list(WT_GSMS)), "WT", "other"))
    a = a[a.obs["group"].isin(["En1-cKO", "WT"])].copy()
    print(f"[load] {a.n_obs} cells, {a.n_vars} genes", flush=True)

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

    rows = []
    for cls in sorted(a.obs["pred_label"].astype(str).unique()):
        mask = (a.obs["pred_label"].astype(str) == cls)
        n_cko = int((mask & (a.obs["group"] == "En1-cKO")).sum())
        n_wt  = int((mask & (a.obs["group"] == "WT")).sum())
        if n_cko < MIN_PER_GROUP or n_wt < MIN_PER_GROUP:
            print(f"[skip] {cls}: n_cKO={n_cko} n_WT={n_wt}", flush=True)
            continue
        sub = a[mask].copy()
        sub.obs["group"] = pd.Categorical(sub.obs["group"].values,
                                          categories=["En1-cKO", "WT"])
        # reference=WT so positive LFC means up in cKO
        try:
            sc.tl.rank_genes_groups(sub, "group", reference="WT",
                                    groups=["En1-cKO"], method="wilcoxon",
                                    n_genes=sub.n_vars, use_raw=False,
                                    pts=True)
            rgg = sub.uns["rank_genes_groups"]
            df = pd.DataFrame({
                "gene":  [x[0] for x in rgg["names"]],
                "lfc":   [x[0] for x in rgg["logfoldchanges"]],
                "padj":  [x[0] for x in rgg["pvals_adj"]],
            })
            up   = int(((df["lfc"] >  LFC_THRESH) & (df["padj"] < PADJ_THRESH)).sum())
            down = int(((df["lfc"] < -LFC_THRESH) & (df["padj"] < PADJ_THRESH)).sum())
            total = up + down
            # persist the COMPLETE significant sets. An earlier .head(10) meant
            # n_DEG could exceed the stored gene list (HF-placode: 12 vs 10),
            # so two of the reported DEGs were never written to any artefact.
            top_up   = df[(df["lfc"] >  LFC_THRESH) & (df["padj"] < PADJ_THRESH)] \
                        .sort_values("lfc", ascending=False)["gene"].tolist()
            top_down = df[(df["lfc"] < -LFC_THRESH) & (df["padj"] < PADJ_THRESH)] \
                        .sort_values("lfc", ascending=True)["gene"].tolist()
        except Exception as e:
            print(f"[fail] {cls}: {e}", flush=True)
            continue
        rows.append({
            "class":    cls,
            "n_cKO":    n_cko,
            "n_WT":     n_wt,
            "n_DEG":    total,
            "n_up":     up,
            "n_down":   down,
            "top_up":   ";".join(top_up),
            "top_down": ";".join(top_down),
        })
        print(f"[ok] {cls}: n_cKO={n_cko} n_WT={n_wt} DEG={total} "
              f"(up={up}, down={down})", flush=True)

    df_out = pd.DataFrame(rows).sort_values("n_DEG", ascending=False)
    df_out.to_csv(OUT / "107_dingwall_class_deg_count.csv", index=False)
    print("\n[rank] classes ordered by DEG count (|LFC|>1 padj<0.05):")
    print(df_out[["class", "n_cKO", "n_WT", "n_DEG", "n_up", "n_down"]]
          .to_string(index=False))

    if not df_out.empty:
        winner = df_out.iloc[0]
        summary = {
            "lfc_threshold": LFC_THRESH, "padj_threshold": PADJ_THRESH,
            "n_classes_tested": int(len(df_out)),
            "top_class":  str(winner["class"]),
            "top_n_DEG":  int(winner["n_DEG"]),
            "top_up":     winner["top_up"],
            "top_down":   winner["top_down"],
            "ranking":    df_out[["class", "n_DEG"]].to_dict("records"),
        }
    else:
        summary = {"error": "no eligible classes"}
    with open(OUT / "107_dingwall_class_deg_count.json", "w") as f:
        json.dump(summary, f, indent=2)
    print("[done]", flush=True)


if __name__ == "__main__":
    main()