"""marker refinement on aldrich zero-shot predictions: per-class wilcoxon DE vs canonical list."""
from __future__ import annotations
from pathlib import Path
import warnings, yaml
warnings.filterwarnings("ignore")

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

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"))
TFS    = Path(str(PANDA_ROOT / "discovery/pan_skin/marker/known_skin_tfs.yaml"))
OUT    = Path(str(PANDA_ROOT / "discovery/pan_skin/marker"))


def main():
    a = ad.read_h5ad(TARGET)
    p = ad.read_h5ad(PROJ)
    a.obs["pred_label"] = p.obs["pred_label"].values
    a.obs["pred_conf"]  = p.obs["pred_conf"].astype(float).values

    with open(TFS) as f:
        tf_cfg = yaml.safe_load(f)
    canon = tf_cfg["classes"]

    keep = a[~a.obs["pred_label"].str.startswith("UNK", na=False)].copy()
    sc.tl.rank_genes_groups(keep, "pred_label", method="wilcoxon",
                            n_genes=100, use_raw=False)

    rows = []
    for cls in keep.obs["pred_label"].unique():
        try:
            names = keep.uns["rank_genes_groups"]["names"][cls]
            lfc   = keep.uns["rank_genes_groups"]["logfoldchanges"][cls]
            padj  = keep.uns["rank_genes_groups"]["pvals_adj"][cls]
        except Exception:
            continue
        top100 = list(names[:100])
        canon_set = set(canon.get(cls, []))
        retained = [g for g in canon_set if g in top100]
        novel = [g for g in top100 if g not in canon_set]
        mask = keep.obs["pred_label"] == cls
        n_cells = int(mask.sum())
        mean_conf = float(keep.obs.loc[mask, "pred_conf"].mean())
        rows.append({
            "class": cls,
            "n_cells_predicted": n_cells,
            "mean_conf": round(mean_conf, 3),
            "n_canonical": len(canon_set),
            "n_canonical_retained_top100": len(retained),
            "retained_canonical": ",".join(retained[:15]),
            "novel_top20": ",".join(novel[:20]),
        })
    df = pd.DataFrame(rows).sort_values("n_cells_predicted", ascending=False)
    df.to_csv(OUT / "52_refined_markers.csv", index=False)
    print(df.to_string(index=False))

    md = ["# Refined pan-skin markers for Aldrich zero-shot predictions\n"]
    md.append("Marker refinement is Wilcoxon DE of each predicted class against all other")
    md.append("predicted cells on raw Aldrich gene expression. `retained_canonical` are canonical")
    md.append("markers recovered in the top-100; `novel_top20` are DE genes not in the canonical list.\n")
    md.append(df.to_markdown(index=False))
    (OUT / "52_refined_markers.md").write_text("\n".join(md))
    print(f"[refine] wrote {OUT}/52_refined_markers.csv and .md")


if __name__ == "__main__":
    main()