"""Build examples/: a small raw-input sample (fold-0 TEST items) + the cached training features of the same items, so
that example_check.py can test embed.py end to end without the full datasets.

    .venv/bin/python -I 54_make_examples.py [n=64]

examples/
  orion_he/*.png, orion_mif.npz (tiles uint8 (n, 16, 333, 333) in embed.MIF_MARKERS order, slide)   ORION H&E + mIF
  hest_he.npy (n, 224, 224, 3) uint8, hest_counts.npz (raw counts, spots x genes), hest_genes.csv          HEST H&E + RNA
  tcga_he/*.png, tcga_captions.csv                                                                    TCGA/PathGen H&E + text
  ref_feats.npz   cached encoder features (view 0) of the same items: the inputs the model was trained / tested on
  items.csv       which items (global indices into cache/)
"""
import os
import shutil
import sys

import numpy as np
import pandas as pd
import scipy.sparse as sp

ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, ROOT)
from common import CACHE, ORION, load_all_tiles  # noqa: E402
from embed import MIF_MARKERS, RNA_E, TXT_E, _load_script  # noqa: E402
from rna_sets import load  # noqa: E402

N = int(sys.argv[1]) if len(sys.argv) > 1 else 64
OUT = os.path.join(ROOT, "examples")
shutil.rmtree(OUT, ignore_errors=True)
os.makedirs(os.path.join(OUT, "orion_he")); os.makedirs(os.path.join(OUT, "tcga_he"))
TI = np.load(os.path.join(ROOT, "models", "ej_ad_mv_fold0", "test_items.npz"))
rng = np.random.default_rng(0)
ld = lambda *p: np.load(os.path.join(*p), mmap_mode="r")
ref, items = {}, []

# ORION: H&E PNG + 16-marker mIF
import tifffile  # noqa: E402

tiles = load_all_tiles()
ot = np.sort(rng.choice(TI["orion"], N, replace=False))
names = [n for n, _ in _load_script("14_extract_if_fm.py", [sys.argv[0], "virtues", "--no_probe"]).MARKERS]
ch = [names.index(m) for m in MIF_MARKERS]
mif = np.stack([tifffile.imread(os.path.join(ORION, p))[..., ch].transpose(2, 0, 1) for p in tiles.target_path.values[ot]])
for i, p in zip(ot, tiles.image_path.values[ot]):
    shutil.copy(os.path.join(ORION, p), os.path.join(OUT, "orion_he", f"{i:05d}.png"))
np.savez_compressed(os.path.join(OUT, "orion_mif.npz"), tiles=mif.astype(np.uint8), slide=tiles.in_slide_name.values[ot])
ref |= {"orion_uni2": ld(CACHE, "he_uni2.npy")[ot, 0], "orion_medvtok": ld(CACHE, "he_medvtok.npy")[ot, 0],
        "orion_virtues": ld(CACHE, "if_virtues_m16.npy")[ot, 0]}
items += [{"set": "orion", "idx": int(i)} for i in ot]

# HEST: spot crops + raw counts
hs = np.sort(rng.choice(TI["hest"], N, replace=False))
np.save(os.path.join(OUT, "hest_he.npy"), ld(CACHE, "hest", "he_img.npy")[hs])
X, genes = load("hest")
sp.save_npz(os.path.join(OUT, "hest_counts.npz"), sp.csr_matrix(X[hs]))
pd.DataFrame({"gene": genes}).to_csv(os.path.join(OUT, "hest_genes.csv"), index=False)
ref |= {"hest_uni2": ld(CACHE, "hest", "he_uni2.npy")[hs, 0], "hest_medvtok": ld(CACHE, "hest", "he_medvtok.npy")[hs, 0],
        **{f"hest_{e}": (ld(CACHE, "hest", "hvg.npy") if e == "hvg" else ld(CACHE, "rnafm", f"{e}_hest.npy"))[hs]
           for e in RNA_E}}
items += [{"set": "hest", "idx": int(i)} for i in hs]

# TCGA / PathGen: patch PNG + its paired caption
k = np.sort(rng.choice(len(TI["tcga_img"]), N, replace=False))
ti, tt = TI["tcga_img"][k], TI["tcga_txt"][k]
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
for i in ti:
    shutil.copy(I.image_path.values[i], os.path.join(OUT, "tcga_he", f"{i:05d}.png"))
pd.DataFrame({"img_idx": ti, "txt_idx": tt, "file": [f"{i:05d}.png" for i in ti],
              "caption": T.caption.fillna("").astype(str).values[tt]}).to_csv(os.path.join(OUT, "tcga_captions.csv"), index=False)
ref |= {"tcga_uni2": ld(CACHE, "allimg_he_uni2.npy")[ti, 0], "tcga_medvtok": ld(CACHE, "allimg_he_medvtok.npy")[ti, 0],
        **{f"tcga_{e}": ld(CACHE, f"alltxt_text_{e}.npy")[tt, 0] for e in TXT_E}}
items += [{"set": "tcga", "idx": int(i), "txt_idx": int(j)} for i, j in zip(ti, tt)]

np.savez_compressed(os.path.join(OUT, "ref_feats.npz"), **{k_: np.asarray(v, np.float32) for k_, v in ref.items()})
pd.DataFrame(items).to_csv(os.path.join(OUT, "items.csv"), index=False)
os.system(f"du -sh {OUT}")