src/data.py

"""
Data loading and preprocessing module for chromatographic RT prediction.
Handles SMILES string to molecular graph conversion and feature extraction.
"""

import pandas as pd
import numpy as np
from rdkit import Chem, DataStructs
from rdkit.Chem import Descriptors, Crippen, Lipinski, rdMolDescriptors, rdFingerprintGenerator
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import KFold, StratifiedKFold
import torch
from torch_geometric.data import Data, DataLoader
from typing import List, Tuple, Dict, Optional, Any
import warnings
warnings.filterwarnings('ignore')


class MolecularFeatureExtractor:
    """Extract molecular features from SMILES strings."""
    
    def __init__(self):
        self.atom_features = [
            'atomic_num', 'degree', 'formal_charge', 'hybridization',
            'is_aromatic', 'total_numHs', 'mass', 'chiral_tag'
        ]
        self.bond_features = ['bond_type', 'conjugated', 'is_ring', 'stereo']
        
    def get_atom_features(self, atom) -> List[float]:
        """Extract atom-level features including chirality."""
        chiral_tag = atom.GetChiralTag()
        # Convert chiral tag to numerical value
        chiral_value = {
            Chem.ChiralType.CHI_UNSPECIFIED: 0.0,
            Chem.ChiralType.CHI_TETRAHEDRAL_CW: 1.0,    # R
            Chem.ChiralType.CHI_TETRAHEDRAL_CCW: 2.0,   # S
            Chem.ChiralType.CHI_OTHER: 3.0,
            Chem.ChiralType.CHI_TETRAHEDRAL: 4.0,
            Chem.ChiralType.CHI_ALLENE: 5.0,
            Chem.ChiralType.CHI_SQUAREPLANAR: 6.0,
            Chem.ChiralType.CHI_TRIGONALBIPYRAMIDAL: 7.0,
            Chem.ChiralType.CHI_OCTAHEDRAL: 8.0
        }.get(chiral_tag, 0.0)
        
        return [
            atom.GetAtomicNum(),
            atom.GetTotalDegree(),
            atom.GetFormalCharge(),
            int(atom.GetHybridization()),
            int(atom.GetIsAromatic()),
            atom.GetTotalNumHs(),
            atom.GetMass() / 100.0,  # Normalized mass
            chiral_value / 8.0       # Normalized chirality
        ]
    
    def get_bond_features(self, bond) -> List[float]:
        """Extract bond-level features including stereochemistry."""
        stereo = bond.GetStereo()
        # Convert stereo to numerical value
        stereo_value = {
            Chem.BondStereo.STEREONONE: 0.0,
            Chem.BondStereo.STEREOANY: 1.0,
            Chem.BondStereo.STEREOZ: 2.0,     # Z (cis)
            Chem.BondStereo.STEREOE: 3.0,     # E (trans)
            Chem.BondStereo.STEREOCIS: 4.0,
            Chem.BondStereo.STEREOTRANS: 5.0
        }.get(stereo, 0.0)
        
        return [
            int(bond.GetBondType()),
            int(bond.GetIsConjugated()),
            int(bond.IsInRing()),
            stereo_value / 5.0  # Normalized stereochemistry
        ]
    
    def smiles_to_graph(self, smiles: str) -> Optional[Data]:
        """Convert SMILES string to PyTorch Geometric graph with stereochemistry."""
        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return None
            
            # Try to assign stereochemistry from SMILES
            try:
                Chem.AssignStereochemistry(mol, cleanIt=True, force=True, flagPossibleStereoCenters=True)
            except:
                pass  # Continue even if stereochemistry assignment fails
                
            # Atom features
            atom_features = []
            for atom in mol.GetAtoms():
                atom_features.append(self.get_atom_features(atom))
            
            # Bond features and edge indices
            edge_indices = []
            edge_attrs = []
            
            for bond in mol.GetBonds():
                i = bond.GetBeginAtomIdx()
                j = bond.GetEndAtomIdx()
                
                edge_indices.extend([[i, j], [j, i]])  # Undirected graph
                bond_feat = self.get_bond_features(bond)
                edge_attrs.extend([bond_feat, bond_feat])
            
            # Convert to tensors
            x = torch.tensor(atom_features, dtype=torch.float)
            edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous() if edge_indices else torch.empty((2, 0), dtype=torch.long)
            edge_attr = torch.tensor(edge_attrs, dtype=torch.float) if edge_attrs else None
            
            return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
            
        except Exception as e:
            print(f"Error processing SMILES {smiles}: {e}")
            return None
    
    def get_molecular_descriptors(self, smiles: str) -> Dict[str, float]:
        """Extract essential molecular descriptors (simplified for speed)."""
        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return {}
                
            descriptors = {
                # Basic descriptors - fast to compute
                'MW': Descriptors.MolWt(mol),  # type: ignore[attr-defined]
                'LogP': Crippen.MolLogP(mol),  # type: ignore[attr-defined]
                'NumHDonors': Lipinski.NumHDonors(mol),  # type: ignore[attr-defined]
                'NumHAcceptors': Lipinski.NumHAcceptors(mol),  # type: ignore[attr-defined]
                'NumRotatableBonds': Descriptors.NumRotatableBonds(mol),  # type: ignore[attr-defined]
                'TPSA': Descriptors.TPSA(mol),  # type: ignore[attr-defined]
                'NumAromaticRings': Descriptors.NumAromaticRings(mol),  # type: ignore[attr-defined]
                'NumAliphaticRings': Descriptors.NumAliphaticRings(mol),  # type: ignore[attr-defined]
                'HeavyAtomCount': mol.GetNumHeavyAtoms(),
                'NumHeteroatoms': Descriptors.NumHeteroatoms(mol),  # type: ignore[attr-defined]
                'RingCount': Descriptors.RingCount(mol),  # type: ignore[attr-defined]
                'BertzCT': Descriptors.BertzCT(mol),  # type: ignore[attr-defined]
                'MolMR': Crippen.MolMR(mol),  # type: ignore[attr-defined]
                'LabuteASA': Descriptors.LabuteASA(mol),  # type: ignore[attr-defined]
                'HallKierAlpha': Descriptors.HallKierAlpha(mol),  # type: ignore[attr-defined]
                'Kappa1': Descriptors.Kappa1(mol),  # type: ignore[attr-defined]
                'Kappa2': Descriptors.Kappa2(mol),  # type: ignore[attr-defined]
                'Kappa3': Descriptors.Kappa3(mol),  # type: ignore[attr-defined]
            }
            
            # Add optional descriptors with error handling
            try:
                descriptors['FractionCsp3'] = Descriptors.FractionCsp3(mol)  # type: ignore[attr-defined]
            except:
                descriptors['FractionCsp3'] = 0.0
            
            # Add simple calculated features
            descriptors['HeavyAtomRatio'] = descriptors['HeavyAtomCount'] / max(mol.GetNumAtoms(), 1)
            descriptors['AromaticRatio'] = descriptors['NumAromaticRings'] / max(descriptors['RingCount'], 1)
            descriptors['RotatableBondRatio'] = descriptors['NumRotatableBonds'] / max(descriptors['HeavyAtomCount'], 1)
            
            return descriptors
            
        except Exception as e:
            print(f"Error calculating descriptors for {smiles}: {e}")
            # Return default values if calculation fails
            return {
                'MW': 200.0, 'LogP': 2.0, 'NumHDonors': 1.0, 'NumHAcceptors': 2.0,
                'NumRotatableBonds': 3.0, 'TPSA': 50.0, 'NumAromaticRings': 1.0,
                'NumAliphaticRings': 0.0, 'HeavyAtomCount': 15.0, 'NumHeteroatoms': 2.0,
                'RingCount': 1.0, 'BertzCT': 100.0, 'MolMR': 60.0, 'LabuteASA': 100.0,
                'HallKierAlpha': 5.0, 'Kappa1': 3.0, 'Kappa2': 2.0, 'Kappa3': 1.0,
                'FractionCsp3': 0.5, 'HeavyAtomRatio': 0.75, 'AromaticRatio': 1.0, 'RotatableBondRatio': 0.2
            }

    def get_morgan_fingerprint(
        self,
        smiles: str,
        *,
        n_bits: int = 2048,
        radius: int = 2,
        use_chirality: bool = True,
        use_features: bool = False,
    ) -> np.ndarray:
        """Generate a Morgan fingerprint as a dense float32 numpy array."""

        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                return np.zeros((n_bits,), dtype=np.float32)

            if use_features:
                atom_inv_gen = rdFingerprintGenerator.GetMorganFeatureAtomInvGen()
            else:
                # Include ring membership information to stay close to RDKit defaults
                atom_inv_gen = rdFingerprintGenerator.GetMorganAtomInvGen(includeRingMembership=True)

            generator = rdFingerprintGenerator.GetMorganGenerator(
                radius=radius,
                fpSize=n_bits,
                includeChirality=use_chirality,
                atomInvariantsGenerator=atom_inv_gen,
            )
            fingerprint = generator.GetFingerprint(mol)
            array = np.zeros((n_bits,), dtype=np.float32)
            DataStructs.ConvertToNumpyArray(fingerprint, array)
            return array

        except Exception as e:
            print(f"Error generating fingerprint for {smiles}: {e}")
            return np.zeros((n_bits,), dtype=np.float32)


class RTDataset:
    """Dataset class for retention time prediction."""
    
    def __init__(self, data_path: str, is_test: bool = False):
        self.data_path = data_path
        self.is_test = is_test
        self.feature_extractor = MolecularFeatureExtractor()
        self.lab_encoder = None
        self.valid_indices: Optional[np.ndarray] = None
        
        # Load data
        self.df = pd.read_csv(data_path)
        print(f"Loaded {'test' if is_test else 'train'} data: {self.df.shape}")
        
    def preprocess_data(self, lab_encoder=None) -> Tuple[List[Data], np.ndarray, Optional[np.ndarray]]:
        """
        Preprocess the dataset.
        
        Returns:
            graphs: List of PyTorch Geometric Data objects
            lab_features: Encoded lab features
            targets: Target RT values (None for test data)
        """
        # Encode lab features
        if lab_encoder is None:
            self.lab_encoder = LabelEncoder()
            lab_features = self.lab_encoder.fit_transform(self.df['Lab'])
        else:
            self.lab_encoder = lab_encoder
            lab_features = self.lab_encoder.transform(self.df['Lab'])
        
        # Convert SMILES to graphs
        graphs = []
        valid_indices = []
        
        print("Converting SMILES to molecular graphs...")
        for idx, smiles in enumerate(self.df['SMILES']):
            graph = self.feature_extractor.smiles_to_graph(smiles)
            if graph is not None:
                graphs.append(graph)
                valid_indices.append(idx)
            else:
                print(f"Failed to process SMILES at index {idx}: {smiles}")

        lab_features = np.asarray(lab_features)
        lab_features = lab_features[valid_indices]

        targets_array = None if self.is_test else np.asarray(self.df['RT'].values)
        targets = None if targets_array is None else targets_array[valid_indices]

        print(f"Successfully processed {len(graphs)} molecules out of {len(self.df)}")

        self.valid_indices = np.array(valid_indices, dtype=int)

        return graphs, lab_features, targets
    
    def get_molecular_descriptors_df(self) -> pd.DataFrame:
        """Get traditional molecular descriptors as DataFrame."""

        descriptors_list: List[Dict[str, Any]] = []

        print("Extracting molecular descriptors...")
        for idx, smiles in enumerate(self.df['SMILES']):
            desc = self.feature_extractor.get_molecular_descriptors(smiles)
            desc_entry: Dict[str, Any] = dict(desc)
            desc_entry['index'] = idx
            desc_entry['Lab'] = self.df.loc[idx, 'Lab']
            if not self.is_test:
                desc_entry['RT'] = self.df.loc[idx, 'RT']
            descriptors_list.append(desc_entry)

        desc_df = pd.DataFrame(descriptors_list)

        # Handle missing values
        numeric_cols = desc_df.select_dtypes(include=[np.number]).columns
        desc_df[numeric_cols] = desc_df[numeric_cols].fillna(desc_df[numeric_cols].median())

        return desc_df

    def get_fingerprint_matrix(
        self,
        *,
        n_bits: int = 2048,
        radius: int = 2,
        use_chirality: bool = True,
        use_features: bool = False,
    ) -> np.ndarray:
        """Return Morgan fingerprints aligned with the valid graph indices."""

        if self.valid_indices is None:
            raise RuntimeError("Call preprocess_data before requesting fingerprints.")

        fingerprints: List[np.ndarray] = []
        for idx in self.valid_indices:
            smiles = str(self.df.loc[idx, 'SMILES'])
            fp = self.feature_extractor.get_morgan_fingerprint(
                smiles,
                n_bits=n_bits,
                radius=radius,
                use_chirality=use_chirality,
                use_features=use_features,
            )
            fingerprints.append(fp)

        if not fingerprints:
            return np.zeros((0, n_bits), dtype=np.float32)

        return np.stack(fingerprints, axis=0)


def create_cross_validation_splits(n_samples: int, n_splits: int = 5, 
                                 stratify_col: Optional[np.ndarray] = None, 
                                 random_state: int = 42) -> List[Tuple[np.ndarray, np.ndarray]]:
    """Create cross-validation splits."""
    if stratify_col is not None:
        # Use stratified k-fold for categorical stratification
        kf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
        splits = list(kf.split(np.arange(n_samples), stratify_col))
    else:
        # Use regular k-fold
        kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)
        splits = list(kf.split(np.arange(n_samples)))
    
    return splits


def create_graph_dataloader(graphs: List[Data], lab_features: np.ndarray, 
                          targets: Optional[np.ndarray] = None, 
                          indices: Optional[np.ndarray] = None,
                          batch_size: int = 32, shuffle: bool = True):
    """Create PyTorch Geometric DataLoader."""
    if indices is not None:
        selected_graphs = [graphs[i] for i in indices]
        selected_lab_features = lab_features[indices]
        selected_targets = targets[indices] if targets is not None else None
    else:
        selected_graphs = graphs
        selected_lab_features = lab_features
        selected_targets = targets
    
    # Add lab features and targets to graph data
    for i, graph in enumerate(selected_graphs):
        graph.lab_feature = torch.tensor([selected_lab_features[i]], dtype=torch.long)
        if selected_targets is not None:
            graph.y = torch.tensor([selected_targets[i]], dtype=torch.float)
    
    return DataLoader(selected_graphs, batch_size=batch_size, shuffle=shuffle)


if __name__ == "__main__":
    # Test the data loading and preprocessing
    train_dataset = RTDataset("train.csv", is_test=False)
    graphs, lab_features, targets = train_dataset.preprocess_data()
    
    print(f"Number of graphs: {len(graphs)}")
    print(f"Lab features shape: {lab_features.shape}")
    if targets is not None:
        print(f"Targets shape: {targets.shape}")
    print(f"Example graph: {graphs[0]}")
    
    # Test descriptor extraction
    desc_df = train_dataset.get_molecular_descriptors_df()
    print(f"Descriptors shape: {desc_df.shape}")
    print(f"Descriptors columns: {list(desc_df.columns)}")