"""primary + secondary EDEN on Dingwall via score_genes against dingwall's own Derm0-11 markers (Data S1C top-30); argmax identity + Fisher cKO depletion."""
from pathlib import Path
import warnings, json, sys, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp
from scipy.stats import fisher_exact
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))
CKO_GSMS = {"GSM6833482", "GSM6833483"}  # CORRECTED: 480/481 are rttaControl (WT), not cKO
WT_GSMS  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}  # CORRECTED: 4 Cre-neg controls per GEO metadata
DERM_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-3.xlsx"
TOP_N = 30  # top N markers per Derm cluster for scoring

# EDEN identity map based on Data S1C + Data S2 CellChat
EDEN_IDENTITY = {
    10: "Secondary_EDEN",
    2: "Primary_EDEN_candidate_1",
    9: "Primary_EDEN_candidate_2",
    3: "EDEN_signalling",
    6: "EDEN_signalling",
}


def load_derm_marker_panels():
    df = pd.read_excel(DERM_MARKERS)
    df = df.sort_values(["cluster", "avg_log2FC"], ascending=[True, False])
    panels = {}
    for cl in sorted(df["cluster"].unique()):
        genes = df[df["cluster"] == cl].head(TOP_N)["gene"].tolist()
        panels[int(cl)] = genes
    return panels


def load_dingwall_v3_dermal():
    raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad")
    pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv")
    pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
    raw.obs["pred_label"] = np.array([pred_map.get(c, "unknown") for c in raw.obs_names.astype(str)])
    raw.obs["genotype"] = np.where(raw.obs["sample"].astype(str).isin(list(CKO_GSMS)), "En1-cKO",
                          np.where(raw.obs["sample"].astype(str).isin(list(WT_GSMS)), "WT", "other"))
    labeled = raw.obs["genotype"].isin(["WT", "En1-cKO"]).values
    raw = raw[labeled].copy()
    dermal_mask = np.isin(raw.obs["pred_label"], ["fibroblast-papillary", "fibroblast-reticular"])
    dermal = raw[dermal_mask].copy()
    return dermal


def score_derm_identities(a, panels):
    for cl, genes in panels.items():
        present = [g for g in genes if g in a.var_names]
        if not present:
            a.obs[f"derm{cl}_score"] = 0.0
            continue
        sc.tl.score_genes(a, gene_list=present, score_name=f"derm{cl}_score",
                          random_state=0, use_raw=False)
    return a


def main():
    print("[eden] loading Data S1C marker panels", flush=True)
    panels = load_derm_marker_panels()
    print(f"[eden] loaded {len(panels)} Derm panels", flush=True)

    print("[eden] loading Dingwall + PANDA-v3 dermal predictions", flush=True)
    dermal = load_dingwall_v3_dermal()
    print(f"[eden] {dermal.n_obs} dermal-fibroblast cells", flush=True)

    print("[eden] normalise + log1p", flush=True)
    sc.pp.normalize_total(dermal, target_sum=1e4); sc.pp.log1p(dermal)

    print("[eden] scoring cells on all 12 Derm identity panels", flush=True)
    dermal = score_derm_identities(dermal, panels)

    score_cols = [f"derm{cl}_score" for cl in sorted(panels.keys())]
    score_mat = dermal.obs[score_cols].values
    argmax = np.argmax(score_mat, axis=1)
    derm_ids = [int(score_cols[i].replace("derm", "").replace("_score", "")) for i in argmax]
    dermal.obs["derm_identity"] = derm_ids
    dermal.obs["max_derm_score"] = score_mat.max(axis=1)

    n_wt_tot = int((dermal.obs["genotype"] == "WT").sum())
    n_cko_tot = int((dermal.obs["genotype"] == "En1-cKO").sum())
    baseline_cko_frac = n_cko_tot / (n_wt_tot + n_cko_tot)
    print(f"[eden] baseline WT={n_wt_tot} cKO={n_cko_tot} (baseline cKO={baseline_cko_frac:.3f})", flush=True)

    rows = []
    for cl in sorted(panels.keys()):
        sub = dermal[dermal.obs["derm_identity"] == cl]
        n_wt = int((sub.obs["genotype"] == "WT").sum())
        n_cko = int((sub.obs["genotype"] == "En1-cKO").sum())
        if n_wt + n_cko == 0:
            continue
        cko_frac = n_cko / (n_wt + n_cko) if (n_wt + n_cko) > 0 else 0
        n_wt_else = n_wt_tot - n_wt
        n_cko_else = n_cko_tot - n_cko
        try:
            odds, p_fisher = fisher_exact([[n_wt, n_wt_else], [n_cko, n_cko_else]],
                                          alternative="two-sided")
        except ValueError:
            odds, p_fisher = 1.0, 1.0
        rows.append({
            "derm_id": cl,
            "identity": EDEN_IDENTITY.get(cl, "other"),
            "n_cells": n_wt + n_cko,
            "n_WT": n_wt, "n_cKO": n_cko,
            "cko_frac": cko_frac,
            "baseline_cko_frac": baseline_cko_frac,
            "cko_delta": cko_frac - baseline_cko_frac,
            "fisher_p": float(p_fisher),
            "odds_ratio": float(odds) if not np.isnan(odds) else None,
            "depletion_direction": "cKO-depleted" if cko_frac < baseline_cko_frac
                                   else "cKO-enriched",
            "top10_markers_dingwall": ", ".join(panels[cl][:10]),
        })
    df = pd.DataFrame(rows).sort_values("cko_delta")

    out = ROOT / "discovery/pan_skin/marker"
    out.mkdir(parents=True, exist_ok=True)
    df.to_csv(out / "101_derm_subcluster_scores.csv", index=False)

    secondary_row = df[df["derm_id"] == 10].iloc[0].to_dict() if 10 in df["derm_id"].values else None
    primary_1_row = df[df["derm_id"] == 2].iloc[0].to_dict() if 2 in df["derm_id"].values else None
    primary_2_row = df[df["derm_id"] == 9].iloc[0].to_dict() if 9 in df["derm_id"].values else None

    summary = {
        "target": "Dingwall_GSE220977",
        "method": "Score PANDA-v3 dermal-fibroblast predictions on Dingwall's own Derm0-11 "
                  "marker panels (Data S1C top-30 genes each); argmax identity per cell; "
                  "Fisher-exact cKO depletion per Derm identity",
        "baseline_cko_frac": baseline_cko_frac,
        "n_dermal_cells_total": int(dermal.n_obs),
        "secondary_eden_Derm10": secondary_row,
        "primary_eden_Derm2": primary_1_row,
        "primary_eden_Derm9": primary_2_row,
        "all_derm_summary": df.to_dict("records"),
    }
    (out / "101_derm_identity_summary.json").write_text(json.dumps(summary, indent=2, default=str))

    print(f"\n[eden] wrote {out}/101_derm_*", flush=True)
    print(f"\n=== SECONDARY EDEN (Derm10) ===", flush=True)
    if secondary_row:
        print(f"  n={secondary_row['n_cells']} (WT {secondary_row['n_WT']} / cKO {secondary_row['n_cKO']}), "
              f"cko_frac={secondary_row['cko_frac']:.3f} vs baseline {baseline_cko_frac:.3f}", flush=True)
        print(f"  {secondary_row['depletion_direction']}, Fisher p={secondary_row['fisher_p']:.2e}", flush=True)

    print(f"\n=== PRIMARY EDEN candidate 1 (Derm2) ===", flush=True)
    if primary_1_row:
        print(f"  n={primary_1_row['n_cells']} (WT {primary_1_row['n_WT']} / cKO {primary_1_row['n_cKO']}), "
              f"cko_frac={primary_1_row['cko_frac']:.3f}", flush=True)
        print(f"  {primary_1_row['depletion_direction']}, Fisher p={primary_1_row['fisher_p']:.2e}", flush=True)

    print(f"\n=== PRIMARY EDEN candidate 2 (Derm9) ===", flush=True)
    if primary_2_row:
        print(f"  n={primary_2_row['n_cells']} (WT {primary_2_row['n_WT']} / cKO {primary_2_row['n_cKO']}), "
              f"cko_frac={primary_2_row['cko_frac']:.3f}", flush=True)
        print(f"  {primary_2_row['depletion_direction']}, Fisher p={primary_2_row['fisher_p']:.2e}", flush=True)


if __name__ == "__main__":
    main()