sample_dataset.py

"""
Main script for sampling molecules from datasets (e.g., CrossDocked).

This script processes multiple protein-ligand pairs from a dataset and generates
drug-like molecules for each pocket using various initial atom selection methods.
"""

import argparse
import json
import os
import shutil
import sys

import numpy as np
import torch
from rdkit import Chem

# Setup paths for PocketHotspot and MolSnapper imports
# Import path_setup using absolute path to avoid circular dependency
_pocket_hotspot_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _pocket_hotspot_root)
from utils.path_setup import setup_paths

mol_snapper_path = setup_paths()

# Local imports from MolSnapper (path added at runtime via setup_paths)
from models.bond_predictor import BondPredictor  # type: ignore[import-untyped]
from models.model import MolDiff  # type: ignore[import-untyped]
from utils.misc import (  # type: ignore[import-untyped]
    EasyDict,
    get_logger,
    get_new_log_dir,
    load_config,
    seed_all
)
from utils.reconstruct import (  # type: ignore[import-untyped]
    MolReconsError,
    reconstruct_from_generated_with_edges
)
from utils.sample import seperate_outputs  # type: ignore[import-untyped]
from utils.transforms import FeaturizeMol, make_data_placeholder  # type: ignore[import-untyped]

# Local imports from PocketHotspot
from utils.cavity_detection import get_cavity_detection_functions
from utils.checkpoint_utils import resolve_checkpoint_path
from utils.hotspot_dataset import get_test_dataloader
from utils.io_utils import (
    save_cavity_points_pdb,
    save_pocket_pdb,
    save_reference_atoms_from_coords
)
from utils.sampling_utils import (
    build_method_params,
    prepare_pharmacophore_batch,
    print_pool_status,
    select_initial_atoms
)


def main():
    parser = argparse.ArgumentParser(
        description='Sampling with datasets (crossdocked)',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        epilog='''
Examples:
  # Basic usage with default settings:
  python sample_dataset.py --dataset_dir ./data/crossdocked

  # Use pharmacophore_locator method:
  python sample_dataset.py --dataset_dir ./data/crossdocked --ref_atoms_method pharmacophore_locator

  # Use hbond_predictor with custom parameters:
  python sample_dataset.py --dataset_dir ./data/crossdocked --ref_atoms_method hbond_predictor --cavity_max_dist 5.0

  # Use score_based method:
  python sample_dataset.py --dataset_dir ./data/crossdocked --ref_atoms_method score_based --atom_fraction 0.3

For more information, use --help or -h to see all available options.
        '''
    )
    # Default config path relative to MolSnapper directory
    parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    default_config = os.path.join(parent_dir, 'MolSnapper', 'configs', 'sample', 'sample_MolDiff.yml')
    
    # Input files
    parser.add_argument(
        '--config',
        type=str,
        default=default_config if os.path.exists(default_config) else './configs/sample/sample_MolDiff.yml',
        help='Model configuration file'
    )
    parser.add_argument(
        '--dataset_dir',
        type=str,
        default='./data/crossdocked',
        help='Directory with dataset files (crossdocked)'
    )
    parser.add_argument('--outdir', type=str, default='./outputs',
                       help='Output directory')
    
    # General configuration
    parser.add_argument('--device', type=str, default='cuda:0',
                       help='Device to use (cuda:0 or cpu)')
    
    # Generation parameters
    parser.add_argument('--batch_size', type=int, default=0,
                       help='Batch size for generation (0 uses config default)')
    parser.add_argument('--clash_rate', type=float, default=0.1,
                       help='Clash rate for pipeline')
    
    # Reference atom selection method
    parser.add_argument(
        '--ref_atoms_method',
        type=str,
        default='pharmacophore_locator',
        choices=['score_based', 'pharmacophore_locator', 'hbond_predictor', 'random'],
        help='Method to select reference atoms'
    )
    
    # Parameters for score_based method
    parser.add_argument('--atom_fraction', type=float, default=0.2, 
                       help='Fraction of best atoms (score_based)')
    
    # Parameters for pharmacophore_locator method
    parser.add_argument('--cutoff', type=float, default=6.0,
                       help='Maximum distance to sum contributions in pharmacophore_locator')
    parser.add_argument('--top_k_per_type', type=int, default=3,
                       help='Maximum number of pharmacophores per type in pharmacophore_locator')
    
    # Parameters for hbond_predictor method
    parser.add_argument(
        '--hbond_model_path',
        type=str,
        default='trained_hbond_predictor/best_model.pt',
        help='Path to trained EGNN model for hbond_predictor'
    )
    parser.add_argument(
        '--cavity_max_dist',
        type=float,
        default=4,
        help='Maximum distance to cavity to filter H-bond candidates (Å)'
    )
    
    # Parameters for random method
    parser.add_argument('--random_num_atoms', type=int, default=5,
                       help='Number of atoms to place randomly')
    parser.add_argument('--random_min_distance', type=float, default=1.5,
                       help='Minimum distance between random atoms in Å')
    parser.add_argument('--random_element_type', type=str, default='random', choices=['O','N','C','random'],
                       help='Element type for random atoms')
    parser.add_argument('--random_seed', type=int, default=None,
                       help='Seed for reproducibility in random selection')
    
    # Pocket/cavity detection
    parser.add_argument(
        '--pocket_detection',
        type=str,
        default='ligand_proximity',
        choices=[
            'kvfinder', 'kvfinder_interactive', 'kvfinder_with_ligand', 'ligand_coords',
            'ligand_proximity', 'auto'
        ],
        help=(
            'Cavity detection mode: a) kvfinder, b) kvfinder_interactive, c) kvfinder_with_ligand, '
            'd) ligand_coords, e) ligand_proximity, f) auto (e->a)'
        )
    )
    parser.add_argument('--ligand_proximity_radius', type=float, default=4.0,
                        help='Radius (Å) to select contact atoms and grid in ligand_proximity mode')
    args = parser.parse_args()

    # Load configs
    config = load_config(args.config)
    config_name = os.path.basename(args.config)[:os.path.basename(args.config).rfind('.')]
    seed_all(config.sample.seed + np.sum([ord(s) for s in args.outdir]))

    # Resolve checkpoint paths relative to MolSnapper
    config.model.checkpoint = resolve_checkpoint_path(config.model.checkpoint, mol_snapper_path)
    if 'bond_predictor' in config:
        config.bond_predictor = resolve_checkpoint_path(config.bond_predictor, mol_snapper_path)

    # Load checkpoint and train config
    ckpt = torch.load(config.model.checkpoint, map_location=args.device)
    train_config = ckpt['config']

    # Base logging (subdirectories will be created per receptor or per uuid)
    log_root = args.outdir.replace('outputs', 'outputs_vscode') if sys.argv[0].startswith('/data') else args.outdir
    base_tag = f"_method_{args.ref_atoms_method}"
    log_dir = get_new_log_dir(log_root, prefix=config_name, tag=base_tag)
    logger = get_logger('sample_dataset', log_dir)
    logger.info(args)
    logger.info(config)
    shutil.copyfile(args.config, os.path.join(log_dir, os.path.basename(args.config)))

    # Transform and loader
    logger.info('Loading data placeholder...')
    featurizer = FeaturizeMol(
        train_config.chem.atomic_numbers,
        train_config.chem.mol_bond_types,
        use_mask_node=train_config.transform.use_mask_node,
        use_mask_edge=train_config.transform.use_mask_edge,
    )
    add_edge = getattr(config.sample, 'add_edge', None)

    pocket_loader = get_test_dataloader(args.dataset_dir, torch.device('cpu:0'))

    # Model
    logger.info('Loading diffusion model...')
    if train_config.model.name != 'diffusion':
        raise NotImplementedError
    model = MolDiff(
        config=train_config.model,
        num_node_types=featurizer.num_node_types,
        num_edge_types=featurizer.num_edge_types,
    ).to(args.device)
    model.load_state_dict(ckpt['model'])
    model.eval()

    # Bond predictor and guidance
    if 'bond_predictor' in config:
        logger.info('Building bond predictor...')
        ckpt_bond = torch.load(config.bond_predictor, map_location=args.device)
        bond_predictor = BondPredictor(
            ckpt_bond['config']['model'],
            featurizer.num_node_types,
            featurizer.num_edge_types - 1,
        ).to(args.device)
        bond_predictor.load_state_dict(ckpt_bond['model'])
        bond_predictor.eval()
    else:
        bond_predictor = None

    guidance = config.sample.guidance if 'guidance' in config.sample else None

    # Load filter (optional - if file doesn't exist, process all dataset)
    filtered_path = os.path.join(args.dataset_dir, 'filtered_uuid.json')
    filtered_uuid = None
    if os.path.exists(filtered_path):
        try:
            with open(filtered_path, 'r') as f:
                data_ids = json.load(f)
                filtered_uuid = set(data_ids)
                logger.info(f"Filter file loaded: {len(filtered_uuid)} UUIDs to process")
        except Exception as e:
            logger.warning(f"Could not load filter file {filtered_path}: {e}. Processing all dataset.")
            filtered_uuid = None
    else:
        logger.info(f"Filter file not found at {filtered_path}. Processing all dataset.")

    
    # Global pool to accumulate all results
    global_pool = EasyDict({'failed': [], 'finished': []})
    
    # Create global SDF directory
    sdf_dir = log_dir + '_SDF'
    os.makedirs(sdf_dir, exist_ok=True)

    for i, data in enumerate(pocket_loader):
        # Skip if filter is active and UUID is not in the filter list
        if filtered_uuid is not None and (data['uuid'][0] not in filtered_uuid):
            continue
           
        # Local pool for this specific receptor
        pool = EasyDict({'failed': [], 'finished': []}) 

        uuid = data['uuid'][0]
        pocket_x = data['pocket_x'][0]
        
        # Prepare pocket data in numpy
        pocket_coords_np = pocket_x.cpu().numpy() if torch.is_tensor(pocket_x) else np.asarray(pocket_x)
        pocket_types = data.get('pocket_types', [['C'] * int(pocket_coords_np.shape[0])])[0]
        pocket_resnames = data.get('pocket_resnames', [None])[0]
        
        # Get number of atoms from reference ligand in dataset
        reference_ligand_atoms = data.get('num_atoms', [30])[0]  # Default 30 if not available

        # Save pocket as PDB
        os.makedirs(os.path.join(sdf_dir, str(uuid)), exist_ok=True)
        pocket_pdb_path = os.path.join(sdf_dir, str(uuid), 'pocket.pdb')
        save_pocket_pdb(pocket_coords_np, pocket_types, pocket_pdb_path, 
                       pocket_resnames=pocket_resnames)

        lig_mol = None
        ligand_coords_np = None
        if args.pocket_detection != 'kvfinder':
            # Get ligand molecule
            lig_mol = data.get('ligand_mol', [None])[0] if 'ligand_mol' in data else None
            if lig_mol is not None:
                try:
                    conf = lig_mol.GetConformer()
                    ligand_coords_np = np.array([list(conf.GetAtomPosition(i)) for i in range(lig_mol.GetNumAtoms())], dtype=float)
                except Exception:
                    ligand_coords_np = None

        kvfinder_coords = None        
        # Cavity detection according to selected mode
        if pocket_coords_np.size > 0:
            mode_order = []
            if args.pocket_detection == 'auto':
                mode_order = ['ligand_proximity', 'kvfinder']
            else:
                mode_order = [args.pocket_detection]

            # Determine if interactive mode is requested
            is_interactive = any(mode.endswith('_interactive') for mode in mode_order)

            # Get cavity detection functions from shared module
            mode_to_fn = get_cavity_detection_functions(
                pocket_pdb_path=pocket_pdb_path,
                lig_mol=lig_mol,
                ligand_coords_np=ligand_coords_np,
                pocket_coords_np=pocket_coords_np,
                ligand_proximity_radius=args.ligand_proximity_radius,
                lig_pdb_path=os.path.join(sdf_dir, str(uuid), 'ligand_dataset.pdb') if lig_mol is not None else None,
                uuid=uuid,
                interactive=is_interactive
            )

            for mode in mode_order:
                # Map interactive modes to base modes
                base_mode = mode.replace('_interactive', '') if mode.endswith('_interactive') else mode
                kv_coords = mode_to_fn[base_mode]()
                if kv_coords is not None and len(kv_coords) > 0:
                    kvfinder_coords = kv_coords.tolist() if hasattr(kv_coords, 'tolist') else kv_coords
                    cavity_pdb_path = os.path.join(sdf_dir, str(uuid), 'cavity_points.pdb')
                    save_cavity_points_pdb(kvfinder_coords, cavity_pdb_path)
                    logger.info(f'Cavity points selected by mode "{mode}" saved to: {cavity_pdb_path}')
                    break
            if kvfinder_coords is None:
                logger.error(f'No cavity found with mode {args.pocket_detection} for UUID {uuid}. Skipping to next target.')
                continue

        # Build receptor_info with cavity_coords from pyKVFinder if available
        receptor_info = {
            'pocket_data': {
                'full_coord': pocket_coords_np.tolist() if hasattr(pocket_coords_np, 'tolist') else pocket_coords_np,
                'full_types': pocket_types,
                'full_resnames': pocket_resnames,
            },
            'cavity_coords': kvfinder_coords if kvfinder_coords is not None else None,
        }

        # Select ligand from dataset for score_based
        ligand_file = None
        if args.ref_atoms_method == 'score_based':
            # Get ligand molecule directly from dataset (already loaded during preprocess)
            # lig_mol = data.get('ligand_mol', [None])[0] if 'ligand_mol' in data else None
            if lig_mol is None:
                logger.warning(f'Ligand not found in dataset for UUID {uuid}; skipping score_based for this UUID')
                continue
            try:
                ligand_file = os.path.join(sdf_dir, str(uuid), 'ligand_dataset.sdf')
                Chem.MolToMolFile(lig_mol, ligand_file)
            except Exception as e:
                logger.warning(f'Could not prepare ligand from dataset for UUID {uuid}: {e}')
                continue

        
        # Build method parameters and select initial atoms
        # score_based method requires receptor_file for SMINA
        if args.ref_atoms_method == 'score_based':
            receptor_info['receptor_file'] = pocket_pdb_path
        method_params = build_method_params(args.ref_atoms_method, args, receptor_info)
        init_coords_list, init_node_types_list, method_info = select_initial_atoms(
            args.ref_atoms_method,
            receptor_info,
            ligand_file,
            train_config.chem.atomic_numbers,
            method_params,
            logger,
            log_dir
        )
        if not init_coords_list or not init_node_types_list:
            logger.error(f"Error in initial atom selection: {method_info.get('error', 'unknown')}")
            continue
        # Convert to tensors compatible with the rest of the flow
        init_coords = torch.tensor(init_coords_list, dtype=torch.float32)
        init_labels = torch.tensor(init_node_types_list, dtype=torch.long)

        # Re-center with respect to selected initial atoms
        mean = init_coords.mean(0)
        pocket_x = (pocket_x - mean).to(args.device)
        ref_positions = (init_coords - mean).to(args.device)
        # Add small offset if ref_positions is all zeros (happens when single atom is re-centered)
        if ref_positions.abs().sum() == 0:
            ref_positions[0, 0] = 1e-5
        ref_node_type = init_labels

        # Save reference atoms ONCE, outside the generation loop
        os.makedirs(os.path.join(sdf_dir, str(uuid)), exist_ok=True)
        
        # Convert tensors to lists for save_reference_atoms_from_coords
        init_coords_list = init_coords.cpu().numpy().tolist()
        init_labels_list = init_labels.cpu().tolist()
        
        # Save reference atoms using the same function as sample_pocket.py
        reference_atoms_file = os.path.join(sdf_dir, str(uuid), 'reference_atoms.sdf')
        sdf_path_arg = save_reference_atoms_from_coords(
            init_coords_list,
            init_labels_list,
            train_config.chem.atomic_numbers,
            out_file=reference_atoms_file,
            generation_info=f"Initial: {args.ref_atoms_method} method ({len(init_coords_list)} atoms)"
        )
        if sdf_path_arg:
            logger.info(f'Reference atoms saved to: {sdf_path_arg}')
        num_atoms_hint = ref_positions.shape[0]
        mean_size = max(num_atoms_hint + 10, reference_ligand_atoms)

        first_loop = True
        while len(pool.finished) < config.sample.num_mols:
            if len(pool.failed) > 3 * (config.sample.num_mols):
                logger.info('Too many failed molecules. Stop sampling.')
                break

            batch_size = args.batch_size if args.batch_size > 0 else config.sample.batch_size
            # Cap n_graphs to reduce memory
            n_graphs = min(batch_size, (config.sample.num_mols - len(pool.finished)) * 2, 10)
            batch_holder = make_data_placeholder(
                n_graphs=n_graphs, device=args.device, mean_size=mean_size
            )
            batch_node, halfedge_index, batch_halfedge = (
                batch_holder['batch_node'],
                batch_holder['halfedge_index'],
                batch_holder['batch_halfedge'],
            )

            # Ensure molecule size accommodates reference atoms
            if 'mol_size' in batch_holder and batch_holder['mol_size'] < ref_positions.shape[0]:
                logger.warning(f"Insufficient molecule size: {batch_holder['mol_size']} < {ref_positions.shape[0]}")
                logger.warning('Limiting reference atoms to available size')
                max_ref_atoms = max(0, batch_holder['mol_size'] - 5)
                if max_ref_atoms > 0:
                    ref_positions = ref_positions[:max_ref_atoms]
                    ref_node_type = ref_node_type[:max_ref_atoms]
                    logger.info(f'Reference atoms limited to {max_ref_atoms}')
                else:
                    logger.error('Molecule size too small for reference atoms')
                    break

            # Do not continue if there are no reference atoms
            if ref_positions.shape[0] == 0:
                logger.error('No reference atoms available after size adjustment')
                break

            ref_coords, ref_type_one_hot = prepare_pharmacophore_batch(
                ref_positions, ref_node_type, batch_node, args.device
            )
            ref_mask = (ref_coords.sum(dim=-1) != 0).to(args.device)

            try:
                outputs = model.sample_in_pocket(
                    n_graphs=n_graphs,
                    batch_node=batch_node,
                    halfedge_index=halfedge_index,
                    batch_halfedge=batch_halfedge,
                    bond_predictor=bond_predictor,
                    guidance=guidance,
                    ref_coords=ref_coords,
                    ref_mask=ref_mask,
                    ref_type=ref_type_one_hot,
                    log_dir=log_dir,
                    condition_clash=True,
                    pocket_x=pocket_x,
                    clash_rate=args.clash_rate,
                )
                outputs = {key: [v.cpu().numpy() for v in value] for key, value in outputs.items()}
            except torch.cuda.OutOfMemoryError:
                logger.error("GPU OOM during sample_in_pocket. Reducing n_graphs and molecule size...")
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                # Retry with n_graphs=1 and smaller mean_size
                n_graphs = 1
                mean_size = max(ref_positions.shape[0] + 5, 20)
                batch_holder = make_data_placeholder(
                    n_graphs=n_graphs, device=args.device, mean_size=mean_size
                )
                batch_node, halfedge_index, batch_halfedge = (
                    batch_holder['batch_node'],
                    batch_holder['halfedge_index'],
                    batch_holder['batch_halfedge'],
                )
                ref_coords, ref_type_one_hot = prepare_pharmacophore_batch(
                    ref_positions, ref_node_type, batch_node, args.device
                )
                ref_mask = (ref_coords.sum(dim=-1) != 0).to(args.device)
                try:
                    outputs = model.sample_in_pocket(
                        n_graphs=n_graphs,
                        batch_node=batch_node,
                        halfedge_index=halfedge_index,
                        batch_halfedge=batch_halfedge,
                        bond_predictor=bond_predictor,
                        guidance=guidance,
                        ref_coords=ref_coords,
                        ref_mask=ref_mask,
                        ref_type=ref_type_one_hot,
                        log_dir=log_dir,
                        condition_clash=True,
                        pocket_x=pocket_x,
                        clash_rate=args.clash_rate,
                    )
                    outputs = {key: [v.cpu().numpy() for v in value] for key, value in outputs.items()}
                except torch.cuda.OutOfMemoryError:
                    logger.error("GPU OOM even with n_graphs=1. Skipping this attempt.")
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                    continue

            # Decode
            batch_node_np, halfedge_index_np, batch_halfedge_np = (
                batch_node.cpu().numpy(),
                halfedge_index.cpu().numpy(),
                batch_halfedge.cpu().numpy(),
            )
            try:
                output_list = seperate_outputs(outputs, n_graphs, batch_node_np, halfedge_index_np, batch_halfedge_np)
            except Exception:
                continue

            gen_list = []

            for i_mol, output_mol in enumerate(output_list):
                mol_info = featurizer.decode_output(
                    pred_node=output_mol['pred'][0],
                    pred_pos=output_mol['pred'][1],
                    pred_halfedge=output_mol['pred'][2],
                    halfedge_index=output_mol['halfedge_index'],
                )

                try:
                    rdmol = reconstruct_from_generated_with_edges(mol_info, add_edge=add_edge)
                    mol_info_shifted = mol_info
                    mol_info_shifted['atom_pos'] = mol_info_shifted['atom_pos'] + mean.cpu().numpy()
                    rdmol_shifted = reconstruct_from_generated_with_edges(mol_info_shifted, add_edge=add_edge)
                except MolReconsError:
                    pool.failed.append(mol_info)
                    logger.warning('Reconstruction error encountered.')
                    continue
                mol_info['rdmol'] = rdmol
                mol_info['rdmol_shifted'] = rdmol_shifted
                smiles = Chem.MolToSmiles(rdmol)
                mol_info['smiles'] = smiles

                if '.' in smiles:
                    logger.warning('Incomplete molecule: %s' % smiles)
                    pool.failed.append(mol_info)
                else:
                    logger.info('Success: %s' % smiles)
                    gen_list.append(mol_info)

            # Save
            # sdf_dir already created globally; here only SMILES and generated molecules
            if first_loop:
                first_loop = False
            with open(os.path.join(log_dir, 'SMILES.txt'), 'a') as smiles_f:
                for i_f, data_finished in enumerate(gen_list):
                    smiles_f.write(data_finished['smiles'] + '\n')
                    base_name = f'{i_f + len(pool.finished)}'
                    rdmol_shifted = data_finished['rdmol_shifted']
                    path_shifted = os.path.join(sdf_dir, str(uuid), f'{base_name}_shifted.sdf')
                    Chem.MolToMolFile(rdmol_shifted, path_shifted)
                    # Save relative SDF path for spreadsheet
                    data_finished['sdf_filename'] = os.path.join(str(uuid), f'{base_name}_shifted.sdf')

            pool.finished.extend(gen_list)
            print_pool_status(pool, logger)

        # Accumulate results in global pool
        global_pool.finished.extend(pool.finished)
        global_pool.failed.extend(pool.failed)
        
        logger.info(f'Receptor {uuid} completed. Total accumulated: {len(global_pool.finished)} molecules generated')

    # Save final pool with all results
    torch.save(global_pool, os.path.join(log_dir, 'samples_all.pt'))


if __name__ == '__main__':
    main()