"""assign canonical pancreas labels via marker scoring."""
from __future__ import annotations
from pathlib import Path
import warnings, yaml
warnings.filterwarnings("ignore")
import numpy as np, pandas as pd, scanpy as sc, anndata as ad

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
CORPUS = Path(str(PANDA_ROOT / "data/corpus/pancreas/harmonized/corpus.h5ad"))
YAML = Path(str(PANDA_ROOT / "scripts/pancreas/known_pancreas_markers.yaml"))

CURATED_MAP = {
    # Bastidas cell_type
    "alpha": "alpha", "beta": "beta", "delta": "delta",
    "progenitor": "endocrine-progenitor", "other": "other",
    # Baron/Muraro/Grun
    "gamma": "gamma", "epsilon": "epsilon", "ductal": "ductal", "acinar": "acinar",
    "endothelial": "endothelial", "immune": "immune", "activated_stellate": "other",
    "quiescent_stellate": "other", "schwann": "other", "mast": "immune", "PP": "gamma",
    "T_cell": "immune", "macrophage": "immune",
}


def main():
    with open(YAML) as f: cfg = yaml.safe_load(f)
    classes = list(cfg["classes"].keys())
    markers = cfg["classes"]
    min_score = cfg["assignment"]["min_score"]
    min_margin = cfg["assignment"]["min_margin"]
    res = cfg["assignment"]["cluster_resolution"]

    a = ad.read_h5ad(CORPUS)
    print(f"[label] corpus: {a.shape}", flush=True)

    a.obs["canonical_label"] = pd.Categorical(["UNK"] * a.n_obs, categories=classes + ["UNK"])
    a.obs["leiden"] = "0"

    for ds_name, sub in a.obs.groupby("dataset"):
        idx = sub.index
        sa = a[idx].copy()
        sc.pp.neighbors(sa, use_rep="X_pca", n_neighbors=15)
        sc.tl.leiden(sa, resolution=res, key_added="leiden")
        a.obs.loc[idx, "leiden"] = ds_name + "_" + sa.obs["leiden"].astype(str)
        for c, m in markers.items():
            present = [g for g in m if g in sa.var_names]
            if not present: sa.obs[f"score_{c}"] = -np.inf; continue
            sc.tl.score_genes(sa, gene_list=present, score_name=f"score_{c}",
                              random_state=0, use_raw=False)
        S = sa.obs[[f"score_{c}" for c in classes]].values
        cluster_labels = sa.obs["leiden"].values
        assigned = np.array(["UNK"] * sa.n_obs, dtype=object)
        for cl in np.unique(cluster_labels):
            mask = cluster_labels == cl
            mean_scores = S[mask].mean(axis=0)
            order = np.argsort(mean_scores)[::-1]
            top, second = mean_scores[order[0]], mean_scores[order[1]]
            if top >= min_score and (top - second) >= min_margin:
                assigned[mask] = classes[order[0]]
        a.obs.loc[idx, "canonical_label"] = assigned
        print(f"[label] {ds_name}: {pd.Series(assigned).value_counts().to_dict()}", flush=True)

    # prefer paper cell_type when present
    if "cell_type" in a.obs.columns:
        mapped = a.obs["cell_type"].astype(str).map(CURATED_MAP)
        ok = mapped.notna()
        a.obs.loc[ok, "canonical_label"] = mapped[ok].astype(a.obs["canonical_label"].dtype)
        print(f"[label] curated override: {int(ok.sum())} cells", flush=True)

    print("[label] final corpus breakdown:")
    print(a.obs.groupby(["dataset", "canonical_label"], observed=True).size().unstack(fill_value=0))
    a.write_h5ad(CORPUS, compression="gzip")


if __name__ == "__main__":
    main()