data.py

"""
Dataset for RNA binding prediction with SAMS-VAE.

Features:
- RhoFoldApoDataset: RhoFold+ apo structures for 3D loss only (z_i=0, z_g=0)
- Coordinates normalized (sample-wise) in encoder
- Preload caching for RhoFold structures to reduce CPU usage
"""

import os
import pickle
import re
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset as BaseDataset
from typing import List, Dict, Tuple, Optional, Set


def detect_embedding_dim(emb_dict: Dict) -> int:
    for k, v in emb_dict.items():
        arr = np.array(v) if not isinstance(v, np.ndarray) else v
        if arr.ndim == 1:
            return arr.shape[0]
        elif arr.ndim == 2:
            return arr.shape[1]
    return 1280


# =============================================================================
# ION DETECTION - Based on formal charge and molecular size
# Criteria: formal_charge != 0 AND heavy_atoms <= 7
# =============================================================================

def is_ion_smiles(smiles: str) -> bool:
    """Check if a SMILES string represents an ion using RDKit.
    
    Criteria: formal_charge != 0 AND heavy_atoms <= 7
    """
    if not isinstance(smiles, str) or not smiles or smiles == 'nan':
        return False
    
    try:
        from rdkit import Chem
        mol = Chem.MolFromSmiles(smiles.strip())
        if mol is None:
            return False
        
        formal_charge = Chem.GetFormalCharge(mol)
        num_heavy = mol.GetNumHeavyAtoms()
        
        return formal_charge != 0 and num_heavy <= 7
    except:
        return False


def detect_ions_from_csv(train_csv: str, test_csv: str) -> Set[str]:
    """Detect all ions from train and test CSV files based on SMILES.
    
    Returns set of uppercase ligand names that are ions.
    Criteria: formal_charge != 0 AND heavy_atoms <= 7
    """
    try:
        from rdkit import Chem
        from rdkit import RDLogger
        RDLogger.DisableLog('rdApp.*')
    except ImportError:
        print("Warning: RDKit not available, using fallback ion detection")
        return _fallback_detect_ions(train_csv, test_csv)
    
    ions = set()
    
    for csv_path in [train_csv, test_csv]:
        if not Path(csv_path).exists():
            continue
        
        df = pd.read_csv(csv_path)
        
        # Handle both 'canonical_smiles' and 'smiles' column names
        smiles_col = 'canonical_smiles' if 'canonical_smiles' in df.columns else 'smiles'
        
        if smiles_col not in df.columns or 'ligand_name' not in df.columns:
            continue
        
        for _, row in df.iterrows():
            smiles = str(row.get(smiles_col, ''))
            name = str(row.get('ligand_name', ''))
            
            # Handle NA which appears as empty string or NaN
            if name == '' or name == 'nan' or pd.isna(row.get('ligand_name')):
                name = 'NA'
            
            mol = Chem.MolFromSmiles(smiles.strip())
            if mol is None:
                continue
            
            formal_charge = Chem.GetFormalCharge(mol)
            num_heavy = mol.GetNumHeavyAtoms()
            
            # Ion criteria: charged AND small
            if formal_charge != 0 and num_heavy <= 7:
                ions.add(name.upper().strip())
    
    return ions


def _fallback_detect_ions(train_csv: str, test_csv: str) -> Set[str]:
    """Fallback ion detection when RDKit is not available."""
    # Hardcoded list based on analysis
    return {
        '2HP', '3CO', 'ACT', 'AU3', 'BA', 'BR', 'CA', 'CAC', 'CD', 'CL',
        'CO', 'CS', 'F', 'GZ6', 'HG', 'IR', 'IR3', 'IRI', 'K', 'MG',
        'MLI', 'MMC', 'MN', 'NA', 'NCO', 'NH4', 'PO4', 'RB', 'SE4', 'SO4',
        'SR', 'TL', 'ZN'
    }


def get_ion_to_idx(train_csv: str = 'data/train_w_smiles.csv', 
                   test_csv: str = 'data/test_w_smiles.csv') -> Dict[str, int]:
    """Get ion_to_idx mapping by detecting ions from CSV files.
    
    This ensures all ions in the data are included, even if data changes.
    """
    ions = detect_ions_from_csv(train_csv, test_csv)
    all_ions = sorted(ions)
    ion_to_idx = {ion: i for i, ion in enumerate(all_ions)}
    print(f"Detected {len(ion_to_idx)} ions from data: {all_ions}")
    return ion_to_idx


def get_ligand_to_idx(train_csv: str = 'data/train_w_smiles.csv',
                      test_csv: str = 'data/test_w_smiles.csv') -> Dict[str, int]:
    """Get ligand_to_idx mapping by detecting non-ion ligands from CSV files."""
    ligands = set()
    for csv_path in [train_csv, test_csv]:
        if not Path(csv_path).exists():
            continue
        df = pd.read_csv(csv_path)
        if 'ligand_name' not in df.columns:
            continue
        for name in df['ligand_name'].dropna().unique():
            name_upper = str(name).upper().strip()
            if name_upper and not is_ion(name_upper):
                ligands.add(name_upper)
    
    all_ligands = sorted(ligands)
    ligand_to_idx = {lig: i for i, lig in enumerate(all_ligands)}
    print(f"Detected {len(ligand_to_idx)} non-ion ligands from data: {all_ligands[:10]}{'...' if len(all_ligands) > 10 else ''}")
    return ligand_to_idx


def get_ligand_to_smiles(train_csv: str = 'data/train_w_smiles.csv',
                         test_csv: str = 'data/test_w_smiles.csv') -> Dict[str, str]:
    """Get ligand_to_smiles mapping (representative SMILES for each ligand type)."""
    ligand_to_smiles = {}
    for csv_path in [train_csv, test_csv]:
        if not Path(csv_path).exists():
            continue
        df = pd.read_csv(csv_path)
        if 'ligand_name' not in df.columns or 'canonical_smiles' not in df.columns:
            continue
        for _, row in df.iterrows():
            name = row.get('ligand_name', '')
            smiles = row.get('canonical_smiles', '')
            if pd.isna(name) or pd.isna(smiles) or not name or not smiles:
                continue
            name_upper = str(name).upper().strip()
            if name_upper and not is_ion(name_upper) and name_upper not in ligand_to_smiles:
                ligand_to_smiles[name_upper] = str(smiles)
    
    print(f"Built ligand_to_smiles mapping for {len(ligand_to_smiles)} ligands")
    return ligand_to_smiles


def is_ion(ligand_name: str, known_ions: Set[str] = None) -> bool:
    """Check if ligand is an ion.
    
    Args:
        ligand_name: Name of the ligand
        known_ions: Optional set of known ion names for direct lookup.
    """
    if not isinstance(ligand_name, str) or not ligand_name:
        return False
    name = ligand_name.upper().strip()
    
    # If known_ions provided, use direct lookup
    if known_ions is not None:
        return name in known_ions
    
    # Fallback: use hardcoded list
    fallback_ions = {
        '2HP', '3CO', 'ACT', 'AU3', 'BA', 'BR', 'CA', 'CAC', 'CD', 'CL',
        'CO', 'CS', 'F', 'GZ6', 'HG', 'IR', 'IR3', 'IRI', 'K', 'MG',
        'MLI', 'MMC', 'MN', 'NA', 'NCO', 'NH4', 'PO4', 'RB', 'SE4', 'SO4',
        'SR', 'TL', 'ZN'
    }
    return name in fallback_ions


# Keep backward compatibility
def is_cation(ligand_name: str) -> bool:
    """Backward compatible alias for is_ion."""
    return is_ion(ligand_name)


def load_embeddings(cache_path: str) -> Tuple[Dict[str, np.ndarray], int]:
    if cache_path and Path(cache_path).exists():
        print(f"Loading embeddings from {cache_path}")
        with open(cache_path, 'rb') as f:
            emb_dict = pickle.load(f)
        dim = detect_embedding_dim(emb_dict)
        print(f"Loaded {len(emb_dict)} embeddings (dim={dim})")
        return emb_dict, dim
    else:
        raise FileNotFoundError(f"Embedding file not found: {cache_path}")


class RNABindingDataset(BaseDataset):
    """RNA-Ligand Binding Dataset (Holo structures)."""
    
    def __init__(
        self, 
        df: pd.DataFrame, 
        path_pdbs: str, 
        rna_emb_dict: Dict, 
        ion_to_idx: Dict[str, int],
        ligand_to_idx: Dict[str, int] = None,
        rna_dim: int = 1280, 
        all_ligands: bool = False,
        train_ligands: set = None,
        rhofold_path: str = None,
    ):
        self.path_pdbs = path_pdbs
        self.rhofold_path = rhofold_path
        self.rna_emb_dict = rna_emb_dict
        self.ion_to_idx = ion_to_idx
        self.num_ions = len(ion_to_idx)
        self.ligand_to_idx = ligand_to_idx or {}
        self.num_ligands = max(len(self.ligand_to_idx), 1)
        self.rna_dim = rna_dim 
        self.all_ligands = all_ligands
        self.target_atoms = ["P", "C4'", "C1'"]
        self.train_ligands = train_ligands or set()
        
        self.samples = []
        self._build_samples(df)
    
    def _get_ligand_labels(self, row, seq_len: int) -> np.ndarray:
        """Get binding labels for a single ligand."""
        labels_str = str(row.get('binding_labels', row.get('labels', '')))
        if not labels_str or labels_str == 'nan':
            return np.zeros(seq_len, dtype=np.int8)
        labels = np.array([int(c) for c in labels_str if c in '01'], dtype=np.int8)
        if len(labels) != seq_len:
            return np.zeros(seq_len, dtype=np.int8)
        return labels
    
    def _compute_combined_labels(self, group: pd.DataFrame, seq_len: int) -> np.ndarray:
        """Compute combined labels (OR of all ligands) for the complex."""
        combined = np.zeros(seq_len, dtype=np.int8)
        for _, row in group.iterrows():
            labels = self._get_ligand_labels(row, seq_len)
            combined = np.maximum(combined, labels)
        return combined
    
    def _build_samples(self, df: pd.DataFrame):
        """
        Build samples with BOTH:
        - combined_labels: OR of all ligand binding sites (for encoder + main decoder loss)
        - per_ligand_labels: dict of {ligand_name: labels} (for auxiliary decoder loss)
        """
        for (pdb_id, chain_id), group in df.groupby(['pdb_id', 'chain_id']):
            rna_seq = group['rna_seq'].iloc[0]
            seq_len = len(rna_seq)
            
            # Combined labels for the whole complex
            combined_labels = self._compute_combined_labels(group, seq_len)
            has_binding = combined_labels.sum() > 0
            
            if not self.all_ligands and not has_binding:
                continue
            
            # Collect per-ligand labels and info
            ligand_names = []
            smiles_list = []
            per_ligand_labels = {}
            
            for _, row in group.iterrows():
                ligand_name = row.get('ligand_name', '')
                if not ligand_name or pd.isna(ligand_name):
                    continue
                
                # Filter out rows with labels sum=0
                labels = self._get_ligand_labels(row, seq_len)
                if labels.sum() == 0:
                    continue
                
                ligand_names.append(ligand_name)
                smiles = row.get('canonical_smiles', '')
                if pd.notna(smiles):
                    smiles_list.append(smiles)
                
                # Store per-ligand labels
                per_ligand_labels[ligand_name] = labels
            
            if not ligand_names:
                continue
            
            # Separate ion_counts and design_vector (both are counts)
            ion_counts = np.zeros(self.num_ions, dtype=np.float32)
            design_vector = np.zeros(self.num_ligands, dtype=np.float32)
            
            # Track unseen ligands (not in ligand_to_idx)
            unseen_ligands = []  # List of (smiles, count)
            unseen_smiles_counts = {}
            
            for lig_name in ligand_names:
                lig_upper = lig_name.upper()
                if is_ion(lig_name):
                    if lig_upper in self.ion_to_idx:
                        ion_counts[self.ion_to_idx[lig_upper]] += 1.0
                else:
                    if lig_upper in self.ligand_to_idx:
                        design_vector[self.ligand_to_idx[lig_upper]] += 1.0
                    else:
                        # Unseen ligand - track SMILES and count
                        smiles = per_ligand_labels.get(lig_name, {})
                        # Get SMILES from the row
                        for _, row in group.iterrows():
                            if row.get('ligand_name', '').upper() == lig_upper:
                                smiles = row.get('canonical_smiles', '')
                                if pd.notna(smiles) and smiles:
                                    if smiles not in unseen_smiles_counts:
                                        unseen_smiles_counts[smiles] = 0.0
                                    unseen_smiles_counts[smiles] += 1.0
                                break
            
            # Convert to list format
            for smiles, count in unseen_smiles_counts.items():
                unseen_ligands.append({'smiles': smiles, 'count': count})
            
            self.samples.append({
                'pdb_id': pdb_id, 
                'chain_id': chain_id, 
                'rna_seq': rna_seq,
                'labels': combined_labels,
                'per_ligand_labels': per_ligand_labels,
                'ion_counts': ion_counts,
                'design_vector': design_vector,
                'ligand_names': ligand_names,
                'smiles_list': smiles_list,
                'unseen_ligands': unseen_ligands,  # NEW: list of {'smiles': str, 'count': float}
            })
    
    def __len__(self):
        return len(self.samples)
    
    def get_class_distribution(self) -> Dict[str, float]:
        total_residues = sum(len(s['labels']) for s in self.samples)
        total_binding = sum((s['labels'] == 1).sum() for s in self.samples)
        return {
            'total_samples': len(self.samples),
            'total_residues': total_residues,
            'total_binding': total_binding,
            'binding_ratio': total_binding / max(total_residues, 1),
            'imbalance_ratio': (total_residues - total_binding) / max(total_binding, 1),
        }
    
    def __getitem__(self, idx):
        s = self.samples[idx]
        seq = s['rna_seq'].upper()
        seq_len = len(seq)
        
        coords = self._load_single_coords(s['pdb_id'], s['chain_id'], seq_len)
        coords_target = coords.clone()
        
        rna_emb = self._get_rna_emb(seq, seq_len)
        
        # Per-ligand labels for auxiliary loss
        per_ligand_labels = self._get_per_ligand_labels(
            s['per_ligand_labels'], s['ligand_names']
        )
        
        return {
            'pdb_id': s['pdb_id'], 
            'chain_id': s['chain_id'], 
            'rna_seq': s['rna_seq'],
            'coords': coords.float(),
            'coords_target': coords_target.float(),
            'labels': torch.tensor(s['labels'], dtype=torch.float32),
            'per_ligand_labels': per_ligand_labels,
            'ligand_names': s['ligand_names'],
            'smiles_list': s['smiles_list'],
            'ion_counts': torch.tensor(s['ion_counts']),
            'design_vector': torch.tensor(s['design_vector']),
            'rna_emb': rna_emb,
            'is_apo': False,
        }
    
    def _load_single_coords(self, pdb_id: str, chain_id: str, seq_len: int) -> torch.Tensor:
        pdb_path = os.path.join(self.path_pdbs, f"{pdb_id.upper()}_chain_{chain_id.upper()}.pdb")
        coords = []
        if os.path.exists(pdb_path):
            with open(pdb_path) as f:
                for line in f:
                    if line.startswith('ATOM') and line[12:16].strip() == "C4'":
                        coords.append([float(line[30:38]), float(line[38:46]), float(line[46:54])])
        
        if not coords:
            return torch.randn(seq_len, 3) * 10
        
        coords = torch.tensor(coords, dtype=torch.float32)
        if coords.shape[0] < seq_len:
            coords = torch.cat([coords, torch.zeros(seq_len - coords.shape[0], 3)])
        return coords[:seq_len]
    
    def _get_rna_emb(self, seq: str, seq_len: int) -> torch.Tensor:
        if seq in self.rna_emb_dict:
            emb = self.rna_emb_dict[seq]
            if not isinstance(emb, torch.Tensor):
                emb = torch.tensor(emb, dtype=torch.float32)
            if emb.shape[0] < seq_len:
                emb = torch.cat([emb, torch.zeros(seq_len - emb.shape[0], self.rna_dim)])
            return emb[:seq_len]
        return torch.zeros(seq_len, self.rna_dim)
    
    def _get_per_ligand_labels(
        self, 
        per_ligand_labels: Dict[str, np.ndarray],
        ligand_names: List[str],
    ) -> List[torch.Tensor]:
        """Get per-ligand labels for auxiliary loss."""
        labels_list = []
        
        # Handle empty case
        if not per_ligand_labels or not ligand_names:
            return labels_list
        
        # Get sequence length from first available label
        first_labels = next(iter(per_ligand_labels.values()), None)
        if first_labels is None:
            return labels_list
        seq_len = len(first_labels)
        
        for lig_name in ligand_names:
            if lig_name in per_ligand_labels:
                labels = torch.tensor(per_ligand_labels[lig_name], dtype=torch.float32)
            else:
                labels = torch.zeros(seq_len, dtype=torch.float32)
            labels_list.append(labels)
        
        return labels_list


class RhoFoldApoDataset(BaseDataset):
    """
    RhoFold+ Apo Structure Dataset.
    
    - 3D loss only (no binding labels)
    - design_vector = 0 (no ions)
    - is_apo = True
    """
    
    def __init__(
        self,
        rhofold_path: str,
        rna_emb_dict: Dict,
        rna_dim: int = 1280,
        num_ions: int = 20,
        num_ligands: int = 1,
    ):
        self.rhofold_path = rhofold_path
        self.rna_emb_dict = rna_emb_dict
        self.rna_dim = rna_dim
        self.num_ions = num_ions
        self.num_ligands = max(num_ligands, 1)
        self.target_atoms = ["P", "C4'", "C1'"]
        
        self.samples = self._scan_rhofold_dirs()
        print(f"[RhoFoldApoDataset] Found {len(self.samples)} apo structures")
    
    def _scan_rhofold_dirs(self) -> List[Dict]:
        samples = []
        if not self.rhofold_path or not os.path.exists(self.rhofold_path):
            print(f"[Warning] rhofold_path not found: {self.rhofold_path}")
            return samples
        
        for entry in sorted(os.listdir(self.rhofold_path)):
            entry_path = os.path.join(self.rhofold_path, entry)
            if not os.path.isdir(entry_path):
                continue
            pdb_file = os.path.join(entry_path, "relaxed_1000_model.pdb")
            fasta_file = os.path.join(entry_path, "0.fasta")
            if os.path.exists(pdb_file):
                seq = self._read_fasta(fasta_file) if os.path.exists(fasta_file) else self._parse_seq_from_pdb(pdb_file)
                if seq and len(seq) > 0:
                    samples.append({'idx': entry, 'pdb_path': pdb_file, 'sequence': seq})
        return samples
    
    def _read_fasta(self, fasta_path: str) -> Optional[str]:
        try:
            with open(fasta_path, 'r') as f:
                lines = f.readlines()
                seq_lines = [l.strip() for l in lines if not l.startswith('>')]
                return ''.join(seq_lines).upper().replace('T', 'U')
        except Exception:
            return None
    
    def _parse_seq_from_pdb(self, pdb_path: str) -> Optional[str]:
        residue_map = {'A': 'A', 'G': 'G', 'C': 'C', 'U': 'U'}
        residues = {}
        try:
            with open(pdb_path) as f:
                for line in f:
                    if line.startswith('ATOM'):
                        res_name = line[17:20].strip()
                        res_num = int(line[22:26])
                        if res_num not in residues:
                            residues[res_num] = residue_map.get(res_name, 'N')
        except Exception:
            return None
        if not residues:
            return None
        return ''.join([residues[k] for k in sorted(residues.keys())])
    
    def __len__(self):
        return len(self.samples)
    
    def get_class_distribution(self) -> Dict[str, float]:
        total_residues = sum(len(s['sequence']) for s in self.samples)
        return {
            'total_samples': len(self.samples),
            'total_residues': total_residues,
            'total_binding': 0,
            'binding_ratio': 0.0,
            'imbalance_ratio': float('inf'),
        }
    
    def __getitem__(self, idx: int) -> Dict:
        s = self.samples[idx]
        seq = s['sequence']
        seq_len = len(seq)
        
        coords = self._parse_single_atom_pdb(s['pdb_path'], seq_len)
        rna_emb = self._compute_rna_emb(s['idx'], seq, seq_len)
        
        return {
            'pdb_id': f"apo_{s['idx']}",
            'chain_id': 'A',
            'rna_seq': seq,
            'coords': coords.float(),
            'coords_target': coords.float(),
            'labels': torch.zeros(seq_len, dtype=torch.float32),
            'ion_counts': torch.zeros(self.num_ions, dtype=torch.float32),
            'design_vector': torch.zeros(self.num_ligands, dtype=torch.float32),
            'rna_emb': rna_emb,
            'is_apo': True,
            'per_ligand_labels': [],
            'ligand_names': [],
            'smiles_list': [],
        }
    
    def _parse_single_atom_pdb(self, pdb_path: str, seq_len: int) -> torch.Tensor:
        coords = []
        try:
            with open(pdb_path) as f:
                for line in f:
                    if line.startswith('ATOM') and line[12:16].strip() == "C4'":
                        coords.append([float(line[30:38]), float(line[38:46]), float(line[46:54])])
        except Exception:
            pass
        if not coords:
            return torch.randn(seq_len, 3) * 10
        coords = torch.tensor(coords, dtype=torch.float32)
        if coords.shape[0] < seq_len:
            coords = torch.cat([coords, torch.zeros(seq_len - coords.shape[0], 3)])
        return coords[:seq_len]
    
    def _compute_rna_emb(self, idx: str, sequence: str, seq_len: int) -> torch.Tensor:
        rna_emb = self.rna_emb_dict.get(idx)
        if rna_emb is None and idx.isdigit():
            rna_emb = self.rna_emb_dict.get(int(idx))
        if rna_emb is None:
            rna_emb = self.rna_emb_dict.get(sequence)
        if rna_emb is None:
            return torch.zeros(seq_len, self.rna_dim)
        if not isinstance(rna_emb, torch.Tensor):
            rna_emb = torch.tensor(rna_emb, dtype=torch.float32)
        if rna_emb.dim() == 1:
            rna_emb = rna_emb.unsqueeze(0).expand(seq_len, -1).clone()
        if rna_emb.shape[0] > seq_len:
            rna_emb = rna_emb[:seq_len]
        elif rna_emb.shape[0] < seq_len:
            pad = torch.zeros(seq_len - rna_emb.shape[0], rna_emb.shape[1])
            rna_emb = torch.cat([rna_emb, pad])
        return rna_emb


Dataset = RNABindingDataset


def seq_to_onehot(seq: str) -> torch.Tensor:
    mapping = {'A': 0, 'U': 1, 'G': 2, 'C': 3, 'T': 1}
    onehot = torch.zeros(len(seq), 4)
    for i, base in enumerate(seq.upper()):
        if base in mapping:
            onehot[i, mapping[base]] = 1.0
        else:
            onehot[i] = 0.25
    return onehot


class CollateFn:
    def __init__(self, ligand_to_idx: Dict[str, int] = None, ligand_to_smiles: Dict[str, str] = None):
        self.ligand_to_idx = ligand_to_idx or {}
        self.ligand_to_smiles = ligand_to_smiles or {}
        # Build ordered SMILES list by ligand index
        num_ligands = len(self.ligand_to_idx)
        self.ordered_smiles = [''] * num_ligands
        for lig_name, idx in self.ligand_to_idx.items():
            if lig_name in self.ligand_to_smiles:
                self.ordered_smiles[idx] = self.ligand_to_smiles[lig_name]
    
    def __call__(self, batch: List[Dict]) -> Dict:
        out = {
            'pdb_id': [x['pdb_id'] for x in batch],
            'chain_id': [x['chain_id'] for x in batch],
            'rna_seq': [x['rna_seq'] for x in batch],
            'coords': torch.cat([x['coords'] for x in batch]),
            'coords_target': torch.cat([x['coords_target'] for x in batch]),
            'labels': torch.cat([x['labels'] for x in batch]),
            'ion_counts': torch.stack([x['ion_counts'] for x in batch]),
            'design_vector': torch.stack([x['design_vector'] for x in batch]),  # kept for backward compat
            'rna_emb': torch.cat([x['rna_emb'] for x in batch]),
            'rna_onehot': torch.cat([seq_to_onehot(x['rna_seq']) for x in batch]),
            'is_apo': torch.tensor([x.get('is_apo', False) for x in batch]),
        }
        out['batch'] = torch.cat([
            torch.full((len(x['rna_seq']),), i, dtype=torch.long) 
            for i, x in enumerate(batch)
        ])
        
        # Per-ligand data for auxiliary loss (list of lists)
        if 'per_ligand_labels' in batch[0]:
            out['per_ligand_labels'] = [x['per_ligand_labels'] for x in batch]
            out['ligand_names'] = [x['ligand_names'] for x in batch]
        
        # Per-sample ligand SMILES (non-ion only) - NEW
        # Each sample has its own list of ligand SMILES
        out['ligand_smiles_list'] = []
        for x in batch:
            sample_smiles = []
            ligand_names = x.get('ligand_names', [])
            smiles_list = x.get('smiles_list', [])
            
            # Match ligand names to SMILES, filter out ions
            for i, name in enumerate(ligand_names):
                if not is_ion(name) and i < len(smiles_list):
                    smiles = smiles_list[i]
                    if smiles and pd.notna(smiles):
                        sample_smiles.append(smiles)
            
            # Add unseen ligand SMILES
            for ul in x.get('unseen_ligands', []):
                if ul.get('smiles'):
                    sample_smiles.append(ul['smiles'])
            
            out['ligand_smiles_list'].append(sample_smiles)
        
        # Legacy: ordered SMILES list for backward compatibility
        out['smiles_list'] = self.ordered_smiles if self.ordered_smiles else ['']
        
        # Unseen ligands (not in ligand_to_idx) - per sample
        if 'unseen_ligands' in batch[0]:
            out['unseen_ligands'] = [x.get('unseen_ligands', []) for x in batch]
        
        return out


def collate_fn(batch: List[Dict]) -> Dict:
    return CollateFn()(batch)


def get_collate_fn(ligand_to_idx: Dict[str, int] = None, ligand_to_smiles: Dict[str, str] = None):
    return CollateFn(ligand_to_idx, ligand_to_smiles)