dataloader.py

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import scipy.sparse as sp
import anndata as ad
import scanpy as sc
import pandas as pd
from typing import Optional, Union
import warnings
from pathlib import Path

def _ensure_csr_matrix(mat):
    if sp.issparse(mat) and not sp.isspmatrix_csr(mat):
        return mat.tocsr()
    return mat

def _get_X_matrix(adata: ad.AnnData):
    
    return adata.raw.X if adata.raw is not None else adata.X

def _load_tf_symbols(tf_list: Union[str, Path]):
    path = Path(tf_list)
    if path.suffix.lower() in {".csv", ".tsv", ".tbl"}:
        sep = "," if path.suffix.lower() == ".csv" else "\t"
        table = pd.read_csv(path, sep=sep, usecols=lambda col: col == "gene_name")
        if "gene_name" in table.columns:
            return set(table["gene_name"].dropna().astype(str))

    with open(path, "r") as f:
        return {line.strip() for line in f if line.strip()}

def _compute_row_sums(adata: ad.AnnData):
    X = _get_X_matrix(adata)
    if sp.issparse(X):
        sums = np.array(X.sum(axis=1)).squeeze()
    else:
        arr = np.asarray(X, dtype=np.float32)
        sums = arr.sum(axis=1)
    return sums  # shape (n_cells,)

def _extract_row(X, cell_idx):
    
    if sp.issparse(X):
        return X[cell_idx, :].toarray().squeeze().astype(np.float32)
    else:
        arr = np.asarray(X, dtype=np.float32)
        return arr[cell_idx, :]


class MultiomeDataset(Dataset):
    def __init__(
        self,
        rna_adata: ad.AnnData,
        atac_adata: ad.AnnData,
        indices: np.ndarray,
        rna_scale: np.ndarray,
        atac_scale: np.ndarray,
        batch_ids: Optional[np.ndarray] = None,
    ):
        
        self.rna_adata = rna_adata
        self.atac_adata = atac_adata
        self.indices = np.array(indices, dtype=int)
        self.rna_scale = rna_scale
        self.atac_scale = atac_scale
        self.batch_ids = None if batch_ids is None else np.asarray(batch_ids, dtype=np.int64)

        # Cache dense copies once to avoid repeated sparse row extractions during training
        rna_X = _get_X_matrix(self.rna_adata)
        atac_X = _get_X_matrix(self.atac_adata)
        if sp.issparse(rna_X):
            self.rna_matrix = rna_X.toarray().astype(np.float32, copy=False)
        else:
            self.rna_matrix = np.asarray(rna_X, dtype=np.float32)
        if sp.issparse(atac_X):
            self.atac_matrix = atac_X.toarray().astype(np.float32, copy=False)
        else:
            self.atac_matrix = np.asarray(atac_X, dtype=np.float32)

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, idx):
        cell_idx = self.indices[idx]

        rna_raw = self.rna_matrix[cell_idx]
        atac_raw = self.atac_matrix[cell_idx]

        rna_norm = rna_raw * self.rna_scale[cell_idx]
        atac_norm = atac_raw * self.atac_scale[cell_idx]

        sample = (
            torch.from_numpy(rna_raw).float(),
            torch.from_numpy(rna_norm).float(),
            torch.from_numpy(atac_raw).float(),
            torch.from_numpy(atac_norm).float(),
        )

        if self.batch_ids is not None:
            sample = sample + (torch.tensor(self.batch_ids[cell_idx], dtype=torch.long),)

        return sample


def collate_multiome(batch):
    if len(batch[0]) == 5:
        rna_raw, rna_norm, atac_raw, atac_norm, batch_ids = zip(*batch)
        return (
            torch.stack(rna_raw, dim=0),
            torch.stack(rna_norm, dim=0),
            torch.stack(atac_raw, dim=0),
            torch.stack(atac_norm, dim=0),
            torch.stack(batch_ids, dim=0),
        )

    rna_raw, rna_norm, atac_raw, atac_norm = zip(*batch)
    return (
        torch.stack(rna_raw, dim=0),
        torch.stack(rna_norm, dim=0),
        torch.stack(atac_raw, dim=0),
        torch.stack(atac_norm, dim=0),
    )


class MultiomeDataModule:
    def __init__(
        self,
        rna: Union[str, ad.AnnData],
        atac: Union[str, ad.AnnData],
        batch_size: int = 128,
        train_frac: float = 0.5,
        val_frac: float = 0.3,
        test_frac: float = 0.2,
        seed: int = 0,
        num_hvg: Optional[int] = None,  # HVG on all data (before split)
        num_hvp: Optional[int] = None,  # variable peaks on all data (before split)
        num_workers: int = 0,
        drop_last: bool = False,
        cell_type_col: str = "cell_type",  # Column name for cell type labels
        batch_col: Optional[str] = None,  # Optional donor/batch covariate column
        merge_subtype: bool = False,  # Whether to merge subtypes into super types
        subtype_separator: str = "_",  # Separator for splitting subtypes
        split_cache_path: Optional[Union[str, Path]] = None,
        include_all_tfs: bool = False,
        tf_list: Optional[Union[str, Path]] = None,
        hvg_flavor: str = "seurat_v3",
    ):
        # Load AnnData (user can pass path or object)
        rna_path = rna if isinstance(rna, str) else None
        atac_path = atac if isinstance(atac, str) else None

        if isinstance(rna, str):
            self.rna_adata = ad.read_h5ad(rna)
        else:
            self.rna_adata = rna
        if isinstance(atac, str):
            self.atac_adata = ad.read_h5ad(atac)
        else:
            self.atac_adata = atac

        # Handle cell_type_col: copy specified column to "cell_type" if different
        if cell_type_col != "cell_type":
            print(f"Copying column '{cell_type_col}' to 'cell_type' in h5ad files...")

            # Update RNA adata
            if cell_type_col not in self.rna_adata.obs.columns:
                raise ValueError(
                    f"Column '{cell_type_col}' not found in RNA h5ad obs. "
                    f"Available columns: {list(self.rna_adata.obs.columns)}"
                )
            self.rna_adata.obs["cell_type"] = self.rna_adata.obs[cell_type_col]

            # Update ATAC adata
            if cell_type_col not in self.atac_adata.obs.columns:
                raise ValueError(
                    f"Column '{cell_type_col}' not found in ATAC h5ad obs. "
                    f"Available columns: {list(self.atac_adata.obs.columns)}"
                )
            self.atac_adata.obs["cell_type"] = self.atac_adata.obs[cell_type_col]

            # Save updated h5ad files if paths were provided
            if rna_path is not None:
                self.rna_adata.write_h5ad(rna_path)
                print(f"Updated {rna_path}")
            if atac_path is not None:
                self.atac_adata.write_h5ad(atac_path)
                print(f"Updated {atac_path}")

        # Handle subtype merging: create super_type column
        if merge_subtype:
            print(f"\n=== Merging Subtypes into Super Types ===")
            print(f"Separator: '{subtype_separator}'")

            def extract_super_type(subtype_name):
                
                parts = str(subtype_name).split(subtype_separator)
                if len(parts) > 1:
                    # Take everything before the last separator
                    return subtype_separator.join(parts[:-1])
                else:
                    # No separator found, return as-is
                    return subtype_name

            # Create super_type column in RNA adata
            self.rna_adata.obs["super_type"] = self.rna_adata.obs["cell_type"].apply(extract_super_type)

            # Create super_type column in ATAC adata
            self.atac_adata.obs["super_type"] = self.atac_adata.obs["cell_type"].apply(extract_super_type)

            # Print mapping
            mapping = pd.DataFrame({
                'subtype': self.rna_adata.obs["cell_type"],
                'super_type': self.rna_adata.obs["super_type"]
            }).drop_duplicates().sort_values('super_type')

            print("\nSubtype -> Super Type Mapping:")
            for super_type in mapping['super_type'].unique():
                subtypes = mapping[mapping['super_type'] == super_type]['subtype'].tolist()
                print(f"  {super_type}: {', '.join(subtypes)}")

            print(f"\nCreated 'super_type' column (original 'cell_type' preserved)")
            print(f"Unique subtypes: {len(self.rna_adata.obs['cell_type'].unique())}")
            print(f"Unique super types: {len(self.rna_adata.obs['super_type'].unique())}")
            print(f"Will use 'super_type' for stratification and prediction")

            # Save updated h5ad files with super_type column if paths were provided
            if rna_path is not None:
                self.rna_adata.write_h5ad(rna_path)
                print(f"Updated {rna_path} with super_type column")
            if atac_path is not None:
                self.atac_adata.write_h5ad(atac_path)
                print(f"Updated {atac_path} with super_type column")

        # Ensure sparse matrices are CSR for safe advanced indexing later.
        # AnnData Raw.X is read-only, so a non-CSR raw matrix must be converted
        # through a temporary AnnData object and assigned back via adata.raw.
        for adata in (self.rna_adata, self.atac_adata):
            if (
                adata.raw is not None
                and sp.issparse(adata.raw.X)
                and not sp.isspmatrix_csr(adata.raw.X)
            ):
                raw_counts = adata.raw.to_adata()
                raw_counts.X = raw_counts.X.tocsr()
                adata.raw = raw_counts
            if sp.issparse(adata.X):
                adata.X = _ensure_csr_matrix(adata.X)

        assert (
            self.rna_adata.n_obs == self.atac_adata.n_obs
        ), "RNA and ATAC must have same number of cells."

        self.batch_size = batch_size
        self.num_workers = num_workers
        self.drop_last = drop_last
        self.seed = seed
        self.batch_col = batch_col
        self.merge_subtype = merge_subtype  # Store for later use

        self.split_cache_path = Path(split_cache_path) if split_cache_path else None
        self.batch_label_to_idx = {}
        self.batch_idx_to_label = []
        self.batch_ids = None
        self.num_batches = 0

        if batch_col is not None:
            if batch_col not in self.rna_adata.obs.columns:
                raise ValueError(
                    f"Column '{batch_col}' not found in RNA h5ad obs. "
                    f"Available columns: {list(self.rna_adata.obs.columns)}"
                )
            if batch_col not in self.atac_adata.obs.columns:
                raise ValueError(
                    f"Column '{batch_col}' not found in ATAC h5ad obs. "
                    f"Available columns: {list(self.atac_adata.obs.columns)}"
                )

            rna_batch = self.rna_adata.obs[batch_col].astype(str).values
            atac_batch = self.atac_adata.obs[batch_col].astype(str).values
            if not np.array_equal(rna_batch, atac_batch):
                raise ValueError(f"RNA and ATAC batch labels differ for column '{batch_col}'")

            self.batch_idx_to_label = sorted(pd.unique(rna_batch).tolist())
            self.batch_label_to_idx = {
                label: idx for idx, label in enumerate(self.batch_idx_to_label)
            }
            self.batch_ids = np.array(
                [self.batch_label_to_idx[label] for label in rna_batch],
                dtype=np.int64
            )
            self.num_batches = len(self.batch_idx_to_label)
            print(f"\n=== Batch Covariate ===")
            print(f"Using '{batch_col}' with {self.num_batches} unique groups")

        n_cells = self.rna_adata.n_obs

        assert abs(train_frac + val_frac + test_frac - 1.0) < 1e-6, "train_frac + val_frac + test_frac must sum to 1"

        # Use stratified split to maintain cell type proportions
        from sklearn.model_selection import train_test_split

        # Get cell type labels for stratification
        # Use super_type if merging, otherwise use cell_type
        if merge_subtype and "super_type" in self.rna_adata.obs.columns:
            cell_type_labels = self.rna_adata.obs["super_type"].values
            stratify_col = "super_type"
        else:
            cell_type_labels = self.rna_adata.obs["cell_type"].values
            stratify_col = "cell_type"

        if self.split_cache_path and self.split_cache_path.exists():
            cached = np.load(self.split_cache_path, allow_pickle=False)
            train_idx = cached["train_idx"].astype(int)
            val_idx = cached["val_idx"].astype(int)
            test_idx = cached["test_idx"].astype(int)
            print(f"\n=== Loaded cached data splits from {self.split_cache_path} ===")
        else:
            # First split: separate test set
            train_val_idx, test_idx = train_test_split(
                np.arange(n_cells),
                test_size=test_frac,
                random_state=seed,
                stratify=cell_type_labels
            )

            # Second split: separate train and val from remaining data
            # Adjust val_frac to account for already removed test set
            val_frac_adjusted = val_frac / (train_frac + val_frac)
            train_idx, val_idx = train_test_split(
                train_val_idx,
                test_size=val_frac_adjusted,
                random_state=seed,
                stratify=cell_type_labels[train_val_idx]
            )
            if self.split_cache_path:
                self.split_cache_path.parent.mkdir(parents=True, exist_ok=True)
                np.savez(
                    self.split_cache_path,
                    train_idx=train_idx.astype(np.int64),
                    val_idx=val_idx.astype(np.int64),
                    test_idx=test_idx.astype(np.int64),
                )
                print(f"\n=== Saved data splits to {self.split_cache_path} ===")

        self.train_idx = train_idx
        self.val_idx = val_idx
        self.test_idx = test_idx

        # Print split statistics
        print(f"\n=== Data Split (Stratified by {stratify_col}) ===")
        print(f"Total cells: {n_cells}")
        print(f"Train: {len(train_idx)} ({len(train_idx)/n_cells*100:.1f}%)")
        print(f"Val: {len(val_idx)} ({len(val_idx)/n_cells*100:.1f}%)")
        print(f"Test: {len(test_idx)} ({len(test_idx)/n_cells*100:.1f}%)")

        # Print per-cell-type distribution
        print(f"\nPer-{stratify_col} distribution:")
        cell_types = np.unique(cell_type_labels)
        print(f"{'Cell Type':<25} {'Total':>8} {'Train':>8} {'Val':>8} {'Test':>8}")
        print("-" * 65)
        for ct in cell_types:
            total = np.sum(cell_type_labels == ct)
            train_ct = np.sum(cell_type_labels[train_idx] == ct)
            val_ct = np.sum(cell_type_labels[val_idx] == ct)
            test_ct = np.sum(cell_type_labels[test_idx] == ct)
            print(f"{ct:<25} {total:>8} {train_ct:>8} {val_ct:>8} {test_ct:>8}")

        # --- HVG on all data, optionally unioned with all TFs, then subset RNA AnnData ---
        if num_hvg is not None:
            counts = _get_X_matrix(self.rna_adata)
            tmp = ad.AnnData(counts.copy() if sp.issparse(counts) else np.array(counts, copy=True))
            sc.pp.normalize_total(tmp, target_sum=1e4)
            sc.pp.log1p(tmp)
            sc.pp.highly_variable_genes(tmp, n_top_genes=num_hvg, flavor=hvg_flavor, subset=False)
            hvg_mask = tmp.var["highly_variable"].values.astype(bool)
        else:
            hvg_mask = np.ones(self.rna_adata.n_vars, dtype=bool)

        if include_all_tfs:
            if tf_list is None:
                raise ValueError("--include_all_tfs requires --tf_list")
            tf_symbols = _load_tf_symbols(tf_list)
            tf_mask = self.rna_adata.var_names.isin(tf_symbols)
            n_hvg = int(hvg_mask.sum())
            n_tf = int(tf_mask.sum())
            hvg_mask = hvg_mask | tf_mask
            print(
                f"include_all_tfs: selected {int(hvg_mask.sum())} RNA genes "
                f"({n_hvg} HVGs via {hvg_flavor} union {n_tf} TFs from {tf_list})"
            )

        # Keep full originals if needed
        self.rna_adata_full = self.rna_adata
        self.rna_adata = self.rna_adata[:, hvg_mask].copy()

        # --- HVP on all data, then subset ATAC AnnData ---
        if num_hvp is not None:
            # For variance of row-normalized ATAC, compute row-sum normalized matrix
            # First compute row sums of original ATAC (before subsetting)
            atac_full_X = _get_X_matrix(self.atac_adata)
            # compute row sums for normalization
            atac_row_sums_full = _compute_row_sums(self.atac_adata)
            eps = 1e-8
            inv_row = 1.0 / (atac_row_sums_full + eps)  # shape (n_cells,)
            # Apply row scaling to get row-sum-normalized ATAC for variance computation
            if sp.issparse(atac_full_X):
                # scale rows via diagonal
                D = sp.diags(inv_row)
                atac_norm_for_var = (D @ atac_full_X).toarray()
            else:
                arr = np.asarray(atac_full_X, dtype=np.float32)
                atac_norm_for_var = (arr.T * inv_row).T  # broadcast per row

            variances = np.var(atac_norm_for_var, axis=0)
            top_idx = np.argsort(-variances)[:num_hvp]
            peak_mask = np.zeros(atac_norm_for_var.shape[1], dtype=bool)
            peak_mask[top_idx] = True
        else:
            peak_mask = np.ones(self.atac_adata.n_vars, dtype=bool)

        self.atac_adata_full = self.atac_adata
        self.atac_adata = self.atac_adata[:, peak_mask].copy()

        # --- Precompute row sums and scales (L1) after subsetting ---
        eps = 1e-8
        self.rna_row_sums = _compute_row_sums(self.rna_adata)
        self.atac_row_sums = _compute_row_sums(self.atac_adata)
        self.rna_scale = 1.0 / (self.rna_row_sums + eps)
        self.atac_scale = 1.0 / (self.atac_row_sums + eps)

        # Create datasets for train/val/test splits
        dataset_params = {
            'rna_adata': self.rna_adata,
            'atac_adata': self.atac_adata,
            'rna_scale': self.rna_scale,
            'atac_scale': self.atac_scale,
            'batch_ids': self.batch_ids,
        }

        self.dataset_train = MultiomeDataset(indices=self.train_idx, **dataset_params)
        self.dataset_val = MultiomeDataset(indices=self.val_idx, **dataset_params)
        self.dataset_test = MultiomeDataset(indices=self.test_idx, **dataset_params)

        # Create DataLoaders with shared parameters
        loader_params = {
            'batch_size': self.batch_size,
            'shuffle': False,
            'collate_fn': collate_multiome,
            'num_workers': self.num_workers,
            'drop_last': self.drop_last,
            'pin_memory': True,
        }
        if self.num_workers > 0:
            loader_params['persistent_workers'] = True
            loader_params['prefetch_factor'] = 2

        # Train loader reshuffles every epoch to decorrelate minibatches across
        # epochs; val/test stay in fixed order for reproducible metrics.
        train_loader_params = {**loader_params, 'shuffle': True}
        self.train_loader = DataLoader(self.dataset_train, **train_loader_params)
        # Fixed-order training loader for embeddings, clustering diagnostics, and
        # visualization. Never use the shuffled optimization loader when rows
        # must align with train_idx and cell labels.
        self.train_eval_loader = DataLoader(self.dataset_train, **loader_params)
        self.val_loader = DataLoader(self.dataset_val, **loader_params)
        self.test_loader = DataLoader(self.dataset_test, **loader_params)
    @property
    def var_names_rna(self):
        return np.array(self.rna_adata.var_names)

    @property
    def var_names_atac(self):
        return np.array(self.atac_adata.var_names)

    def get_split_adatas(self):
        
        train_rna_adata = self.rna_adata[self.train_idx]
        train_atac_adata = self.atac_adata[self.train_idx]
        val_rna_adata = self.rna_adata[self.val_idx]
        val_atac_adata = self.atac_adata[self.val_idx]
        test_rna_adata = self.rna_adata[self.test_idx]
        test_atac_adata = self.atac_adata[self.test_idx]

        return train_rna_adata, train_atac_adata, val_rna_adata, val_atac_adata, test_rna_adata, test_atac_adata

    def get_train_indices(self):
        return self.train_idx

    def get_test_indices(self):
        return self.test_idx

    def get_val_indices(self):
        return self.val_idx

    def encode_batch_labels(self, labels) -> np.ndarray:
        if self.batch_col is None:
            raise ValueError("No batch_col configured for this MultiomeDataModule")

        labels = np.asarray(labels, dtype=str)
        unknown = sorted(set(labels) - set(self.batch_label_to_idx))
        if unknown:
            raise ValueError(f"Unknown batch labels encountered: {unknown[:5]}")

        return np.array([self.batch_label_to_idx[label] for label in labels], dtype=np.int64)