"""en1-cKO vs WT contrast on aldrich predictions."""
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 import stats

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"))
NOVEL  = Path(str(PANDA_ROOT / "discovery/pan_skin/marker/51_aldrich_novel_annotation.csv"))
OUT    = Path(str(PANDA_ROOT / "discovery/pan_skin/marker"))


def main():
    a = ad.read_h5ad(TARGET)
    p = ad.read_h5ad(PROJ)
    for c in ["pred_label", "pred_conf", "abstain", "pred_bbse_label"]:
        if c in p.obs.columns:
            a.obs[c] = p.obs[c].values
    if NOVEL.exists():
        nov = pd.read_csv(NOVEL, index_col=0)
        a.obs["novel_cluster"] = nov.reindex(a.obs_names)["novel_cluster"].values

    print(f"[cko] target: {a.shape}, genotype: {a.obs['genotype'].value_counts().to_dict()}",
          flush=True)

    rows = []
    for cls in a.obs["pred_bbse_label"].unique():
        mask = a.obs["pred_bbse_label"] == cls
        n_cko = int((mask & (a.obs["genotype"] == "En1-cKO")).sum())
        n_wt  = int((mask & (a.obs["genotype"] == "WT")).sum())
        total_cko = (a.obs["genotype"] == "En1-cKO").sum()
        total_wt  = (a.obs["genotype"] == "WT").sum()
        contingency = np.array([[n_cko, total_cko - n_cko],
                                [n_wt,  total_wt  - n_wt]])
        odds, p_val = stats.fisher_exact(contingency)
        f_cko = (n_cko + 1) / (total_cko + 2)
        f_wt  = (n_wt  + 1) / (total_wt  + 2)
        log2_fc = np.log2(f_cko / f_wt)
        rows.append({
            "class": cls,
            "n_En1cKO": n_cko,
            "n_WT": n_wt,
            "pct_En1cKO": round(100 * n_cko / total_cko, 2),
            "pct_WT": round(100 * n_wt / total_wt, 2),
            "log2_fold_enrich_cKO_vs_WT": round(log2_fc, 3),
            "fisher_pvalue": p_val,
        })
    df = pd.DataFrame(rows).sort_values("log2_fold_enrich_cKO_vs_WT")
    print("\n[cko] class enrichment (cKO vs WT):")
    print(df.to_string(index=False))

    df.to_csv(OUT / "53_en1_cko_class_enrichment.csv", index=False)

    de_rows = []
    for cls in sorted(a.obs["pred_bbse_label"].unique()):
        cls_mask = a.obs["pred_bbse_label"] == cls
        if cls_mask.sum() < 50:
            continue
        sub = a[cls_mask].copy()
        vc = sub.obs["genotype"].value_counts()
        if not {"En1-cKO", "WT"}.issubset(vc.index) or vc.min() < 15:
            continue
        try:
            sc.tl.rank_genes_groups(sub, "genotype", method="wilcoxon",
                                    n_genes=40, use_raw=False)
            for grp in ["En1-cKO", "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])
                for g, lf in zip(names, lfcs):
                    de_rows.append({
                        "class": cls, "up_in": grp, "gene": g, "logfc": round(float(lf), 3),
                    })
        except Exception as exc:
            print(f"[cko] DE failed for {cls}: {exc}")
            continue
    de_df = pd.DataFrame(de_rows)
    de_df.to_csv(OUT / "53_en1_cko_wilcoxon_within_class.csv", index=False)

    md = ["# Aldrich En1-cKO vs WT contrast on PANDA-MLP predictions\n",
          f"Total cells: {a.n_obs:,} ({int((a.obs['genotype']=='En1-cKO').sum())} En1-cKO, "
          f"{int((a.obs['genotype']=='WT').sum())} WT).\n",
          "## Class-level cKO/WT enrichment (BBSE-corrected predictions)\n",
          df.to_markdown(index=False), "",
          "The direction of `log2_fold_enrich_cKO_vs_WT` indicates whether a class is over-represented",
          "in En1-cKO (positive) or WT (negative). Fisher exact p-value tests significance vs the",
          "background genotype ratio (~40% cKO / 60% WT).\n",
          "## Per-class Wilcoxon DE (En1-cKO vs WT within each class)\n",
          "Top genes differentially expressed BETWEEN genotypes WITHIN a predicted class. Genes up in",
          "cKO reveal En1-loss-responsive programs specific to that cell type; genes up in WT are the",
          "opposite.\n",
          ]
    if len(de_df):
        for cls in sorted(de_df["class"].unique()):
            md.append(f"\n### {cls}\n")
            for grp in ["En1-cKO", "WT"]:
                sub = de_df[(de_df["class"] == cls) & (de_df["up_in"] == grp)]
                if not len(sub):
                    continue
                md.append(f"**Up in {grp}**: " + ", ".join(sub["gene"].tolist()))
    md += ["", "## Novel population x genotype cross-tab\n"]
    if "novel_cluster" in a.obs.columns:
        xt = pd.crosstab(a.obs["novel_cluster"], a.obs["genotype"])
        md.append(xt.to_markdown())
    (OUT / "53_en1_cko_contrast.md").write_text("\n".join(md))
    print(f"[cko] wrote {OUT}/53_en1_cko_contrast.md")


if __name__ == "__main__":
    main()