REDDA/prediction.py

"""
Drug Repositioning Prediction Script for REDDA Model

This script performs drug repurposing prediction using the trained REDDA model.
It takes an input CSV with drug SMILES and calculates similarity to existing drugs
in the knowledge graph, then predicts drug-disease associations.

Usage:
    python prediction.py --input_csv <path_to_input_csv> --target_disease <disease_id> --output_csv <output_path>
"""

import os
import sys
import argparse
import numpy as np
import pandas as pd
import torch as th
from rdkit import Chem
from rdkit.Chem import DataStructs, rdMolDescriptors
from rdkit import RDLogger
import dgl
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

# Suppress RDKit warnings
RDLogger.DisableLog('rdApp.*')

# Import local modules
from model import Model
from load_data import load
from utils import set_seed


def calculate_tanimoto_similarity(smiles1, smiles2):
    """
    Calculate Tanimoto similarity between two SMILES strings using Morgan fingerprints.
    
    Args:
        smiles1 (str): First SMILES string
        smiles2 (str): Second SMILES string
    
    Returns:
        float: Tanimoto similarity score (0-1)
    """
    try:
        mol1 = Chem.MolFromSmiles(smiles1)
        mol2 = Chem.MolFromSmiles(smiles2)
        
        if mol1 is None or mol2 is None:
            return 0.0
        
        # Generate Morgan fingerprints (radius=2, 2048 bits)
        fp1 = rdMolDescriptors.GetMorganFingerprintAsBitVect(mol1, radius=2, nBits=2048)
        fp2 = rdMolDescriptors.GetMorganFingerprintAsBitVect(mol2, radius=2, nBits=2048)
        
        # Calculate Tanimoto similarity
        similarity = DataStructs.TanimotoSimilarity(fp1, fp2)
        return similarity
    except:
        return 0.0


from concurrent.futures import ThreadPoolExecutor
from functools import partial

def calculate_similarity_for_drug(query_smiles, row):
    """
    Calculate similarity for a single drug.
    
    Args:
        query_smiles (str): Query drug SMILES
        row (pd.Series): Drug row from DataFrame
    
    Returns:
        tuple: (drug_id, similarity_score) or None if invalid SMILES
    """
    kg_smiles = row['SMILES']
    drug_id = row['Inter_ID']
    
    # Skip if SMILES is empty or NaN
    if pd.isna(kg_smiles) or kg_smiles == '':
        return None
        
    similarity = calculate_tanimoto_similarity(query_smiles, kg_smiles)
    if similarity > 0:  # Only keep non-zero similarities
        return (drug_id, similarity)
    return None

def find_similar_drugs(query_smiles, drug_df, top_k=20, n_jobs=8):
    """
    Find top-k most similar drugs to the query SMILES from the knowledge graph.
    Uses parallel processing for faster computation.
    
    Args:
        query_smiles (str): Query drug SMILES
        drug_df (pd.DataFrame): DataFrame containing drug information with SMILES
        top_k (int): Number of top similar drugs to return
        n_jobs (int): Number of parallel threads to use
    
    Returns:
        list: List of tuples (drug_inter_id, similarity_score)
    """
    # Create a partial function with fixed query_smiles
    calc_func = partial(calculate_similarity_for_drug, query_smiles)
    
    # Calculate similarities in parallel
    similarities = []
    with ThreadPoolExecutor(max_workers=n_jobs) as executor:
        results = list(executor.map(calc_func, [row for _, row in drug_df.iterrows()]))
        
    # Filter out None results and sort
    similarities = [x for x in results if x is not None]
    similarities.sort(key=lambda x: x[1], reverse=True)
    return similarities[:top_k]




def create_drug_id_mapping(drug_df):
    """
    Create mapping from Inter_ID to index for drugs.
    
    Args:
        drug_df (pd.DataFrame): Drug dataframe
    
    Returns:
        dict: Mapping from Inter_ID to index
    """
    # Create mapping and also store the reverse mapping
    id_to_idx = {drug_id: idx for idx, drug_id in enumerate(drug_df['Inter_ID'])}
    return id_to_idx


def create_disease_id_mapping(disease_df):
    """
    Create mapping from Inter_ID to index for diseases.
    
    Args:
        disease_df (pd.DataFrame): Disease dataframe
    
    Returns:
        dict: Mapping from Inter_ID to index
    """
    return {disease_id: idx for idx, disease_id in enumerate(disease_df['Inter_ID'])}


def augment_graph_with_new_drugs(input_drugs_df, drug_df, drug_id_mapping, kg_dir, output_folder):
    """
    Augment the graph by creating temporary node and edge files with new drugs.
    
    Args:
        input_drugs_df (pd.DataFrame): Input drugs DataFrame
        drug_df (pd.DataFrame): Knowledge graph drugs DataFrame
        drug_id_mapping (dict): Mapping from drug ID to index
        kg_dir (str): Base directory of knowledge graph data
        output_folder (str): Path to output folder
    
    Returns:
        tuple: (new_drug_indices, node_path, edge_path)
    """
    global temp_dir
    print("Augmenting graph with new drugs...")
    
    # Get the last drug index
    last_drug_id = int(drug_df['Inter_ID'].str.extract(r'DRUG_(\d+)').astype(int).max())
    
    # Create new drug entries
    new_drugs = []
    new_drug_indices = {}  # Maps input index to numerical index in augmented graph
    base_idx = len(drug_df)  # Start new indices after existing drugs
    
    for idx, row in input_drugs_df.iterrows():
        new_drug_id = f"DRUG_{last_drug_id + idx + 1}"
        new_drug_indices[idx] = base_idx + idx  # Store numerical index
        
        new_drug = {
            'id': f"NEW_{new_drug_id}",
            'name': f"New_Drug_{idx}",
            'source': 'Prediction',
            'type': 'drug',
            'SMILES': row['SMILES'],
            'Inter_ID': new_drug_id
        }
        new_drugs.append(new_drug)

    # Create augmented drug node file
    augmented_drug_df = pd.concat([
        drug_df,
        pd.DataFrame(new_drugs)
    ], ignore_index=True)
    
    # Create new edges based on similarity
    new_edges = []
    
    # Add similarity-based edges for new drugs
    for idx, row in tqdm(input_drugs_df.iterrows(), total=len(input_drugs_df), desc="Finding similar drugs"):
        query_smiles = row['SMILES']
        new_drug_id = new_drug_indices[idx]
        
        # Find top-20 similar drugs
        similar_drugs = find_similar_drugs(query_smiles, drug_df, top_k=20)
        
        for similar_drug_id, similarity in similar_drugs:
            # Get similar drug info
            similar_drug_row = drug_df[drug_df['Inter_ID'] == similar_drug_id].iloc[0]
            similar_drug_name = similar_drug_row['name']
            similar_drug_idx = drug_id_mapping[similar_drug_id]
            
            # Add bidirectional edges with similarity score
            new_edges.append({
                'relation': 'drug_drug',
                'display_relation': 'structural_similarity',
                'x_index': base_idx + idx,  # New drug index
                'x_id': f'DRUG_{new_drug_id}',
                'x_type': 'drug',
                'x_name': f"New_Drug_{idx}",
                'x_source': 'Prediction',
                'y_index': similar_drug_idx,  # Existing drug index
                'y_id': similar_drug_id,
                'y_type': 'drug',
                'y_name': similar_drug_name,
                'y_source': 'DrugBank',
                'similarity_score': similarity
            })
            new_edges.append({
                'relation': 'drug_drug',
                'display_relation': 'structural_similarity',
                'x_index': similar_drug_idx,  # Existing drug index
                'x_id': similar_drug_id,
                'x_type': 'drug',
                'x_name': similar_drug_name,
                'x_source': 'DrugBank',
                'y_index': base_idx + idx,  # New drug index
                'y_id': f'DRUG_{new_drug_id}',
                'y_type': 'drug',
                'y_name': f"New_Drug_{idx}",
                'y_source': 'Prediction',
                'similarity_score': similarity
            })
    
    # Create new edge DataFrame with correct format
    new_edge_df = pd.DataFrame(new_edges)

    # Create temporary directory in output folder
    temp_dir = os.path.join(output_folder, 'temp')
    os.makedirs(temp_dir, exist_ok=True)
    
    # Save augmented drug nodes and edges
    temp_node_path = os.path.join(temp_dir, 'drug.csv')
    temp_edge_path = os.path.join(temp_dir, 'drug_drug.csv')
    
    augmented_drug_df.to_csv(temp_node_path, index=False)
    
    # Read original drug_drug edges and combine with new edges
    orig_edges = pd.read_csv(os.path.join(kg_dir, 'edge/drug_drug.csv'))
    
    # Ensure new edges have all required columns
    for col in orig_edges.columns:
        if col not in new_edge_df.columns:
            new_edge_df[col] = ''  # Fill missing columns with empty string
    
    # Combine original and new edges
    combined_edges = pd.concat([orig_edges, new_edge_df], ignore_index=True)
    
    # Save combined edges
    combined_edges.to_csv(temp_edge_path, index=False)
    
    # Create symlinks for all other files
    for file_name in os.listdir(os.path.join(kg_dir, 'node')):
        if file_name != 'drug.csv':
            src = os.path.join(kg_dir, 'node', file_name)
            dst = os.path.join(temp_dir, file_name)
            if not os.path.exists(dst):
                os.symlink(src, dst)
    
    for file_name in os.listdir(os.path.join(kg_dir, 'edge')):
        if file_name != 'drug_drug.csv':
            src = os.path.join(kg_dir, 'edge', file_name)
            dst = os.path.join(temp_dir, file_name)
            if not os.path.exists(dst):
                os.symlink(src, dst)
    
    return new_drug_indices, temp_dir, temp_dir


def load_models(model_dir, g, feature_dim, device, hidden_feats=128, num_heads=5, dropout=0.0):
    """
    Load all 5 trained model checkpoints.
    
    Args:
        model_dir (str): Directory containing model checkpoints
        g (dgl.DGLHeteroGraph): Graph for model initialization
        feature_dim (int): Feature dimension
        device (torch.device): Device to load models on
        hidden_feats (int): Hidden layer dimension (default: 64)
        num_heads (int): Number of attention heads (default: 5)
        dropout (float): Dropout rate (default: 0.0)
    
    Returns:
        list: List of loaded models
    """
    models = []
    
    for fold in range(1, 6):  # model_1.pth to model_5.pth
        model_path = os.path.join(model_dir, f"model_{fold}.pth")
        
        if not os.path.exists(model_path):
            print(f"Warning: Model checkpoint {model_path} not found")
            continue
        
        # Initialize model with same architecture as training
        model = Model(
            etypes=g.etypes,
            ntypes=g.ntypes,
            in_feats=feature_dim,
            hidden_feats=hidden_feats,
            num_heads=num_heads,
            dropout=dropout,
        ).to(device)
        
        # Load checkpoint
        checkpoint = th.load(model_path, map_location=device)
        model.load_state_dict(checkpoint)
        model.eval()
        
        models.append(model)
        print(f"Loaded model from {model_path}")
    
    return models


def ensemble_predict(models, g, feature, device):
    """
    Perform ensemble prediction using multiple model checkpoints.
    
    Args:
        models (list): List of trained models
        g (dgl.DGLHeteroGraph): Input graph
        feature (dict): Node features
        device (torch.device): Device for computation
    
    Returns:
        tuple: (mean_predictions, std_predictions)
    """
    predictions = []
    
    print("Performing ensemble prediction...")
    with th.no_grad():
        for i, model in enumerate(models):
            print(f"Predicting with model {i+1}/{len(models)}")
            pred = th.sigmoid(model(g, feature)).cpu().numpy()
            predictions.append(pred)
    
    # Calculate mean and standard deviation
    predictions = np.array(predictions)
    mean_pred = np.mean(predictions, axis=0)
    std_pred = np.std(predictions, axis=0)
    
    return mean_pred, std_pred


# Global variables
temp_dir = None

def main():
    global temp_dir
    parser = argparse.ArgumentParser(description='REDDA Drug Repositioning Prediction')
    parser.add_argument('--input_csv', type=str, required=True,
                        help='Path to input CSV file with drug SMILES')
    parser.add_argument('--target_disease', type=str, required=True,
                        help='Target disease Inter_ID (e.g., DISEASE_123)')
    parser.add_argument('--output_folder', type=str, required=True,
                        help='Path to output folder for results and temporary files')
    parser.add_argument('--device_id', type=str, default=None,
                        help='GPU device ID (default: CPU)')
    parser.add_argument('--dataset', type=str, default='KGdataset',
                        choices=['KGdataset', 'KGdataset_tiny'],
                        help='Dataset to use')
    parser.add_argument('--model_dir', type=str, default='result_kg_42',
                        help='Directory containing model checkpoints')
    
    pred_args = parser.parse_args()
    
    # Set device
    if pred_args.device_id:
        device = th.device(f"cuda:{pred_args.device_id}")
        print(f"Using GPU: {pred_args.device_id}")
    else:
        device = th.device("cpu")
        print("Using CPU")
    
    # Set seed for reproducibility
    set_seed(42)
    
    print("Loading input data...")
    # Load input CSV
    input_df = pd.read_csv(pred_args.input_csv)
    
    # Validate input CSV format
    required_columns = ['SMILES']
    if not all(col in input_df.columns for col in required_columns):
        raise ValueError(f"Input CSV must contain columns: {required_columns}")
    
    print(f"Input contains {len(input_df)} drugs")
    
    # Load knowledge graph data
    print("Loading knowledge graph data...")
    drug_df = pd.read_csv('/vast/yg3191/AIVS/kg/node/drug.csv')
    disease_df = pd.read_csv('/vast/yg3191/AIVS/kg/node/disease.csv')
    
    # Create ID mappings
    drug_id_mapping = create_drug_id_mapping(drug_df)
    disease_id_mapping = create_disease_id_mapping(disease_df)
    
    # Validate target disease
    if pred_args.target_disease not in disease_id_mapping:
        available_diseases = list(disease_id_mapping.keys())[:10]  # Show first 10
        raise ValueError(f"Target disease '{pred_args.target_disease}' not found in knowledge graph. "
                        f"Available diseases (first 10): {available_diseases}")
    
    target_disease_idx = disease_id_mapping[pred_args.target_disease]
    print(f"Target disease: {pred_args.target_disease} (index: {target_disease_idx})")
    
    # Load original graph
    print("Loading original graph...")
    g = load(pred_args.dataset).to(device)
    
    print(f"Original graph: {g.num_nodes('drug')} drugs, {g.num_nodes('disease')} diseases")
    
    # Create output folder if not exists
    os.makedirs(pred_args.output_folder, exist_ok=True)
    
    # Create temporary files with augmented graph data
    kg_dir = '/vast/yg3191/AIVS/kg'
    new_drug_indices, node_path, edge_path = augment_graph_with_new_drugs(
        input_df, drug_df, drug_id_mapping, kg_dir, pred_args.output_folder
    )
    
    # Build augmented graph using _build_heterograph
    from load_data import _build_heterograph
    
    # Define node and edge files (same as in load_KGdataset)
    node_files = {
        'drug': 'drug.csv',
        'disease': 'disease.csv',
        'protein': 'protein.csv',
        # 'bioprocess': 'bioprocess.csv',
        # 'cellcomp': 'cellcomp.csv',
        # 'molfunc': 'molfunc.csv',
        # 'pathway': 'pathway.csv',
        # 'phenotype': 'phenotype.csv',
        # 'exposure': 'exposure.csv',
        # 'effect': 'effect.csv',
    }
    
    edge_files = {
        'drug_drug': 'drug_drug.csv',
        # 'drug_effect': 'drug_effect.csv',
        'drug_protein': 'drug_protein.csv',
        'drug_disease': 'drug_disease_indication.csv',
        'protein_protein': 'protein_protein.csv',
        # 'protein_bioprocess': 'protein_bioprocess.csv',
        # 'protein_cellcomp': 'protein_cellcomp.csv',
        # 'protein_molfunc': 'protein_molfunc.csv',
        # 'protein_pathway': 'protein_pathway.csv',
        'protein_disease': 'protein_disease.csv',
        'disease_disease': 'disease_disease.csv',
        # 'disease_phenotype': 'disease_phenotype_positive.csv',
        # 'disease_exposure': 'disease_exposure.csv',
        # 'bioprocess_bioprocess': 'bioprocess_bioprocess.csv',
        # 'cellcomp_cellcomp': 'cellcomp_cellcomp.csv',
        # 'molfunc_molfunc': 'molfunc_molfunc.csv',
        # 'pathway_pathway': 'pathway_pathway.csv',
        # 'phenotype_phenotype': 'phenotype_phenotype.csv',
    }
    
    augmented_g = _build_heterograph(
        node_path=node_path,
        edge_path=edge_path,
        node_files=node_files,
        edge_files=edge_files,
        feature_dim=128
    ).to(device)
    
    # Extract features
    feature = {ntype: augmented_g.nodes[ntype].data['h'] for ntype in augmented_g.ntypes}
    feature_dim = feature['drug'].shape[1]
    
    # Load trained models
    print("Loading trained models...")
    # Use default model parameters that match the training configuration
    models = load_models(pred_args.model_dir, augmented_g, feature_dim, device, 
                        hidden_feats=128, num_heads=8, dropout=0.4)
    
    if not models:
        raise ValueError("No model checkpoints found!")
    
    print(f"Loaded {len(models)} model checkpoints")
    
    # Perform ensemble prediction
    mean_pred, std_pred = ensemble_predict(models, augmented_g, feature, device)
    
    # Extract predictions for input drugs and target disease
    print("Extracting predictions...")
    results = []
    print(mean_pred.shape)
    print(std_pred.shape)
    print(new_drug_indices)
    for idx, row in input_df.iterrows():
        # Get the numerical index for this drug in the augmented graph
        drug_idx = new_drug_indices[idx]  # This is now a numerical index
        
        # Get prediction Score for this drug-disease pair
        Score_mean = float(mean_pred[drug_idx, target_disease_idx])
        Score_std = float(std_pred[drug_idx, target_disease_idx])
        
        # Create result row
        result_row = row.copy()
        result_row['Target_Disease'] = pred_args.target_disease
        result_row['Score_Mean'] = Score_mean
        result_row['Score_Std'] = Score_std
        
        results.append(result_row)
    
    # Create output DataFrame
    output_df = pd.DataFrame(results)
    
    # Create output folder if not exists
    os.makedirs(pred_args.output_folder, exist_ok=True)
    
    # Save results
    output_path = os.path.join(pred_args.output_folder, 'result.csv')
    output_df.to_csv(output_path, index=False)
    print(f"Predictions saved to: {output_path}")
    
    # Print summary statistics
    print("\nPrediction Summary:")
    print(f"Mean Score: {output_df['Score_Mean'].mean():.4f}")
    print(f"Std Score: {output_df['Score_Mean'].std():.4f}")
    print(f"Max Score: {output_df['Score_Mean'].max():.4f}")
    print(f"Min Score: {output_df['Score_Mean'].min():.4f}")
    
    # Show top predictions
    print(f"\nTop 5 predictions:")
    top_predictions = output_df.nlargest(5, 'Score_Mean')
    for idx, row in top_predictions.iterrows():
        print(f"  SMILES: {row['SMILES'][:50]}... | Score: {row['Score_Mean']:.4f} ± {row['Score_Std']:.4f}")
    
    # Clean up temporary files (optional, comment out to keep temp files for debugging)
    try:
        # Remove symlinks first
        for file_name in os.listdir(temp_dir):
            file_path = os.path.join(temp_dir, file_name)
            if os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isfile(file_path):
                os.remove(file_path)
        # Remove temp directory
        os.rmdir(temp_dir)
        print("\nTemporary files cleaned up")
    except Exception as e:
        print(f"\nWarning: Could not clean up temporary files: {e}")
        print(f"Temporary files remain in: {temp_dir}")


if __name__ == "__main__":
    main()