"""dingwall UMAP: (a) PANDA-predicted class, (b) en1 genotype."""
from pathlib import Path
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import anndata as ad
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import umap

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
FIG = Path(str(PANDA_ROOT / "figures"))
FIG.mkdir(exist_ok=True)

CLASS_COLORS = {
    "fibroblast-reticular": "#66c2a5",
    "basal-IFE": "#fc8d62",
    "endothelial": "#8da0cb",
    "immune": "#e78ac3",
    "fibroblast-papillary": "#a6d854",
    "melanocyte": "#ffd92f",
    "spinous": "#e5c494",
    "granular": "#b3b3b3",
    "HF-ORS": "#1b9e77",
    "HF-DP": "#d95f02",
    "HF-placode": "#7570b3",
    "eccrine-duct": "#666666",
    "eccrine-placode": "#000000",
}


def main():
    p = ad.read_h5ad(str(PANDA_ROOT / "discovery/pan_skin/marker/50_aldrich_projections.h5ad"))
    print(f"[umap] Dingwall projection shape {p.shape}", flush=True)
    if "Z_projection" not in p.obsm:
        print(f"[umap] no Z_projection in obsm; keys: {list(p.obsm.keys())}")
        return
    Z = np.asarray(p.obsm["Z_projection"])
    print(f"[umap] Z shape {Z.shape}", flush=True)

    print(f"[umap] fitting UMAP …", flush=True)
    reducer = umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=42,
                       metric="cosine", n_components=2)
    emb = reducer.fit_transform(Z)
    print(f"[umap] UMAP done, emb shape {emb.shape}", flush=True)

    pred = p.obs["pred_bbse_label"] if "pred_bbse_label" in p.obs else p.obs.get("pred_label")
    genotype = p.obs.get("genotype", pd.Series(index=p.obs.index, data="unknown"))

    fig, axes = plt.subplots(1, 2, figsize=(14, 6))

    ax = axes[0]
    for cls in pred.unique():
        m = pred == cls
        ax.scatter(emb[m, 0], emb[m, 1], s=1.5, alpha=0.4,
                   c=CLASS_COLORS.get(cls, "#999999"), label=f"{cls} (n={int(m.sum())})")
    ax.set_title("Dingwall (25,344 cells) — PANDA-predicted class (BBSE)")
    ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")
    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=7,
              markerscale=6, frameon=False)

    ax = axes[1]
    gcolors = {"WT": "#2b83ba", "En1-cKO": "#d7191c", "unknown": "#999999"}
    for g in ["WT", "En1-cKO"]:
        m = genotype == g
        ax.scatter(emb[m, 0], emb[m, 1], s=1.5, alpha=0.35,
                   c=gcolors[g], label=f"{g} (n={int(m.sum())})")
    ax.set_title("Dingwall — En1 genotype")
    ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2")
    ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", markerscale=6, frameon=False)

    plt.tight_layout()
    plt.savefig(FIG / "fig5_dingwall_umap.pdf", bbox_inches="tight")
    plt.close()
    print(f"[umap] wrote {FIG}/fig5_dingwall_umap.pdf")


if __name__ == "__main__":
    main()