"""marker-scored 18-class labels per corpus cell. writes canonical_label + leiden to corpus.h5ad."""
from __future__ import annotations
from pathlib import Path
import warnings, yaml, sys
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import scanpy as sc
import 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/pan_skin/harmonized/corpus.h5ad"))
YAML   = Path(str(PANDA_ROOT / "scripts/pan_skin/known_skin_tfs.yaml"))


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"]
    resolution = cfg["assignment"]["cluster_resolution"]

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

    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"):
        print(f"[label] {ds_name}: n={sub.shape[0]}")
        idx = sub.index
        sa = a[idx].copy()

        sc.pp.neighbors(sa, use_rep="X_pca", n_neighbors=15)
        sc.tl.leiden(sa, resolution=resolution, 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"  {ds_name} label breakdown:",
              pd.Series(assigned).value_counts().to_dict())

    print("[label] corpus-wide breakdown:")
    print(a.obs.groupby(["dataset", "canonical_label"], observed=True).size().unstack(fill_value=0))

    a.write_h5ad(CORPUS, compression="gzip")
    print(f"[label] wrote {CORPUS}")


if __name__ == "__main__":
    main()