"""materialise held-out h5ad slices for nestorowa (hsc) and sulic (skin) using anchor split rules."""
from pathlib import Path
import warnings, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp
warnings.filterwarnings("ignore"); sc.settings.verbosity = 0

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
ROOT = Path(str(PANDA_ROOT))


def make_nestorowa_test():
    src = ROOT / "data/raw/nestorowa_combined.h5ad"
    print(f"[nestorowa] loading {src}", flush=True)
    a = ad.read_h5ad(src)
    # anchor split from scripts/hematopoiesis/09_retrain_with_nestorowa_anchor.py:
    # 150 LT-HSC + 600 HSPC stratified (seed=0); rest is test.
    rng = np.random.default_rng(0)
    lt = np.where(a.obs["cell_type"].astype(str).values == "LT-HSC")[0]
    hs = np.where(a.obs["cell_type"].astype(str).values == "HSPC")[0]
    lt_anchor = rng.choice(lt, size=min(150, len(lt)), replace=False)
    hs_anchor = rng.choice(hs, size=min(600, len(hs)), replace=False)
    anchor_ix = np.concatenate([lt_anchor, hs_anchor])
    test_ix = np.setdiff1d(np.arange(a.n_obs), anchor_ix)
    a_test = a[test_ix].copy()
    # coerce obs cols to string so h5ad write survives mixed dtypes
    for c in list(a_test.obs.columns):
        try: a_test.obs[c] = a_test.obs[c].astype(str)
        except Exception: del a_test.obs[c]
    out_dir = ROOT / "data/corpus/hematopoiesis/held_out_labeled"
    out_dir.mkdir(parents=True, exist_ok=True)
    out = out_dir / "nestorowa_GSE81682_test.h5ad"
    a_test.write_h5ad(out)
    print(f"[nestorowa] wrote {out} ({a_test.n_obs:,} cells; "
          f"gate dist: {a_test.obs['cell_type'].value_counts().to_dict()})", flush=True)


def make_sulic_test():
    corpus = ROOT / "data/corpus/pan_skin/harmonized/corpus.h5ad"
    print(f"[sulic] loading {corpus} (legacy, has sulic + labels)", flush=True)
    a = ad.read_h5ad(corpus)
    a_sul = a[a.obs["dataset"] == "sulic_GSE212673"].copy()
    print(f"[sulic] {a_sul.n_obs} sulic cells "
          f"({a_sul.obs['canonical_label'].value_counts().to_dict()})", flush=True)
    # anchor split from scripts/pan_skin/92_retrain_with_sulic_anchor.py:
    # 300 HF-placode + 200 basal-IFE stratified (seed=0); rest (4183) is test.
    rng = np.random.default_rng(0)
    hf = np.where(a_sul.obs["canonical_label"].astype(str).values == "HF-placode")[0]
    bi = np.where(a_sul.obs["canonical_label"].astype(str).values == "basal-IFE")[0]
    hf_anchor = rng.choice(hf, size=min(300, len(hf)), replace=False)
    bi_anchor = rng.choice(bi, size=min(200, len(bi)), replace=False)
    anchor_ix = np.concatenate([hf_anchor, bi_anchor])
    test_ix = np.setdiff1d(np.arange(a_sul.n_obs), anchor_ix)
    a_test = a_sul[test_ix].copy()
    for c in list(a_test.obs.columns):
        try: a_test.obs[c] = a_test.obs[c].astype(str)
        except Exception: del a_test.obs[c]
    out_dir = ROOT / "data/corpus/pan_skin/held_out_labeled"
    out_dir.mkdir(parents=True, exist_ok=True)
    out = out_dir / "sulic_GSE212673_test.h5ad"
    a_test.write_h5ad(out)
    print(f"[sulic] wrote {out} ({a_test.n_obs:,} cells; "
          f"labels: {a_test.obs['canonical_label'].value_counts().to_dict()})", flush=True)


if __name__ == "__main__":
    make_nestorowa_test()
    make_sulic_test()