"""Shared RNA inputs for the RNA foundation-model extractors (43_*.py).

load(name) -> (X, genes): X = scipy CSR float32 raw(-like) counts (n_spot, n_gene), genes = HGNC symbols (np.array str)
  hest     : HEST-1k Visium real counts (9484 spots, 14,802 common genes)             -> training / test of H&E<->RNA
  hest_v   : STPath virtual RNA of the same 9484 HEST spots (calibration only)
  orion_v  : STPath virtual RNA of the 3187 ORION tiles (test only)
  tcga_v   : STPath virtual RNA of the 6144 TCGA text patches (test only)
  dlpfc    : DLPFC external Visium real counts (482 spots, 33,538 genes)            -> external test only
STPath outputs log1p counts; virtual counts = round(expm1(pred)) so that the zero pattern means "not detected"
as in real Visium (models that tokenise expressed genes would otherwise see all 14,791 genes as expressed).
Output convention of every extractor: cache/rnafm/<model>_<set>.npy, float16 (n_spot, D), row order = the set's.
"""
import os

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

from common import CACHE

SETS = ["hest", "hest_v", "orion_v", "tcga_v", "dlpfc"]
OUT = os.path.join(CACHE, "rnafm")
_STP = {"hest_v": "hest_val", "orion_v": "orion", "tcga_v": "tcga_text"}


def load(name):
    if name == "hest":
        X = sp.load_npz(os.path.join(CACHE, "hest", "counts.npz")).tocsr().astype(np.float32)
        g = pd.read_csv(os.path.join(CACHE, "hest", "genes.csv")).gene.values
    elif name in _STP:
        p = np.load(os.path.join(CACHE, "stpath", f"{_STP[name]}_pred.npy")).astype(np.float32)
        X = sp.csr_matrix(np.rint(np.expm1(p).clip(0)).astype(np.float32))
        g = pd.read_csv(os.path.join(CACHE, "stpath", f"{_STP[name]}_genes.csv")).gene.values
    elif name == "dlpfc":
        X = sp.csr_matrix(np.load(os.path.join(CACHE, "external", "dlpfc_counts_full.npy")).astype(np.float32))
        g = pd.read_csv(os.path.join(CACHE, "external", "dlpfc_genes_full.csv")).gene.values
    else:
        raise KeyError(name)
    X.eliminate_zeros()
    return X, np.asarray(g, dtype=str)


def save(model, name, E):
    os.makedirs(OUT, exist_ok=True)
    E = np.asarray(E, dtype=np.float32)
    assert np.isfinite(E).all(), (model, name)
    np.save(os.path.join(OUT, f"{model}_{name}.npy"), E.astype(np.float16))
    print(f"[{model}] {name}: saved {E.shape}", flush=True)


if __name__ == "__main__":
    for s in SETS:
        X, g = load(s)
        nnz = np.diff(X.indptr)
        print(f"{s:8s} {X.shape}  genes/spot median {np.median(nnz):.0f}  UMI/spot median {np.median(np.asarray(X.sum(1)).ravel()):.0f}")