PDGrapher/src/gidflow/data/sciplex_dataset.py

"""SciPlex3 drug perturbation dataset for population-level training.

Loads the Srivatsan & Trapnell 2020 SciPlex3 dataset, which contains
single-cell RNA-seq profiles of cells treated with small molecule drugs
at multiple doses and time points.

Data structure
--------------
    SciPlex3 contains:
      - 189 drug perturbations (including vehicle control)
      - 5 dose levels: 0, 10, 100, 1000, 10000 nM
      - 2 time points: 24h, 72h
      - 3 cell lines: MCF7 (breast), A549 (lung), K562 (CML)
      - ChEMBL IDs for 155 drugs
      - Protein target annotations (87 target categories)

    Training pairs are constructed as:
      source = vehicle control cells at the same dose/time/cell_line
      target = drug-treated cells

    This ensures the source/target difference is purely the drug effect,
    not batch or technical variation.

Parameters
----------
h5ad_path : str
    Path to SciPlex3 h5ad file.
n_hvg : int
    Number of highly variable genes to select (default 2000).
min_cells_per_cond : int
    Minimum cells per (drug, dose, time) condition (default 30).
max_source_cells : int
    Cells sampled per source (vehicle) population (default 64).
max_target_cells : int
    Cells sampled per target (drug-treated) population (default 64).
cell_lines : list of str, optional
    Which cell lines to include. Default: all (MCF7, A549, K562).
doses : list of float, optional
    Which doses to include. Default: all non-zero doses.
times : list of float, optional
    Which time points to include. Default: all.
include_control : bool
    If True, include vehicle-only conditions as training samples (no drug).
    Default: False.
target_sum : float
    Normalisation target (default 1e4).
seed : int
    Random seed.

Usage
-----
    ds = Sciplex3Dataset("SrivatsanTrapnell2020_sciplex3.h5ad")
    sample = ds[0]
    # source_cells: [Ns, G]  vehicle control
    # target_cells: [Nt, G]  drug-treated
    # drug_emb:     [D_d]    drug molecular embedding
    # dose:         scalar    drug concentration
    # perturbation: [G]       drug target genes (multi-hot)
    # chembl_id:    str       ChEMBL identifier
    # target:       str       protein target name
    # pathway:      str       pathway annotation
"""

from __future__ import annotations

import logging
import os
import csv
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

logger = logging.getLogger(__name__)


class Sciplex3Dataset(Dataset):
    """SciPlex3 drug perturbation dataset for population-level training."""

    def __init__(
        self,
        h5ad_path: str,
        n_hvg: int = 2000,
        min_cells_per_cond: int = 30,
        max_source_cells: int = 64,
        max_target_cells: int = 64,
        cell_lines: Optional[List[str]] = None,
        doses: Optional[List[float]] = None,
        times: Optional[List[float]] = None,
        include_control: bool = False,
        target_sum: float = 1e4,
        seed: int = 42,
        # Drug encoder settings
        drug_encoding: str = "morgan",
        drug_emb_dim: int = 128,
        # Pre-computed drug embeddings (optional, avoids importing DrugEncoder)
        precomputed_drug_embeddings: Optional[Dict[str, torch.Tensor]] = None,
        # Drug-target mapping: external gene targets per drug
        drug_target_mapping: Optional[Dict[str, List[str]]] = None,
        # Preprocessed data path (faster than loading from h5ad)
        preprocessed_path: Optional[str] = None,
        # Align gene dimension to this value (pad with zeros if needed)
        target_num_genes: Optional[int] = None,
        # Path to ChEMBL SMILES CSV (chembl_id,smiles)
        drug_smiles_csv: str = "",
    ) -> None:
        super().__init__()
        self.max_source_cells = max_source_cells
        self.max_target_cells = max_target_cells
        self.seed = seed
        self.rng = np.random.default_rng(seed)
        self.target_num_genes = target_num_genes
        self._smiles_cache: Dict[str, str] = {}
        self._load_smiles_cache(drug_smiles_csv)

        # Load from preprocessed file (fast path) or h5ad (slow path)
        if preprocessed_path and os.path.exists(preprocessed_path):
            self._load_preprocessed(preprocessed_path, drug_emb_dim, precomputed_drug_embeddings)
            # Try to load gene symbol mapping (ENSEMBL → symbol)
            map_path = preprocessed_path.replace('.pt', '_gene_map.json')
            if os.path.exists(map_path):
                self._load_gene_mapping(map_path)
            return

        print(f"[Sciplex3Dataset] Loading {os.path.basename(h5ad_path)} ...")
        try:
            import anndata as ad
        except ImportError:
            raise ImportError("pip install anndata")

        adata = ad.read_h5ad(h5ad_path)
        print(f"  raw shape: {adata.shape}")

        # ------------------------------------------------------------------
        # 1. Filter by metadata BEFORE extracting expression matrix
        #    SciPlex3 has 799K cells × 110K genes — must filter before dense load
        # ------------------------------------------------------------------
        obs = adata.obs.copy()

        # Filter by cell line
        if cell_lines:
            obs = obs[obs['cell_line'].isin(cell_lines)].copy()
            print(f"  Filtered to cell lines: {cell_lines} → {len(obs)} cells")

        # Filter by dose (exclude dose=0 = vehicle)
        if doses:
            obs = obs[obs['dose_value'].isin(doses)].copy()
        else:
            obs = obs[obs['dose_value'] > 0].copy()

        if times:
            obs = obs[obs['time'].isin(times)].copy()

        # Subset adata to filtered cells
        keep_idx = obs.index.values
        print(f"  After filtering: {len(keep_idx)} cells")
        print(f"  Cell lines: {obs['cell_line'].value_counts().to_dict()}")
        print(f"  Doses: {sorted(obs['dose_value'].unique())}")
        print(f"  Drugs: {obs['perturbation'].nunique()}")

        adata = adata[keep_idx].copy()
        obs = adata.obs.copy()
        obs['_cell_idx'] = np.arange(len(obs))
        # NOTE on drug-name whitespace: 11 SciPlex3 perturbations carry a
        # trailing space in the h5ad ('Busulfan ', 'Mesna ', ...). We deliberately
        # KEEP the raw name in cond['drug_name'] because the split pair_ids were
        # generated from the raw (spaced) names, so pair_id<->split matching needs
        # the raw form. Stripping is applied ONLY at the drug_order / protein-index
        # lookup boundaries (see _get_batch_protein_labels and eval query builder).

        # ------------------------------------------------------------------
        # 2. Normalize + HVG selection (on filtered data)
        #    Work with sparse matrices to avoid 300+ GiB dense allocation.
        # ------------------------------------------------------------------
        import scipy.sparse as sp

        # Capture gene names before deleting adata
        var_names_all = list(adata.var.index) if hasattr(adata, 'var') else [f"Gene_{i}" for i in range(adata.shape[1])]

        X_sp = adata.X  # sparse: [n_cells, n_genes]
        if not sp.issparse(X_sp):
            X_sp = sp.csr_matrix(X_sp)
        elif not sp.isspmatrix_csr(X_sp):
            X_sp = X_sp.tocsr()

        # Free the anndata object (keep only the sparse matrix)
        del adata

        # --- Sparse library-size normalization (no full densification) ---
        row_sums = np.asarray(X_sp.sum(axis=1)).ravel().clip(min=1.0)  # [n_cells]
        # Scale factor per cell; apply via row-wise multiplication in sparse form
        scale = (target_sum / row_sums).astype(np.float32)  # [n_cells]
        X_sp_norm = X_sp.multiply(scale[:, np.newaxis])  # still sparse [n_cells, n_genes]
        del X_sp  # free raw sparse

        if n_hvg > 0:
            # Compute per-gene variance in sparse normalized space without densifying
            # E[x^2] - E[x]^2; works for nonnegative (pre-log) counts
            n_cells = X_sp_norm.shape[0]
            mean_g = np.asarray(X_sp_norm.mean(axis=0)).ravel()  # [n_genes]
            # E[x^2]: sum of squares per column
            X_sq = X_sp_norm.copy()
            X_sq.data **= 2
            mean_sq_g = np.asarray(X_sq.mean(axis=0)).ravel()  # [n_genes]
            del X_sq
            var_g = mean_sq_g - mean_g ** 2  # [n_genes]
            n_select = min(n_hvg, X_sp_norm.shape[1])
            hvg_idx = np.argsort(var_g)[::-1][:n_select].copy()
            gene_names = np.array(var_names_all)[hvg_idx]

            # Densify only the HVG columns: [n_cells, n_hvg]
            # Slice column-wise in CSC for efficiency
            X_sp_csc = X_sp_norm.tocsc()
            del X_sp_norm
            X_hvg_pre = np.asarray(X_sp_csc[:, hvg_idx].todense(), dtype=np.float32)
            del X_sp_csc
        else:
            # Densify full matrix (only safe for small datasets)
            gene_names = np.array(var_names_all)
            X_hvg_pre = np.asarray(X_sp_norm.todense(), dtype=np.float32)
            del X_sp_norm

        # Apply log1p after slimming (much cheaper on [n_cells, n_hvg])
        X_hvg = np.log1p(X_hvg_pre)
        del X_hvg_pre

        self.num_genes = len(gene_names)
        self.gene_names = gene_names
        self._gene_to_idx: Dict[str, int] = {g: i for i, g in enumerate(gene_names)}
        print(f"  Using {self.num_genes} genes")

        # Align gene dimension to target_num_genes if specified
        if target_num_genes is not None and self.num_genes != target_num_genes:
            self._align_genes(target_num_genes)

        # ------------------------------------------------------------------
        # 3. Separate vehicle and drug-treated cells
        # ------------------------------------------------------------------
        vehicle_mask = obs['perturbation'].str.lower().str.contains('vehicle', na=False) | \
                       (obs['dose_value'] == 0)

        # If include_control=True, also use dose=0 as target (reconstruction task)
        vehicle_obs = obs[vehicle_mask].copy()
        drug_obs = obs[~vehicle_mask].copy()

        # If no vehicle cells found, use all cells as drug-treated and
        # create pseudo-vehicle from the same condition's minimum dose
        if len(vehicle_obs) == 0:
            logger.warning("No vehicle cells found. Using minimum dose as pseudo-vehicle.")
            min_dose = obs['dose_value'].min()
            vehicle_obs = obs[obs['dose_value'] == min_dose].copy()
            drug_obs = obs[obs['dose_value'] > min_dose].copy()

        print(f"  Vehicle cells: {len(vehicle_obs)}, Drug cells: {len(drug_obs)}")

        # ------------------------------------------------------------------
        # 4. Build drug-target mapping
        # ------------------------------------------------------------------
        # From SciPlex3 target annotations (protein-level → gene-level)
        # Plus optional external mapping (DGIdb)
        self._drug_target_map = self._build_drug_target_map(
            drug_obs=drug_obs,
            external_mapping=drug_target_mapping,
            gene_names=list(gene_names),
        )

        # ------------------------------------------------------------------
        # 5. Build training conditions
        # ------------------------------------------------------------------
        # Each condition = (drug, dose, time, cell_line) tuple
        conditions: List[Dict] = []

        # Group drug-treated cells by (perturbation, dose, time, cell_line)
        drug_groups = drug_obs.groupby(['perturbation', 'dose_value', 'time', 'cell_line'])

        for (drug_name, dose, time, cell_line), group in drug_groups:
            if len(group) < min_cells_per_cond:
                continue

            # Find matching vehicle cells: same cell_line, same dose=0, same plate if possible
            vehicle_same_cl = vehicle_obs[vehicle_obs['cell_line'] == cell_line]
            if len(vehicle_same_cl) < min_cells_per_cond:
                # Fall back to any vehicle
                vehicle_same_cl = vehicle_obs

            # Build drug target vector in HVG space
            pert_vec = self._build_perturbation_vector(drug_name, drug_target_mapping)

            # Get ChEMBL ID
            chembl_id = group['chembl-ID'].iloc[0] if 'chembl-ID' in group.columns else ""
            if pd.isna(chembl_id):
                chembl_id = ""

            # Get target annotation
            target_name = group['target'].iloc[0] if 'target' in group.columns else ""
            pathway_l1 = group['pathway_level_1'].iloc[0] if 'pathway_level_1' in group.columns else ""
            pathway_l2 = group['pathway_level_2'].iloc[0] if 'pathway_level_2' in group.columns else ""

            conditions.append({
                "pert_name": f"{drug_name}_{dose}nM",
                "drug_name": drug_name,
                "cell_line": cell_line,
                "dose": float(dose),
                "time": float(time) if not pd.isna(time) else 24.0,
                "drug_cell_idx": group['_cell_idx'].values,
                "vehicle_cell_idx": vehicle_same_cl['_cell_idx'].values,
                "pert_vec": pert_vec,
                "chembl_id": chembl_id,
                "target": target_name,
                "pathway_l1": pathway_l1,
                "pathway_l2": pathway_l2,
            })

        self._conditions = conditions
        self._X = X_hvg  # [n_cells, n_hvg]

        print(f"  Conditions (≥{min_cells_per_cond} cells): {len(conditions)}")

        # Add drug_name → SMILES entries to cache (via ChEMBL ID lookup)
        for cond in self._conditions:
            name = cond['drug_name']
            chembl = cond.get('chembl_id', '')
            if chembl and chembl in self._smiles_cache:
                self._smiles_cache[name] = self._smiles_cache[chembl]

        # Align gene dimension after conditions are built (updates pert_vecs)
        if self.target_num_genes is not None and self.num_genes != self.target_num_genes:
            self._align_genes(self.target_num_genes)

        # Drug name list for encoding
        drug_names = list(set(c['drug_name'] for c in conditions))
        print(f"  Unique drugs: {len(drug_names)}")

        # ------------------------------------------------------------------
        # 6. Pre-compute drug embeddings
        # ------------------------------------------------------------------
        self.drug_emb_dim = drug_emb_dim
        self._drug_name_to_idx = {name: i for i, name in enumerate(sorted(drug_names))}

        # SMILES lookup table (populated by _load_smiles_cache in __init__)
        # Add drug_name → SMILES entries via ChEMBL ID lookup
        for cond in conditions:
            name = cond['drug_name']
            chembl = cond.get('chembl_id', '')
            if chembl and chembl in self._smiles_cache:
                self._smiles_cache[name] = self._smiles_cache[chembl]

        # Pre-computed drug embeddings
        self._drug_embeddings: Optional[torch.Tensor] = None
        if precomputed_drug_embeddings is not None:
            self._load_precomputed_embeddings(precomputed_drug_embeddings)
        else:
            # Compute embeddings using a lightweight Morgan fingerprint encoder
            self._compute_drug_embeddings()

    # ------------------------------------------------------------------
    def __len__(self) -> int:
        return len(self._conditions)

    def __getitem__(self, idx: int) -> Dict:
        cond = self._conditions[idx]

        # Sample source (vehicle) cells
        ns = min(self.max_source_cells, len(cond['vehicle_cell_idx']))
        src_idx = self.rng.choice(cond['vehicle_cell_idx'], size=ns, replace=False)
        source_cells = torch.from_numpy(self._X[src_idx]).float()  # [ns, G]

        # Sample target (drug-treated) cells
        nt = min(self.max_target_cells, len(cond['drug_cell_idx']))
        tgt_idx = self.rng.choice(cond['drug_cell_idx'], size=nt, replace=False)
        target_cells = torch.from_numpy(self._X[tgt_idx]).float()  # [nt, G]

        perturbation = torch.from_numpy(cond['pert_vec']).float()  # [G]
        drug_smiles = self._get_drug_smiles(cond['drug_name'])  # str

        return {
            "source_cells": source_cells,
            "target_cells": target_cells,
            "perturbation": perturbation,
            "drug_smiles": drug_smiles,
            "dose": torch.tensor(cond['dose'], dtype=torch.float32),
            "cell_line": cond['cell_line'],
            "chembl_id": cond['chembl_id'],
            "target": cond['target'],
            "pathway_l1": cond['pathway_l1'],
            "pathway_l2": cond['pathway_l2'],
            "metadata": {
                "pert_name": cond['pert_name'],
                "drug_name": cond['drug_name'],
                "num_source_cells": ns,
                "num_target_cells": nt,
            },
        }

    # ------------------------------------------------------------------
    # ------------------------------------------------------------------
    def _get_drug_smiles(self, drug_name: str) -> str:
        """Get SMILES string for a drug by name."""
        # Try direct lookup
        smiles = self._smiles_cache.get(drug_name, "")
        if smiles:
            return smiles
        # Try ChEMBL ID lookup
        idx = self._drug_name_to_idx.get(drug_name, 0)
        cond = self._conditions[0] if self._conditions else {}
        # Try to find ChEMBL ID from conditions
        for c in self._conditions:
            if c['drug_name'] == drug_name:
                chembl_id = c.get('chembl_id', '')
                if chembl_id:
                    smiles = self._smiles_cache.get(chembl_id, "")
                    if smiles:
                        return smiles
                break
        return ""

    def _get_drug_embedding(self, drug_name: str) -> torch.Tensor:
        """Get pre-computed drug embedding by name (legacy, for backward compat)."""
        if self._drug_embeddings is None:
            device = torch.device("cpu")
            return torch.zeros(self.drug_emb_dim, device=device)

        idx = self._drug_name_to_idx.get(drug_name, 0)
        return self._drug_embeddings[idx]

    def _load_preprocessed(
        self,
        path: str,
        drug_emb_dim: int = 128,
        precomputed_drug_embeddings: Optional[Dict[str, torch.Tensor]] = None,
    ) -> None:
        """Load from preprocessed .pt file (much faster than h5ad)."""
        print(f"  Loading preprocessed data from {path} ...")
        data = torch.load(path, weights_only=False)

        self._X = data['X'].numpy() if isinstance(data['X'], torch.Tensor) else data['X']
        self.gene_names = data['gene_names']
        self.num_genes = len(self.gene_names)
        self._gene_to_idx = {g: i for i, g in enumerate(self.gene_names)}
        obs = data['obs']

        print(f"  Loaded: {self._X.shape[0]} cells, {self.num_genes} genes")

        # Try to load sidecar gene mapping (ENSEMBL → symbol)
        map_path = path.replace('.pt', '_gene_map.json')
        if os.path.exists(map_path):
            self._load_gene_mapping(map_path)

        # ------------------------------------------------------------------
        # Separate vehicle and drug-treated cells
        # ------------------------------------------------------------------
        vehicle_mask = obs['perturbation'].str.lower().str.contains('vehicle', na=False) | \
                       (obs['dose_value'] == 0)
        vehicle_obs = obs[vehicle_mask].copy()
        drug_obs = obs[~vehicle_mask].copy()

        if len(vehicle_obs) == 0:
            min_dose = obs['dose_value'].min()
            vehicle_obs = obs[obs['dose_value'] == min_dose].copy()
            drug_obs = obs[obs['dose_value'] > min_dose].copy()

        print(f"  Vehicle cells: {len(vehicle_obs)}, Drug cells: {len(drug_obs)}")

        # ------------------------------------------------------------------
        # Build conditions
        # ------------------------------------------------------------------
        drug_target_mapping = getattr(self, '_drug_target_mapping', None)
        # ------------------------------------------------------------------
        # Build conditions
        # ------------------------------------------------------------------
        # Pre-build drug_target_map from target categories in the data
        self._drug_target_map = {}
        for drug_name in drug_obs['perturbation'].unique():
            target_cat = drug_obs[drug_obs['perturbation'] == drug_name]['target'].iloc[0]
            gene_targets = self._target_category_to_genes(
                target_cat, list(self.gene_names),
                self._gene_to_idx, getattr(self, '_symbol_to_ensembl', None)
            )
            self._drug_target_map[drug_name] = gene_targets

        conditions: List[Dict] = []

        drug_groups = drug_obs.groupby(['perturbation', 'dose_value', 'time', 'cell_line'])
        for (drug_name, dose, time, cell_line), group in drug_groups:
            if len(group) < self.max_source_cells:
                continue

            vehicle_same_cl = vehicle_obs[vehicle_obs['cell_line'] == cell_line]
            if len(vehicle_same_cl) < self.max_source_cells:
                vehicle_same_cl = vehicle_obs

            pert_vec = self._build_perturbation_vector(drug_name, None)
            chembl_id = group['chembl-ID'].iloc[0] if 'chembl-ID' in group.columns else ""
            if pd.isna(chembl_id):
                chembl_id = ""
            target_name = group['target'].iloc[0] if 'target' in group.columns else ""
            pathway_l1 = group['pathway_level_1'].iloc[0] if 'pathway_level_1' in group.columns else ""
            pathway_l2 = group['pathway_level_2'].iloc[0] if 'pathway_level_2' in group.columns else ""

            conditions.append({
                "pert_name": f"{drug_name}_{dose}nM",
                "drug_name": drug_name,
                "cell_line": cell_line,
                "dose": float(dose),
                "time": float(time) if not pd.isna(time) else 24.0,
                "drug_cell_idx": group['_cell_idx'].values,
                "vehicle_cell_idx": vehicle_same_cl['_cell_idx'].values,
                "pert_vec": pert_vec,
                "chembl_id": chembl_id,
                "target": target_name,
                "pathway_l1": pathway_l1,
                "pathway_l2": pathway_l2,
            })

        self._conditions = conditions
        print(f"  Conditions: {len(conditions)}")

        # Add drug_name → SMILES entries to cache (via ChEMBL ID lookup)
        for cond in self._conditions:
            name = cond['drug_name']
            chembl = cond.get('chembl_id', '')
            if chembl and chembl in self._smiles_cache:
                self._smiles_cache[name] = self._smiles_cache[chembl]

        # Align gene dimension after conditions are built (updates pert_vecs)
        if self.target_num_genes is not None and self.num_genes != self.target_num_genes:
            self._align_genes(self.target_num_genes)

        drug_names = list(set(c['drug_name'] for c in conditions))
        print(f"  Unique drugs: {len(drug_names)}")

        # Drug embeddings
        self.drug_emb_dim = drug_emb_dim
        self._drug_name_to_idx = {name: i for i, name in enumerate(sorted(drug_names))}
        # _smiles_cache already populated by _load_smiles_cache() in __init__
        self._drug_embeddings: Optional[torch.Tensor] = None
        if precomputed_drug_embeddings is not None:
            self._load_precomputed_embeddings(precomputed_drug_embeddings)
        else:
            self._compute_drug_embeddings()

    def _compute_drug_embeddings(self) -> None:
        """Compute Morgan fingerprint embeddings for all unique drugs."""
        drug_names = sorted(self._drug_name_to_idx.keys())
        n_drugs = len(drug_names)
        self._drug_embeddings = torch.zeros(n_drugs, self.drug_emb_dim)

        # Build name → ChEMBL ID mapping from conditions metadata
        name_to_chembl = self._build_name_chembl_map()

        try:
            from rdkit import Chem
            from rdkit.Chem import AllChem
            # Test that RDKit actually works (may fail with NumPy 2.x)
            Chem.MolFromSmiles("C")
        except (ImportError, AttributeError, Exception) as e:
            logger.warning("RDKit not available (%s) — drug embeddings will be zero vectors.", e)
            return

        success_count = 0
        for i, name in enumerate(drug_names):
            smiles = ""

            # Strategy 1: Look up via ChEMBL ID (most reliable)
            chembl_id = name_to_chembl.get(name, "")
            if chembl_id:
                smiles = self._smiles_cache.get(chembl_id, "")

            # Strategy 2: Try drug name directly as SMILES (works for simple names like "C")
            if not smiles:
                smiles = name

            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                if chembl_id:
                    logger.warning("No SMILES for drug '%s' (ChEMBL: %s, tried: %s)", name, chembl_id, smiles[:40] if smiles else "")
                continue

            fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048)
            fp_arr = np.array(list(fp), dtype=np.float32)

            # Simple MLP projection: 2048 → 512 → emb_dim
            fp_tensor = torch.from_numpy(fp_arr).float()
            hidden = torch.nn.functional.linear(
                fp_tensor,
                torch.randn(512, 2048) * 0.01,
            )
            hidden = torch.nn.functional.gelu(hidden)
            emb = torch.nn.functional.linear(
                hidden,
                torch.randn(self.drug_emb_dim, 512) * 0.01,
            )
            self._drug_embeddings[i] = emb.detach()
            success_count += 1

        print(f"  Pre-computed drug embeddings: {self._drug_embeddings.shape} "
              f"({success_count}/{n_drugs} drugs with fingerprints)")

    def _load_precomputed_embeddings(self, embeddings: Dict[str, torch.Tensor]) -> None:
        """Load pre-computed drug embeddings from a dictionary."""
        drug_names = sorted(self._drug_name_to_idx.keys())
        n_drugs = len(drug_names)
        self._drug_embeddings = torch.zeros(n_drugs, self.drug_emb_dim)

        for i, name in enumerate(drug_names):
            if name in embeddings:
                emb = embeddings[name]
                if emb.dim() == 1:
                    emb = emb[:self.drug_emb_dim]
                elif emb.dim() == 2:
                    emb = emb[0, :self.drug_emb_dim]
                self._drug_embeddings[i] = emb

        print(f"  Loaded pre-computed drug embeddings: {self._drug_embeddings.shape}")

    # ------------------------------------------------------------------
    def set_smiles(self, smiles_dict: Dict[str, str]) -> None:
        """Update SMILES for drugs and recompute embeddings.

        Parameters
        ----------
        smiles_dict : dict mapping drug_name → SMILES string
        """
        self._smiles_cache.update(smiles_dict)
        self._compute_drug_embeddings()

    # ------------------------------------------------------------------
    def _load_gene_mapping(self, map_path: str) -> None:
        """Load ENSEMBL→gene symbol mapping from JSON file."""
        import json
        try:
            with open(map_path) as f:
                mapping = json.load(f)
            # Build reverse map: symbol → list of ENSEMBL IDs
            self._symbol_to_ensembl: Dict[str, List[str]] = {}
            for eid, sym in mapping.items():
                sym = sym.upper()
                if sym not in self._symbol_to_ensembl:
                    self._symbol_to_ensembl[sym] = []
                self._symbol_to_ensembl[sym].append(eid)
            print(f"  Loaded gene mapping: {len(mapping)} ENSEMBL IDs → {len(self._symbol_to_ensembl)} symbols")
        except Exception as e:
            print(f"  Warning: Failed to load gene mapping ({e})")
            self._symbol_to_ensembl = {}

    def _load_smiles_cache(self, csv_path: str) -> None:
        """Load SMILES from a ChEMBL CSV into _smiles_cache.

        CSV format: chembl_id,smiles
        Keys in _smiles_cache: both chembl_id and (later) drug_name.
        """
        if not csv_path or not os.path.exists(csv_path):
            return

        import csv as csv_mod
        try:
            with open(csv_path, newline="", encoding="utf-8") as f:
                reader = csv_mod.DictReader(f)
                for row in reader:
                    cid = row.get("chembl_id", "").strip()
                    smi = row.get("smiles", "").strip()
                    if cid and smi:
                        self._smiles_cache[cid] = smi
            n_loaded = sum(1 for k in self._smiles_cache if k.startswith("CHEMBL"))
            print(f"  Loaded {n_loaded} SMILES from {os.path.basename(csv_path)}")
        except Exception as e:
            print(f"  Warning: Failed to load SMILES CSV ({e})")

    def _build_name_chembl_map(self) -> Dict[str, str]:
        """Build drug_name → ChEMBL ID mapping from loaded conditions."""
        mapping: Dict[str, str] = {}
        for cond in getattr(self, '_conditions', []):
            name = cond.get('drug_name', '')
            chembl = cond.get('chembl_id', '')
            if name and chembl and not pd.isna(chembl):
                mapping[name] = str(chembl).strip()
        return mapping

    def _align_genes(self, target_num_genes: int) -> None:
        """Pad or truncate gene dimension to match target_num_genes.

        When SciPlex3 has a different number of genes than the CRISPRi
        dataset (e.g., due to different HVG selection), this method pads
        the expression matrix with zeros to match the target dimension.

        Parameters
        ----------
        target_num_genes : int
            Target number of genes (from the model / other dataset).
        """
        if self.num_genes == target_num_genes:
            return

        if self.num_genes > target_num_genes:
            # Truncate: keep first target_num_genes genes
            print(f"  Aligning genes: truncating {self.num_genes} → {target_num_genes}")
            self._X = self._X[:, :target_num_genes]
            self.gene_names = self.gene_names[:target_num_genes]
            self.num_genes = target_num_genes
            self._gene_to_idx = {g: i for i, g in enumerate(self.gene_names)}
        else:
            # Pad with zeros
            pad_width = target_num_genes - self.num_genes
            print(f"  Aligning genes: padding {self.num_genes} → {target_num_genes} (+{pad_width} zeros)")
            self._X = np.pad(self._X, ((0, 0), (0, pad_width)), mode='constant')
            self.gene_names = np.concatenate([
                self.gene_names,
                np.array([f"PAD_{i}" for i in range(pad_width)]),
            ])
            self.num_genes = target_num_genes
            self._gene_to_idx = {g: i for i, g in enumerate(self.gene_names)}

        # Update perturbation vectors in conditions (if they exist)
        if hasattr(self, '_conditions') and self._conditions:
            for cond in self._conditions:
                old_vec = cond['pert_vec']
                if len(old_vec) < target_num_genes:
                    new_vec = np.zeros(target_num_genes, dtype=np.float32)
                    new_vec[:len(old_vec)] = old_vec
                    cond['pert_vec'] = new_vec
                elif len(old_vec) > target_num_genes:
                    cond['pert_vec'] = old_vec[:target_num_genes]

    # ------------------------------------------------------------------
    def _build_drug_target_map(
        self,
        drug_obs,
        external_mapping: Optional[Dict[str, List[str]]],
        gene_names: List[str],
    ) -> Dict[str, np.ndarray]:
        """Build per-drug target gene vectors in HVG space."""
        drug_target_map: Dict[str, np.ndarray] = {}

        # From SciPlex3 target annotations
        for drug_name in drug_obs['perturbation'].unique():
            target_cat = drug_obs[drug_obs['perturbation'] == drug_name]['target'].iloc[0]
            gene_targets = self._target_category_to_genes(
                target_cat, gene_names,
                self._gene_to_idx, getattr(self, '_symbol_to_ensembl', None)
            )
            drug_target_map[drug_name] = gene_targets

        # Merge with external mapping (DGIdb)
        if external_mapping:
            for drug_name, gene_list in external_mapping.items():
                targets = np.zeros(len(gene_names), dtype=np.float32)
                for g in gene_list:
                    if g in self._gene_to_idx:
                        targets[self._gene_to_idx[g]] = 1.0
                drug_target_map[drug_name] = targets

        return drug_target_map

    def _build_perturbation_vector(
        self, drug_name: str, external_mapping: Optional[Dict]
    ) -> np.ndarray:
        """Build multi-hot perturbation vector for a drug."""
        pert_vec = np.zeros(self.num_genes, dtype=np.float32)

        # From internal mapping
        if drug_name in self._drug_target_map:
            pert_vec = self._drug_target_map[drug_name].copy()

        # From external mapping
        if external_mapping and drug_name in external_mapping:
            for g in external_mapping[drug_name]:
                if g in self._gene_to_idx:
                    pert_vec[self._gene_to_idx[g]] = 1.0

        return pert_vec

    # ------------------------------------------------------------------
    # Utility methods
    # ------------------------------------------------------------------
    @staticmethod
    def _normalize_counts(X: np.ndarray, target_sum: float = 1e4) -> np.ndarray:
        """Library-size normalization + log1p."""
        from ..data.gene_selection import normalize_counts
        return normalize_counts(X, target_sum=target_sum, log1p=True)

    @staticmethod
    def _select_hvg(X: np.ndarray, n_genes: int = 2000) -> np.ndarray:
        """Select highly variable genes."""
        from ..data.gene_selection import select_hvg_by_variance
        return select_hvg_by_variance(X, n_genes=n_genes)

    @staticmethod
    def _target_category_to_genes(
        target_category: str,
        gene_names: List[str],
        gene_to_idx: Optional[Dict[str, int]] = None,
        symbol_to_ensembl: Optional[Dict[str, List[str]]] = None,
    ) -> np.ndarray:
        """Map SciPlex3 protein target category to gene targets.

        This is a simplified mapping. For production use, replace with
        a curated mapping from ChEMBL/DGIdb.

        Parameters
        ----------
        target_category : str
            SciPlex3 target annotation, e.g. "Aurora Kinase", "HDAC", "CDK"
        gene_names : list of str
            All gene names in the HVG space.

        Returns
        -------
        gene_targets : [G] binary vector
        """
        # Curated mapping: target category → known target genes
        # This should be loaded from a proper database in production
        TARGET_GENE_MAP = {
            "Aurora Kinase": ["AURKA", "AURKB", "AURKC"],
            "EGFR": ["EGFR"],
            "HDAC": ["HDAC1", "HDAC2", "HDAC3", "HDAC4", "HDAC5", "HDAC6", "HDAC7", "HDAC8", "HDAC9", "HDAC10", "HDAC11"],
            "DNA alkylator": [],
            "Bcl-2": ["BCL2", "BCL2L1"],
            "PKC": ["PRKCA", "PRKCB", "PRKCG", "PRKCD", "PRKCE", "PRKCH", "PRKCI", "PRKCL", "PRKCM", "PRKCN", "PRKDC", "PRKD1", "PRKD2", "PRKD3"],
            "CDK": ["CDK1", "CDK2", "CDK4", "CDK6", "CDK7", "CDK8", "CDK9", "CDK12", "CDK13"],
            "Histone Methyltransferase": ["EZH2", "DOT1L", "PRMT1", "PRMT5", "SETD2", "SETDB1", "SUV39H1", "SUV39H2"],
            "DNA/RNA Synthesis": ["TOP1", "TOP2A", "TOP2B", "RRM1", "RRM2", "TYMS", "DHFR"],
            "IGF-1R": ["IGF1R"],
            "PARP": ["PARP1", "PARP2", "PARP3", "PARP4", "PARP6", "PARP8", "PARP10", "PARP11", "PARP12", "PARP14", "PARP15"],
            "Sirtuin": ["SIRT1", "SIRT2", "SIRT3", "SIRT4", "SIRT5", "SIRT6", "SIRT7"],
            "TNF-alpha": ["TNF", "TNFRSF1A", "TNFRSF2"],
            "JAK": ["JAK1", "JAK2", "JAK3", "TYK2"],
            "Topoisomerase": ["TOP1", "TOP2A", "TOP2B"],
            "CCR": ["CCR1", "CCR2", "CCR3", "CCR4", "CCR5", "CCR6", "CCR7", "CCR8", "CCR9", "CCR10"],
            "HSP (e.g. HSP90)": ["HSP90AA1", "HSP90AB1", "HSPA1A", "HSPA1B", "HSPA5", "HSPA8", "HSPB1", "HSPD1", "HSPE1"],
            "HIF": ["HIF1A", "EPAS1", "HIF3A"],
            "Microtubule Associated": ["TUBB", "TUBA1A", "TUBA1B", "TUBA1C", "TUBB2A", "TUBB2B", "TUBB3", "TUBB4B"],
            "VEGF": ["VEGFA", "VEGFB", "VEGFC", "VEGFD"],
            "NF-kB": ["NFKB1", "NFKB2", "RELA", "RELB", "REL"],
            "PI3K": ["PIK3CA", "PIK3CB", "PIK3CD", "PIK3CG", "PIK3R1", "PIK3R2", "PIK3R3"],
            "AKT": ["AKT1", "AKT2", "AKT3"],
            "mTOR": ["MTOR"],
            "RAR": ["RARA", "RARB", "RARG"],
            "RXR": ["RXRA", "RXRB", "RXRG"],
            "BET": ["BRD2", "BRD3", "BRD4", "BRDT"],
            "IAP": ["BIRC2", "BIRC3", "BIRC5", "BIRC6", "BIRC7", "XIAP"],
            "Telomerase": ["TERT", "TERC"],
            "Mdm2": ["MDM2", "MDM4"],
            "TRAIL": ["TNFSF10", "TRAILR1", "TRAILR2", "TRAILR3", "TRAILR4"],
            "TGF-beta/Smad": ["TGFB1", "TGFB2", "TGFB3", "SMAD1", "SMAD2", "SMAD3", "SMAD4", "SMAD5", "SMAD6", "SMAD7"],
            "Wnt": ["CTNNB1", "APC", "AXIN1", "AXIN2", "GSK3B", "LRP5", "LRP6", "FZD1", "FZD2", "FZD3", "FZD4", "FZD5", "FZD6", "FZD7", "FZD8", "FZD9", "FZD10"],
            "Notch": ["NOTCH1", "NOTCH2", "NOTCH3", "NOTCH4", "JAG1", "JAG2", "DLL1", "DLL3", "DLL4"],
            "Hedgehog": ["SHH", "IHH", "DHH", "PTCH1", "PTCH2", "SMO", "GLI1", "GLI2", "GLI3"],
            "p53": ["TP53"],
            "Src": ["SRC", "FYN", "YES1", "FGR", "LCK", "HCK", "LYN", "BLK", "HCK", "MATK"],
            "BCL-2": ["BCL2", "BCL2L1", "MCL1"],
            "JAK/STAT": ["JAK1", "JAK2", "JAK3", "TYK2", "STAT1", "STAT2", "STAT3", "STAT4", "STAT5A", "STAT5B", "STAT6"],
            "PKA": ["PRKACA", "PRKACB", "PRKACG", "PRKAR1A", "PRKAR1B", "PRKAR2A", "PRKAR2B"],
            "PIP3": ["PIK3CA", "PIK3CB", "PIK3CD", "PIK3CG", "PTEN"],
            "GSK-3": ["GSK3A", "GSK3B"],
            "RAF": ["ARAF", "BRAF", "RAF1"],
            "MEK": ["MAP2K1", "MAP2K2"],
            "ERK": ["MAPK1", "MAPK3"],
            "p38": ["MAPK11", "MAPK12", "MAPK13", "MAPK14"],
            "JNK": ["MAPK8", "MAPK9", "MAPK10"],
            "Caspase": ["CASP1", "CASP2", "CASP3", "CASP4", "CASP5", "CASP6", "CASP7", "CASP8", "CASP9", "CASP10"],
            "Proteasome": ["PSMA1", "PSMA2", "PSMA3", "PSMA4", "PSMA5", "PSMA6", "PSMA7", "PSMB1", "PSMB2", "PSMB3", "PSMB4", "PSMB5", "PSMB6", "PSMB7", "PSMC1", "PSMC2", "PSMC3", "PSMC4", "PSMC5", "PSMC6", "PSMD1", "PSMD2", "PSMD3", "PSMD4", "PSMD5", "PSMD6", "PSMD7", "PSMD8", "PSMD9", "PSMD10", "PSMD11", "PSMD12", "PSMD13", "PSMD14"],
            "Calcineurin": ["PPP3CA", "PPP3CB", "PPP3CC", "PPP3R1", "PPP3R2"],
            "Carbonic Anhydrase": ["CA1", "CA2", "CA3", "CA4", "CA5A", "CA6", "CA7", "CA8", "CA9", "CA10", "CA11", "CA12", "CA13", "CA14"],
            "Glucocorticoid Receptor": ["NR3C1"],
            "Estrogen Receptor": ["ESR1", "ESR2"],
            "Androgen Receptor": ["AR"],
            "Progesterone Receptor": ["PGR"],
            "Vitamin D Receptor": ["VDR"],
            "Retinoic acid receptor": ["RARA", "RARB", "RARG"],
            "FXR": ["NR1H4"],
            "LXR": ["NR1H2", "NR1H3"],
            "PPAR": ["PPARA", "PPARD", "PPARG"],
            "Serotonin Receptor": ["HTR1A", "HTR1B", "HTR1D", "HTR1E", "HTR1F", "HTR2A", "HTR2B", "HTR2C", "HTR3A", "HTR3B", "HTR3C", "HTR3D", "HTR3E", "HTR4", "HTR5A", "HTR6", "HTR7"],
            "Dopamine Receptor": ["DRD1", "DRD2", "DRD3", "DRD4", "DRD5"],
            "Histamine Receptor": ["HRH1", "HRH2", "HRH3", "HRH4"],
            "Adrenergic Receptor": ["ADRA1A", "ADRA1B", "ADRA1D", "ADRA2A", "ADRA2B", "ADRA2C", "ADRB1", "ADRB2", "ADRB3"],
            "Sigma Receptor": ["SIGMAR1", "SIGMAR2"],
            "Cannabinoid Receptor": ["CNR1", "CNR2"],
            "Adenosine Receptor": ["ADORA1", "ADORA2A", "ADORA2B", "ADORA3"],
            "Imidazoline Receptor": ["ADRA1A"],
            "GABA-A Receptor": ["GABRA1", "GABRA2", "GABRA3", "GABRA4", "GABRA5", "GABRA6", "GABRB1", "GABRB2", "GABRB3", "GABRG1", "GABRG2", "GABRG3", "GABRD", "GABRE", "GABRP", "GABRQ", "GABRR1", "GABRR2", "GABRR3"],
            "Glutamate Receptor": ["GRIA1", "GRIA2", "GRIA3", "GRIA4", "GRID1", "GRID2", "GRIN1", "GRIN2A", "GRIN2B", "GRIN2C", "GRIN2D", "GRIN3A", "GRIN3B", "GRIK1", "GRIK2", "GRIK3", "GRIK4", "GRIK5", "GRM1", "GRM2", "GRM3", "GRM4", "GRM5", "GRM6", "GRM7", "GRM8"],
            "Muscarinic Acetylcholine Receptor": ["CHRM1", "CHRM2", "CHRM3", "CHRM4", "CHRM5"],
            "Nicotinic Acetylcholine Receptor": ["CHRNA1", "CHRNA2", "CHRNA3", "CHRNA4", "CHRNA5", "CHRNA6", "CHRNA7", "CHRNA9", "CHRNA10", "CHRNB1", "CHRNB2", "CHRNB3", "CHRNB4", "CHRND", "CHRNE", "CHRNG"],
            "Opioid Receptor": ["OPRD1", "OPRK1", "OPRL1", "OPRM1"],
            "Angiotensin Receptor": ["AGTR1", "AGTR2"],
            "Chemokine Receptor": ["CCR1", "CCR2", "CCR3", "CCR4", "CCR5", "CCR6", "CCR7", "CCR8", "CCR9", "CCR10", "CXCR1", "CXCR2", "CXCR3", "CXCR4", "CXCR5", "CXCR6", "CXCR7", "XCR1", "CX3CR1"],
            "Cytokine Receptor": ["IL1R1", "IL2RA", "IL2RB", "IL2RG", "IL3RA", "IL4R", "IL5RA", "IL6R", "IL7R", "CSF1R", "CSF2R", "CSF3R", "EPOR", "FLT3", "KIT", "MPL"],
            "Tachykinin Receptor": ["TACR1", "TACR2", "TACR3"],
            "Platelet-derived growth factor receptor": ["PDGFRA", "PDGFRB"],
            "Fibroblast growth factor receptor": ["FGFR1", "FGFR2", "FGFR3", "FGFR4"],
            "Vascular endothelial growth factor receptor": ["FLT1", "KDR", "FLT4"],
            "Folate Receptor": ["FOLR1", "FOLR2", "FOLR3"],
            "NMDA Receptor": ["GRIN1", "GRIN2A", "GRIN2B", "GRIN2C", "GRIN2D"],
            "AMPA Receptor": ["GRIA1", "GRIA2", "GRIA3", "GRIA4"],
            "Kainate Receptor": ["GRIK1", "GRIK2", "GRIK3", "GRIK4", "GRIK5"],
            "Neurotrophin Receptor": ["NTRK1", "NTRK2", "NTRK3"],
            "Trk Receptor": ["NTRK1", "NTRK2", "NTRK3"],
            "Neuropeptide Receptor": [],
            "Protein kinase": [],
            "DNA polymerase": [],
            "RNA polymerase": [],
            "Ribosome": [],
            "Tubulin": ["TUBB", "TUBA1A", "TUBA1B", "TUBA1C", "TUBB2A", "TUBB2B", "TUBB3", "TUBB4B"],
            "Farnesyltransferase": ["FNTA", "FNTB"],
            "Geranylgeranyltransferase": ["PGGT1B", "RABGGTA", "RABGGTB"],
            "Dihydrofolate reductase": ["DHFR"],
            "Bacterial cell wall": [],
            "Fungal cell wall": [],
            "Bacterial ribosome": [],
            "Fungal ribosome": [],
            "Mitochondrial ribosome": ["MRPL1", "MRPL2", "MRPL3", "MRPL4", "MRPL5", "MRPL6", "MRPL7", "MRPL8", "MRPL9", "MRPL10", "MRPL11", "MRPL12", "MRPL13", "MRPL14", "MRPL15", "MRPL16", "MRPL17", "MRPL18", "MRPL19", "MRPL20", "MRPL21", "MRPL22", "MRPL23", "MRPL24", "MRPL27", "MRPL28", "MRPL30", "MRPL32", "MRPL33", "MRPL34", "MRPL35", "MRPL36", "MRPL37", "MRPL38", "MRPL39", "MRPL40", "MRPL41", "MRPL42", "MRPL43", "MRPL44", "MRPL45", "MRPL48", "MRPL49", "MRPL50", "MRPS2", "MRPS5", "MRPS6", "MRPS7", "MRPS9", "MRPS10", "MRPS11", "MRPS12", "MRPS14", "MRPS15", "MRPS16", "MRPS17", "MRPS18A", "MRPS18B", "MRPS18C", "MRPS21", "MRPS22", "MRPS23", "MRPS24", "MRPS25", "MRPS26", "MRPS27", "MRPS28", "MRPS30", "MRPS31", "MRPS33", "MRPS34", "MRPS35"],
        }

        pert_vec = np.zeros(len(gene_names), dtype=np.float32)
        target_genes = TARGET_GENE_MAP.get(target_category, [])

        # Build index if not provided
        if gene_to_idx is None:
            gene_to_idx = {g: i for i, g in enumerate(gene_names)}

        # If gene_names are ENSEMBL IDs, try to convert target symbols
        gene_name_set = set(gene_names)
        if symbol_to_ensembl and any(str(g).startswith('ENSG') for g in gene_names[:10]):
            converted_targets = []
            for g in target_genes:
                g_upper = g.upper()
                if g_upper in symbol_to_ensembl:
                    for eid in symbol_to_ensembl[g_upper]:
                        if eid in gene_name_set:
                            converted_targets.append(eid)
            target_genes = converted_targets

        for g in target_genes:
            if g in gene_to_idx:
                pert_vec[gene_to_idx[g]] = 1.0

        return pert_vec

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------
    @property
    def condition_names(self) -> List[str]:
        return [c['pert_name'] for c in self._conditions]

    @property
    def unique_drugs(self) -> List[str]:
        return sorted(set(c['drug_name'] for c in self._conditions))

    @property
    def cell_line_list(self) -> List[str]:
        return sorted(set(c['cell_line'] for c in self._conditions))


# ---------------------------------------------------------------------------
# Collate function for SciPlex3 data
# ---------------------------------------------------------------------------
def sciplex_collate_fn(batch: List[Dict], num_genes: Optional[int] = None) -> Dict[str, torch.Tensor]:
    """Collate function for SciPlex3Dataset batches.

    Handles variable-size cell populations by padding to the max size
    in the batch. Optionally pads gene dimension to match `num_genes`
    for cross-dataset training (e.g., CRISPRi + SciPlex3).

    Parameters
    ----------
    batch : list of dicts
    num_genes : int, optional
        Target gene dimension. If provided and larger than batch genes,
        pads with zeros.
    """
    B = len(batch)
    max_ns = max(b['source_cells'].size(0) for b in batch)
    max_nt = max(b['target_cells'].size(0) for b in batch)
    G = batch[0]['source_cells'].size(1)

    device = batch[0]['source_cells'].device

    # Target gene dimension
    target_G = num_genes if num_genes is not None else G
    gene_pad = max(0, target_G - G)

    source_cells = torch.zeros(B, max_ns, target_G, device=device)
    target_cells = torch.zeros(B, max_nt, target_G, device=device)
    source_mask = torch.zeros(B, max_ns, device=device)
    target_mask = torch.zeros(B, max_nt, device=device)
    perturbation = torch.zeros(B, target_G, device=device)
    drug_smiles = [b['drug_smiles'] for b in batch]  # list of str
    dose = torch.stack([b['dose'] for b in batch])  # [B]

    for i, b in enumerate(batch):
        ns = b['source_cells'].size(0)
        nt = b['target_cells'].size(0)
        # Copy gene data (existing genes)
        source_cells[i, :ns, :G] = b['source_cells']
        target_cells[i, :nt, :G] = b['target_cells']
        # Copy perturbation vector
        pert = b['perturbation']
        perturbation[i, :pert.size(0)] = pert
        source_mask[i, :ns] = 1.0
        target_mask[i, :nt] = 1.0

    return {
        "source_cells": source_cells,
        "target_cells": target_cells,
        "source_mask": source_mask,
        "target_mask": target_mask,
        "perturbation": perturbation,
        "drug_smiles": drug_smiles,
        "dose": dose,
        "metadata": [b['metadata'] for b in batch],
    }