scRep_pretrain/dataset.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

from collections import OrderedDict
import os

import torch
import anndata as ad
import numpy as np
import scipy.sparse as sp
from torch.utils.data import IterableDataset, get_worker_info

from .bin_cache import load_value_bin_cache


# =============================================================================
# Debug / Logging
# =============================================================================
DEBUG_PRINT: bool = False


def _log(msg: str) -> None:
    if DEBUG_PRINT:
        print(f"[sc-token] {msg}")


def _safe_close_anndata(adata_obj: Any) -> None:
    f = getattr(adata_obj, "file", None)
    if f is not None:
        try:
            f.close()
        except Exception:
            pass


# =============================================================================
# Per-file materialized reader
# =============================================================================
class H5ADFileReader:
    """
    Per-file reader.

    Matrix semantics:
      - use_raw=True  : adata.raw.X counts -> log1p(CP10K)
      - use_raw=False : adata.X is assumed to already be log1p(CP10K)

    Design notes:
      - This reader always returns ALL observed nonzero genes for a cell.
      - No top-k truncation happens here anymore.
      - Top-k is now applied later inside the collator.
      - Gene order is preserved here; any shuffling happens in the collator.
    """

    def __init__(
        self,
        path: str,
        dataset_id: int,
        *,
        gene_name_to_id: Dict[str, int],
        meta_maps: Dict[str, Dict[str, int]],
        meta_keys: List[str],
        use_raw: bool,
        cp_target_sum: float = 1e4,
        backed: Optional[str] = "r",
        prefer_csr: bool = True,
        strict_gene: bool = True,
        filter_missing_genes: bool = False,
        value_bin_cache_dir: str = "",
        require_value_bin_cache: bool = False,
    ) -> None:
        self.path = path
        self.dataset_id = int(dataset_id)

        self.gene_name_to_id = gene_name_to_id
        self.meta_maps = meta_maps
        self.meta_keys = list(meta_keys)

        self.use_raw = bool(use_raw)
        self.cp_target_sum = float(cp_target_sum)
        self.prefer_csr = bool(prefer_csr)
        self.filter_missing_genes = bool(filter_missing_genes)
        self.value_bin_cache_dir = str(value_bin_cache_dir or "")
        self.require_value_bin_cache = bool(require_value_bin_cache)
        self._bin_cache_indptr = None
        self._bin_cache_ids = None

        _log(f"open: {path} (use_raw={self.use_raw}, backed={backed})")
        self.adata = ad.read_h5ad(path, backed=backed)
        self.n_cells = int(self.adata.n_obs)

        if self.use_raw:
            if (self.adata.raw is None) or (self.adata.raw.X is None):
                raise ValueError(f"[{path}] use_raw=True but adata.raw.X missing.")
            self.X = self.adata.raw.X
            self.var_names = self.adata.raw.var_names
        else:
            if self.adata.X is None:
                raise ValueError(f"[{path}] use_raw=False but adata.X missing.")
            self.X = self.adata.X
            self.var_names = self.adata.var_names

        self.local_to_global = self._build_local_to_global(strict=strict_gene)
        self.meta_int_cols = self._build_meta_int_cols()
        self._load_value_bin_cache()

        self._is_csr = sp.isspmatrix_csr(self.X)
        if self._is_csr and self.prefer_csr:
            self._csr_indptr = self.X.indptr
            self._csr_indices = self.X.indices
            self._csr_data = self.X.data

        _log(f"[{path}] n_cells={self.n_cells}")

    def close(self) -> None:
        _safe_close_anndata(self.adata)

    def _load_value_bin_cache(self) -> None:
        if not self.value_bin_cache_dir:
            return
        cache = load_value_bin_cache(
            self.value_bin_cache_dir,
            self.path,
            expected_n_cells=self.n_cells,
            mmap_mode="r",
        )
        if cache is None:
            if self.require_value_bin_cache:
                raise FileNotFoundError(
                    f"Missing value bin cache for {self.path} under {self.value_bin_cache_dir}"
                )
            return
        self._bin_cache_indptr = cache["indptr"]
        self._bin_cache_ids = cache["bin_ids"]

    def has_value_bin_cache(self) -> bool:
        return self._bin_cache_indptr is not None and self._bin_cache_ids is not None

    def get_cached_bin_ids(self, cell_idx: int) -> Optional[np.ndarray]:
        if not self.has_value_bin_cache():
            return None
        i = int(cell_idx)
        start = int(self._bin_cache_indptr[i])
        end = int(self._bin_cache_indptr[i + 1])
        return np.asarray(self._bin_cache_ids[start:end], dtype=np.int64)

    def _build_local_to_global(self, *, strict: bool) -> np.ndarray:
        ids = np.fromiter(
            (self.gene_name_to_id.get(str(g), 0) for g in list(self.var_names)),
            dtype=np.int64,
            count=len(self.var_names),
        )
        if strict and int(np.sum(ids == 0)) > 0:
            missing = int(np.sum(ids == 0))
            raise ValueError(
                f"[{self.path}] {missing} genes missing in global gene vocab. "
                "Rebuild vocab with full manifest or check gene naming consistency."
            )
        return ids

    def _build_meta_int_cols(self) -> Dict[str, np.ndarray]:
        out: Dict[str, np.ndarray] = {}
        for k in self.meta_keys:
            if k not in self.adata.obs.columns:
                raise ValueError(f"[{self.path}] obs has no column '{k}'")
            if k not in self.meta_maps:
                raise ValueError(f"meta_maps missing key='{k}'. Build meta vocabs first.")

            gmap = self.meta_maps[k]
            col = self.adata.obs[k].astype("category")
            cats = col.cat.categories.astype(str).tolist()
            codes = col.cat.codes.to_numpy()

            cat_ids = np.array([gmap.get(c, 0) for c in cats], dtype=np.int32)
            ids = np.zeros((self.n_cells,), dtype=np.int32)
            valid = codes >= 0
            ids[valid] = cat_ids[codes[valid]]
            out[k] = ids
        return out

    def _csr_row_indices_data(self, cell_idx: int) -> Tuple[np.ndarray, np.ndarray]:
        i = int(cell_idx)
        start = int(self._csr_indptr[i])
        end = int(self._csr_indptr[i + 1])
        return self._csr_indices[start:end], self._csr_data[start:end]

    def _row_indices_data(self, cell_idx: int) -> Tuple[np.ndarray, np.ndarray]:
        if self._is_csr and self.prefer_csr:
            idx, val = self._csr_row_indices_data(cell_idx)
        else:
            row = self.X[int(cell_idx)]
            if sp.issparse(row):
                if not sp.isspmatrix_csr(row):
                    row = row.tocsr()
                idx, val = row.indices, row.data
            else:
                dense = np.asarray(row).ravel()
                idx = np.flatnonzero(dense)
                val = dense[idx]

        if val.size > 0:
            keep = val != 0
            if not np.all(keep):
                idx = idx[keep]
                val = val[keep]

        return idx, val

    def get_cell_tokens(self, cell_idx: int) -> Tuple[List[int], List[float]]:
        idx, val = self._row_indices_data(cell_idx)

        if idx.size == 0:
            return [], []

        row_sum = float(val.sum())
        if self.use_raw:
            if row_sum <= 0:
                return [], []
            expr = np.log1p(val.astype(np.float32) * (self.cp_target_sum / row_sum))
        else:
            expr = val.astype(np.float32)

        gene_ids = self.local_to_global[idx.astype(np.int64)].astype(np.int64)
        if self.filter_missing_genes:
            keep = gene_ids > 0
            if not np.all(keep):
                gene_ids = gene_ids[keep]
                expr = expr[keep]
            if gene_ids.size == 0:
                return [], []
        return gene_ids.tolist(), expr.tolist()

    def get_cell_observation_stats(self, cell_idx: int) -> Tuple[float, int]:
        idx, val = self._row_indices_data(cell_idx)
        detected_genes = int(idx.size)
        if self.use_raw:
            return float(val.sum()), detected_genes

        obs = self.adata.obs
        row = obs.iloc[int(cell_idx)]

        total_counts = None
        for key in ("total_counts", "n_counts"):
            if key not in obs.columns:
                continue
            candidate = float(row[key])
            if np.isfinite(candidate):
                total_counts = candidate
                break
        if total_counts is None:
            total_counts = float(detected_genes)

        for key in ("n_genes_by_counts", "n_genes"):
            if key not in obs.columns:
                continue
            candidate = float(row[key])
            if np.isfinite(candidate):
                detected_genes = int(candidate)
                break
        return total_counts, detected_genes

    def get_meta_batch(self, cell_ids: np.ndarray, *, include_dataset_id: bool = False) -> Dict[str, List[int]]:
        ids = cell_ids.astype(np.int64)
        out: Dict[str, List[int]] = {}

        if include_dataset_id:
            out["dataset_id"] = [self.dataset_id + 1] * int(ids.size)

        for k in self.meta_keys:
            out[k] = self.meta_int_cols[k][ids].astype(np.int64).tolist()

        return out


# =============================================================================
# Global pseudo-batch style index metadata
# =============================================================================
@dataclass
class _PseudoBatchGroupIndex:
    """
    One synthetic group inside one file.

    For scRep we keep one group containing all cells in the file.
    """

    raw_cell_ids: np.ndarray   # int32
    n_batches: int
    usable_cells: int


@dataclass
class _FileIndex:
    path: str
    dataset_id: int
    groups: List[_PseudoBatchGroupIndex]


# =============================================================================
# Small LRU reader cache
# =============================================================================
class _ReaderCache:
    """
    Small LRU cache of open H5ADFileReader objects.

    Each worker / rank gets its own cache inside __iter__.
    """

    def __init__(
        self,
        *,
        max_open_files: int,
        gene_name_to_id: Dict[str, int],
        meta_maps: Dict[str, Dict[str, int]],
        meta_keys: List[str],
        use_raw: bool,
        cp_target_sum: float,
        backed: Optional[str],
        prefer_csr: bool,
        strict_gene: bool,
        filter_missing_genes: bool,
        value_bin_cache_dir: str,
        require_value_bin_cache: bool,
    ) -> None:
        self.max_open_files = int(max_open_files)
        self.gene_name_to_id = gene_name_to_id
        self.meta_maps = meta_maps
        self.meta_keys = meta_keys
        self.use_raw = bool(use_raw)
        self.cp_target_sum = float(cp_target_sum)
        self.backed = backed
        self.prefer_csr = bool(prefer_csr)
        self.strict_gene = bool(strict_gene)
        self.filter_missing_genes = bool(filter_missing_genes)
        self.value_bin_cache_dir = str(value_bin_cache_dir or "")
        self.require_value_bin_cache = bool(require_value_bin_cache)

        self._cache: "OrderedDict[int, H5ADFileReader]" = OrderedDict()

    def get(self, file_id: int, file_meta: _FileIndex) -> H5ADFileReader:
        if file_id in self._cache:
            reader = self._cache.pop(file_id)
            self._cache[file_id] = reader
            return reader

        reader = H5ADFileReader(
            path=file_meta.path,
            dataset_id=file_meta.dataset_id,
            gene_name_to_id=self.gene_name_to_id,
            meta_maps=self.meta_maps,
            meta_keys=self.meta_keys,
            use_raw=self.use_raw,
            cp_target_sum=self.cp_target_sum,
            backed=self.backed,
            prefer_csr=self.prefer_csr,
            strict_gene=self.strict_gene,
            filter_missing_genes=self.filter_missing_genes,
            value_bin_cache_dir=self.value_bin_cache_dir,
            require_value_bin_cache=self.require_value_bin_cache,
        )

        self._cache[file_id] = reader

        while len(self._cache) > self.max_open_files:
            _, old_reader = self._cache.popitem(last=False)
            try:
                old_reader.close()
            except Exception:
                pass

        return reader

    def close_all(self) -> None:
        for _, reader in self._cache.items():
            try:
                reader.close()
            except Exception:
                pass
        self._cache.clear()


# =============================================================================
# Multi-file single-cell stream
# =============================================================================
class MultiH5ADCellStream(IterableDataset):
    """
    Globally shuffled single-cell stream.

    Changes from the old implementation:
    - Build a static global cell-slot index once.
    - Sample slots by a global permutation, so cells are drawn with weight
      proportional to file size (number of cells), not by equal-file coin flips.
    - Still materializes one cell per yielded example.
    """

    def __init__(
        self,
        manifest: List[Dict[str, Any]],
        *,
        gene_name_to_id: Dict[str, int],
        meta_maps: Dict[str, Dict[str, int]],
        meta_keys: Optional[List[str]] = None,
        buffer_size: int = 16,
        seed: int = 42,
        use_raw: bool = True,
        cp_target_sum: float = 1e4,
        backed: Optional[str] = "r",
        prefer_csr: bool = True,
        sort_cell_ids_for_io: bool = False,
        include_dataset_id: bool = False,
        skip_bad_files: bool = False,
        filter_missing_genes: bool = False,
        value_bin_cache_dir: str = "",
        require_value_bin_cache: bool = False,
    ) -> None:
        self.manifest = list(manifest)
        self.gene_name_to_id = gene_name_to_id
        self.meta_maps = meta_maps

        self.meta_keys = list(meta_keys or [])
        self.full_meta_keys = list(self.meta_keys)

        # Kept for compatibility with existing call sites / configs.
        self.buffer_size = int(buffer_size)
        self.seed = int(seed)

        # Equivalent to sampling one-cell pseudo-batches.
        self.N = 1

        self.use_raw = bool(use_raw)
        self.cp_target_sum = float(cp_target_sum)
        self.backed = backed
        self.prefer_csr = bool(prefer_csr)
        self.filter_missing_genes = bool(filter_missing_genes)
        self.strict_gene = not self.filter_missing_genes

        self.sort_cell_ids_for_io = bool(sort_cell_ids_for_io)
        self.include_dataset_id = bool(include_dataset_id)
        self.skip_bad_files = bool(skip_bad_files)
        self.value_bin_cache_dir = str(value_bin_cache_dir or "")
        self.require_value_bin_cache = bool(require_value_bin_cache)

        # reader cache
        self.max_open_files = 1024

        self.cycle = True
        self._iter_id = 0

        if len(self.manifest) == 0:
            raise RuntimeError("Empty manifest.")

        # Built once
        self.file_index: List[_FileIndex] = []
        # shape: (num_global_pseudo_batches, 3), columns = [file_id, group_id, batch_id]
        self.batch_refs: np.ndarray = np.empty((0, 3), dtype=np.int32)

        self._build_global_index()

    # -------------------------------------------------------------------------
    # Global index construction
    # -------------------------------------------------------------------------
    def _scan_one_file(self, item: Dict[str, Any], file_id: int) -> Optional[_FileIndex]:
        path = item["path"]
        dataset_id = int(item.get("dataset_id", 0))

        adata = None
        try:
            adata = ad.read_h5ad(path, backed=self.backed)
            n_cells = int(adata.n_obs)
            if n_cells < self.N:
                _log(f"skip file (n_cells < N): {path}")
                return None

            # Keep failure mode aligned with old stream: invalid/missing columns fail at scan
            # rather than mid-iteration.
            for k in self.full_meta_keys:
                if k not in adata.obs.columns:
                    raise ValueError(f"[{path}] obs has no required column '{k}'.")

            raw_cells = np.arange(n_cells, dtype=np.int32)
            n_batches = int(raw_cells.size // self.N)
            if n_batches <= 0:
                _log(f"skip file (no full pseudo-batches): {path}")
                return None

            usable_cells = int(n_batches * self.N)
            group = _PseudoBatchGroupIndex(
                raw_cell_ids=raw_cells,
                n_batches=n_batches,
                usable_cells=usable_cells,
            )
            return _FileIndex(path=path, dataset_id=dataset_id, groups=[group])
        finally:
            if adata is not None:
                _safe_close_anndata(adata)

    def _build_global_index(self) -> None:
        file_index: List[_FileIndex] = []
        batch_refs: List[Tuple[int, int, int]] = []

        for src_file_id, item in enumerate(self.manifest):
            try:
                fidx = self._scan_one_file(item, src_file_id)
            except Exception as e:
                if self.skip_bad_files:
                    _log(f"skip bad file: {item.get('path')} | {e}")
                    continue
                raise

            if fidx is None:
                continue

            new_file_id = len(file_index)
            file_index.append(fidx)

            for group_id, g in enumerate(fidx.groups):
                for batch_id in range(int(g.n_batches)):
                    batch_refs.append((new_file_id, group_id, batch_id))

        if len(file_index) == 0 or len(batch_refs) == 0:
            raise RuntimeError(
                "No eligible cells found. "
                "Check manifest, requested meta keys, and skip_bad_files setting."
            )

        self.file_index = file_index
        self.batch_refs = np.asarray(batch_refs, dtype=np.int32)

        total_files = len(self.file_index)
        total_groups = int(sum(len(f.groups) for f in self.file_index))
        total_batches = int(self.batch_refs.shape[0])
        total_cells_usable = int(sum(g.usable_cells for f in self.file_index for g in f.groups))

        _log(
            f"global index built: files={total_files}, groups={total_groups}, "
            f"slots={total_batches}, usable_cells={total_cells_usable}"
        )

    # -------------------------------------------------------------------------
    # Introspection helpers
    # -------------------------------------------------------------------------
    @property
    def num_global_pseudo_batches(self) -> int:
        return int(self.batch_refs.shape[0])

    def estimate_batches_per_shard(self, num_shards: int) -> int:
        if num_shards <= 0:
            return self.num_global_pseudo_batches
        return (self.num_global_pseudo_batches + num_shards - 1) // num_shards

    # -------------------------------------------------------------------------
    # Per-epoch group reshuffle
    # -------------------------------------------------------------------------
    def _build_epoch_group_batches(self, epoch_seed: int) -> List[List[np.ndarray]]:
        """
        For the current epoch/pass, rebuild each file/group by:
          - reshuffling cell ids
          - dropping tail
          - reshaping to (n_batches, N)

        Returns:
          epoch_group_batches[file_id][group_id] = np.ndarray of shape (n_batches, N)
        """
        epoch_group_batches: List[List[np.ndarray]] = []

        for file_id, fidx in enumerate(self.file_index):
            file_batches: List[np.ndarray] = []

            for group_id, gidx in enumerate(fidx.groups):
                group_seed = (
                    int(epoch_seed)
                    + 1000003 * int(file_id + 1)
                    + 9176 * int(group_id + 1)
                ) % (2**63 - 1)

                rng = np.random.default_rng(group_seed)

                raw = gidx.raw_cell_ids
                perm = rng.permutation(raw.size)
                shuffled = raw[perm]

                usable = shuffled[: gidx.usable_cells]
                batch_cells = usable.reshape(gidx.n_batches, self.N).astype(np.int32, copy=False)
                file_batches.append(batch_cells)

            epoch_group_batches.append(file_batches)

        return epoch_group_batches

    # -------------------------------------------------------------------------
    # Materialization
    # -------------------------------------------------------------------------
    def _materialize_one_batch(
        self,
        *,
        file_id: int,
        group_id: int,
        batch_id: int,
        epoch_group_batches: List[List[np.ndarray]],
        cache: _ReaderCache,
    ) -> Dict[str, Any]:
        fmeta = self.file_index[file_id]
        cell_ids = epoch_group_batches[file_id][group_id][batch_id]

        if self.sort_cell_ids_for_io:
            cell_ids = np.sort(cell_ids)

        reader = cache.get(file_id, fmeta)

        meta_acc = reader.get_meta_batch(
            cell_ids,
            include_dataset_id=self.include_dataset_id,
        )

        cid = int(cell_ids[0])
        gene_ids, expr = reader.get_cell_tokens(cid)
        total_counts, n_genes_by_counts = reader.get_cell_observation_stats(cid)
        out = {
            "gene_ids": gene_ids,
            "expr": expr,
            "meta": meta_acc,
            "total_counts": float(total_counts),
            "n_genes_by_counts": int(n_genes_by_counts),
        }
        cached_bin_ids = reader.get_cached_bin_ids(cid)
        if cached_bin_ids is not None:
            if int(cached_bin_ids.shape[0]) != len(gene_ids):
                raise ValueError(
                    f"Cached bin length mismatch for {reader.path} cell={cid}: "
                    f"bin_ids={cached_bin_ids.shape[0]} genes={len(gene_ids)}"
                )
            out["bin_ids"] = cached_bin_ids.astype(np.int64, copy=False).tolist()
        return out

    # -------------------------------------------------------------------------
    # Iteration
    # -------------------------------------------------------------------------
    def __iter__(self):
        # -------------------------
        # rank / world
        # -------------------------
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            rank = torch.distributed.get_rank()
            world_size = torch.distributed.get_world_size()
        else:
            rank = int(os.environ.get("RANK", "0"))
            world_size = int(os.environ.get("WORLD_SIZE", "1"))

        # -------------------------
        # worker info
        # -------------------------
        info = get_worker_info()
        if info is None:
            worker_id, num_workers = 0, 1
        else:
            worker_id, num_workers = info.id, info.num_workers

        shard_id = rank * num_workers + worker_id
        num_shards = world_size * num_workers

        iter_id = getattr(self, "_iter_id", 0)
        self._iter_id = iter_id + 1

        base_seed = int(self.seed)
        mix = (base_seed + 10000019 * iter_id) % (2**63 - 1)
        rng = np.random.default_rng(mix)

        if DEBUG_PRINT:
            _log(
                f"rank={rank}/{world_size} worker={worker_id}/{num_workers} "
                f"iter_id={iter_id} seed={mix} shard={shard_id}/{num_shards} "
                f"global_pseudo_batches={self.num_global_pseudo_batches}"
            )

        cache = _ReaderCache(
            max_open_files=self.max_open_files,
            gene_name_to_id=self.gene_name_to_id,
            meta_maps=self.meta_maps,
            meta_keys=self.full_meta_keys,
            use_raw=self.use_raw,
            cp_target_sum=self.cp_target_sum,
            backed=self.backed,
            prefer_csr=self.prefer_csr,
            strict_gene=self.strict_gene,
            filter_missing_genes=self.filter_missing_genes,
            value_bin_cache_dir=self.value_bin_cache_dir,
            require_value_bin_cache=self.require_value_bin_cache,
        )

        try:
            epoch_idx = 0
            while True:
                epoch_seed = (mix + 1000003 * epoch_idx) % (2**63 - 1)

                # 1) Re-shuffle every group for this epoch
                epoch_group_batches = self._build_epoch_group_batches(epoch_seed)

                # 2) Global shuffle over static slot table
                global_order = rng.permutation(self.num_global_pseudo_batches)

                # 3) Shard by rank * worker
                local_order = global_order[shard_id::num_shards]

                if DEBUG_PRINT:
                    _log(
                        f"new epoch/pass: epoch_idx={epoch_idx}, total={self.num_global_pseudo_batches}, "
                        f"local={len(local_order)}, shard={shard_id}/{num_shards}"
                    )

                for ref_idx in local_order.tolist():
                    file_id, group_id, batch_id = self.batch_refs[int(ref_idx)].tolist()
                    yield self._materialize_one_batch(
                        file_id=int(file_id),
                        group_id=int(group_id),
                        batch_id=int(batch_id),
                        epoch_group_batches=epoch_group_batches,
                        cache=cache,
                    )

                epoch_idx += 1
                if not self.cycle:
                    break

        finally:
            cache.close_all()