sample_pocket.py

"""
Main script for sampling molecules in a single protein pocket.

This script generates drug-like molecules using structure-based methods with
various initial atom selection strategies (pharmacophore locator, H-bond predictor,
score-based, or random).
"""

import argparse
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
from models.bond_predictor import BondPredictor
from models.model import MolDiff
from utils.misc import (
    EasyDict,
    get_logger,
    get_new_log_dir,
    load_config,
    seed_all
)
from utils.reconstruct import (
    MolReconsError,
    reconstruct_from_generated_with_edges
)
from utils.sample import seperate_outputs
from utils.transforms import FeaturizeMol, make_data_placeholder

# Local imports from PocketHotspot
from utils.cavity_detection import get_cavity_detection_functions
from utils.checkpoint_utils import resolve_checkpoint_path
from utils.io_utils import (
    save_cavity_points_pdb,
    save_pocket_pdb,
    save_reference_atoms_from_coords
)
from utils.pocket_sample_utils import create_dummy_molecule
from utils.prepare_pocket import get_full_receptor, get_pocket
from utils.sampling_utils import (
    build_method_params,
    prepare_pharmacophore_batch,
    print_pool_status,
    select_initial_atoms
)


def main():  

    parser = argparse.ArgumentParser(
        description="PocketHotspot - Structure-based drug design tool",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        epilog='''
Examples:
  # Basic usage with default settings:
  python sample_pocket.py --receptor data/example_5NGZ/5ngz_A_rec.pdb --ligand data/example_5NGZ/5ngz_A_rec_5ngz_2bg_lig_tt_min_0.sdf

  # Use pharmacophore_locator method:
  python sample_pocket.py --receptor receptor.pdb --ligand ligand.sdf --ref_atoms_method pharmacophore_locator

  # Use hbond_predictor with custom device:
  python sample_pocket.py --receptor receptor.pdb --ligand ligand.sdf --ref_atoms_method hbond_predictor --device cuda:0

  # Use score_based method with custom atom fraction:
  python sample_pocket.py --receptor receptor.pdb --ligand ligand.sdf --ref_atoms_method score_based --atom_fraction 0.3

  # Use ligand_proximity pocket detection:
  python sample_pocket.py --receptor receptor.pdb --ligand ligand.sdf --pocket_detection ligand_proximity --ligand_proximity_radius 5.0

For more information, use --help or -h to see all available options.
        '''
    )
    # Input files
    parser.add_argument(
        "--receptor",
        required=True,
        help="Receptor file (PDBQT or PDB)"
    )
    parser.add_argument(
        "--ligand",
        help="Ligand file (SDF or PDBQT or MOL2) - Required for score-based"
    )
    parser.add_argument(
        '--config',
        type=str,
        default='../MolSnapper/configs/sample/sample_MolDiff.yml',
        help='Model configuration file'
    )

    # 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=8,
        help='Batch size for generation'
    )
    parser.add_argument(
        '--mol_size',
        type=int,
        default=20,
        help='Target molecule size'
    )
    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='hbond_predictor',
        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 atoms with best affinity to select (e.g., 0.2 for 20 percent)'
    )

    # 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'
    )

    # 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_coords',
        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 configurations
    outdir_arg = os.path.dirname(args.receptor)
    config = load_config(args.config)
    seed_all(config.sample.seed + np.sum([ord(s) for s in outdir_arg]))

    # 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 training configuration
    ckpt = torch.load(config.model.checkpoint, map_location=args.device)
    train_config = ckpt['config']

    # Get receptor name (without extension) for log directory prefix
    receptor_name = os.path.basename(args.receptor)
    if '.' in receptor_name:
        receptor_name = receptor_name[:receptor_name.rfind('.')]

    # Configure logging
    log_root = outdir_arg.replace('outputs', 'outputs_vscode') if sys.argv[0].startswith('/data') else outdir_arg
    log_dir = get_new_log_dir(log_root, prefix=receptor_name, tag=args.ref_atoms_method)
    logger = get_logger('sample', log_dir)
    logger.info(args)
    logger.info(config)
    shutil.copyfile(args.config, os.path.join(log_dir, os.path.basename(args.config)))

    # Validate that ligand is provided when necessary
    lig_mol = None
    ligand_coords_np = None

    ligand_required_methods = ['score_based']
    ligand_required_detection = ['kvfinder_with_ligand', 'ligand_coords', 'ligand_proximity']

    requires_ligand = (
        (args.ref_atoms_method in ligand_required_methods) or
        (args.pocket_detection in ligand_required_detection)
    )

    if requires_ligand and not args.ligand:
        if args.ref_atoms_method in ligand_required_methods:
            print(f"ERROR: --ligand is required for ref_atoms_method '{args.ref_atoms_method}'")
        if args.pocket_detection in ligand_required_detection:
            print(f"ERROR: --ligand is required for pocket_detection '{args.pocket_detection}'")
        sys.exit(1)

    # Load ligand if provided
    if args.ligand:
        if not os.path.exists(args.ligand):
            logger.error(f"Ligand file not found: {args.ligand}")
            if requires_ligand:
                logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}'")
                sys.exit(1)
        else:
            try:
                lig_mol = Chem.SDMolSupplier(args.ligand)[0]
                if lig_mol is None:
                    logger.error(f"Failed to load ligand from file: {args.ligand}")
                    if requires_ligand:
                        logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}'")
                        sys.exit(1)
                else:
                    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 as e:
                        logger.error(f"Failed to extract coordinates from ligand: {e}")
                        if requires_ligand:
                            logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}'")
                            sys.exit(1)
            except Exception as e:
                logger.error(f"Error loading ligand file '{args.ligand}': {e}")
                if requires_ligand:
                    logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}'")
                    sys.exit(1)

    # Validate ligand is available when required
    if requires_ligand and (lig_mol is None or ligand_coords_np is None):
        logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}' but could not be loaded")
        logger.error("Please provide a valid ligand file with --ligand")
        sys.exit(1)

    receptor_info = get_full_receptor(args.receptor)

    # Validate ligand is available when required for pocket detection BEFORE attempting detection
    if args.ref_atoms_method != 'score_based':
        if args.pocket_detection in ligand_required_detection:
            if lig_mol is None:
                logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}' but could not be loaded")
                logger.error("Please provide a valid ligand file with --ligand")
                sys.exit(1)
            if ligand_coords_np is None or len(ligand_coords_np) == 0:
                logger.error(f"Ligand file '{args.ligand}' has no valid coordinates")
                logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}'")
                logger.error("Please provide a valid ligand file that matches the receptor")
                sys.exit(1)
            if lig_mol.GetNumAtoms() == 0:
                logger.error(f"Ligand file '{args.ligand}' contains no atoms")
                logger.error(f"Ligand is required for pocket_detection '{args.pocket_detection}'")
                sys.exit(1)
            # Quick validation: check if ligand is reasonably close to receptor
            # This helps detect cases where default ligand doesn't match the receptor
            receptor_coords = np.array(receptor_info['full_coord'])
            if len(receptor_coords) > 0 and len(ligand_coords_np) > 0:
                receptor_center = receptor_coords.mean(axis=0)
                ligand_center = ligand_coords_np.mean(axis=0)
                distance = np.linalg.norm(receptor_center - ligand_center)
                # If ligand is more than 50Å away from receptor center, it's likely wrong
                if distance > 50.0:
                    logger.error(f"Ligand file '{args.ligand}' appears to be too far from receptor (distance: {distance:.1f}Å)")
                    logger.error(f"This usually means the ligand file does not match the receptor '{args.receptor}'")
                    logger.error("Please provide a valid ligand file that corresponds to this receptor")
                    sys.exit(1)
    
    # Detect pocket according to selected mode (if not score_based)
    if args.ref_atoms_method != 'score_based':
        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
        lig_pdb_path_local = os.path.join(log_dir, 'ligand_for_kvfinder.pdb') if lig_mol is not None else None
        mode_to_fn = get_cavity_detection_functions(
            pocket_pdb_path=args.receptor,
            lig_mol=lig_mol,
            ligand_coords_np=ligand_coords_np,
            pocket_coords_np=np.array(receptor_info['full_coord']),
            ligand_proximity_radius=args.ligand_proximity_radius,
            lig_pdb_path=lig_pdb_path_local,
            uuid=None,
            interactive=is_interactive
        )

        kvfinder_coords = None

        for mode in mode_order:
            logger.info(f'Attempting to detect cavity using mode: {mode}')
            # 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(log_dir, 'cavity_points.pdb')
                save_cavity_points_pdb(kvfinder_coords, cavity_pdb_path)
                # Update receptor_info with cavity_coords
                receptor_info['cavity_coords'] = kvfinder_coords

                # Generate pocket using detected cavity or ligand coordinates
                pocket_data = None
                if mode == 'ligand_coords':
                    # Use ligand directly for pocket generation
                    if lig_mol is None:
                        logger.error("Mode 'ligand_coords' requires a valid ligand, but none was loaded")
                        logger.error("Please provide a valid ligand file with --ligand")
                        sys.exit(1)
                    receptor_info['skip_neighbor_check'] = True
                    mol_for_pocket = lig_mol
                    logger.info("Mode 'ligand_coords': generating pocket around ligand...")
                else:
                    # Create dummy molecule with cavity coordinates
                    mol_for_pocket = create_dummy_molecule(kvfinder_coords)
                    logger.info("Generating pocket using detected cavity coordinates...")

                try:
                    pocket_data = get_pocket(mol_for_pocket, args.receptor, receptor_info)
                except Exception as e:
                    logger.warning(f"Error generating pocket: {e}")

                if pocket_data is not None:
                    num_atoms = len(pocket_data['full_coord'])

                    # Validate that pocket has atoms, especially for ligand_coords mode
                    if num_atoms == 0:
                        if mode == 'ligand_coords':
                            logger.error(f"Pocket generated with 0 atoms using mode '{mode}'")
                            logger.error("This usually means the ligand file is invalid or does not match the receptor")
                            logger.error(f"Please check that --ligand points to a valid ligand file for receptor '{args.receptor}'")
                            sys.exit(1)
                        else:
                            logger.warning(f"Pocket generated with 0 atoms using mode '{mode}', trying next mode...")
                            continue

                    # Save pocket to PDB file (optional, for visualization)
                    pocket_pdb_path = os.path.join(log_dir, 'pocket.pdb')
                    save_pocket_pdb(
                        pocket_data['full_coord'],
                        pocket_data['full_types'],
                        pocket_pdb_path,
                        pocket_resnames=pocket_data.get('full_resnames', None)
                    )

                    # Update receptor_info with generated pocket (store only pocket_data)
                    receptor_info['pocket_data'] = pocket_data

                    logger.info(f"Pocket generated with {num_atoms} atoms using mode '{mode}'")
                    # Successfully found cavity and generated pocket, exit the loop
                    break
                else:
                    logger.warning(f"Could not generate pocket using mode '{mode}', trying next mode...")
                    # Continue to next mode if pocket generation failed
                    continue
            else:
                # No cavity detected with this mode, try next mode
                logger.warning(f"No cavity detected with mode '{mode}', trying next mode...")
                continue

        # If we get here, no mode succeeded
        if kvfinder_coords is None:
            logger.error(f'No cavity found with any mode in {mode_order}')
            logger.error(f'Method {args.ref_atoms_method} requires a detected pocket to work correctly')
            logger.error('Please verify that:')
            logger.error('  1. The receptor file is valid')
            if any(m in mode_order for m in ['kvfinder_with_ligand', 'ligand_coords', 'ligand_proximity']):
                logger.error('  2. The ligand file is available and valid')
            if any(m in mode_order for m in ['kvfinder', 'kvfinder_with_ligand', 'kvfinder_interactive']):
                logger.error('  3. pyKVFinder is installed and working correctly')
            logger.error('  4. Consider trying another detection mode with --pocket_detection')
            sys.exit(1)
        else:
            # Cavity was detected but pocket could not be generated with any mode
            if receptor_info.get('pocket_data') is None:
                logger.error('Cavity was detected but pocket could not be generated with any mode')
                sys.exit(1)
    else:
        receptor_info['receptor_file'] = args.receptor

    # Create reference atom selection method parameters
    method_params = build_method_params(args.ref_atoms_method, args, receptor_info)

    ligand_file = args.ligand if args.ligand else None
    initial_coords, initial_node_types, 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 initial_coords or not initial_node_types:
        logger.error(f"Error in initial atom selection: {method_info.get('error', 'Unknown error')}")
        sys.exit(1)

    # Save reference atoms
    save_reference_atoms_from_coords(
        initial_coords,
        initial_node_types,
        train_config.chem.atomic_numbers,
        out_file=os.path.join(log_dir, 'reference_atoms.sdf'),
        generation_info=f"Initial: {args.ref_atoms_method} method ({len(initial_coords)} atoms)"
    )

    # Configure transformer
    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)

    # Load 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()

    # Configure 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

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

    # Initialize molecule pool
    pool = EasyDict({
        'failed': [],
        'finished': [],
    })

    # Global counter for SDF files
    global_sdf_counter = 0

    # Configure reference atoms
    ref_node_type = torch.tensor(initial_node_types, dtype=torch.long)
    ref_positions = torch.tensor(initial_coords, dtype=torch.float32)
    # Get pocket coordinates from receptor_info
    if receptor_info.get('pocket_data') is not None:
        pocket_coords = receptor_info['pocket_data']['full_coord']
    else:
        pocket_coords = receptor_info['full_coord']
    pocket_x = torch.tensor(pocket_coords, dtype=torch.float32)

    mean = ref_positions.mean(0)
    if ref_positions.shape[0] == 1:
        mean = mean + torch.randn_like(mean)

    pocket_x = (pocket_x - mean).to(args.device)
    ref_positions = (ref_positions - mean).to(args.device)

    # Initialize variables for tracking reference atoms
    logger.info(f'Reference atoms: {len(initial_coords)} atoms')
    logger.info('Using input reference atoms for generation')
    
    # Main generation loop
    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

        # Prepare batch
        batch_size = args.batch_size if args.batch_size > 0 else config.sample.batch_size
        remaining_mols = max(0, config.sample.num_mols - len(pool.finished))
        if remaining_mols == 0:
            logger.info('All required molecules already generated.')
            break

        # Reduce batch size to use less memory
        n_graphs = min(batch_size, remaining_mols * 2, 10)
        n_graphs = max(1, n_graphs)

        # Ensure molecule size is sufficient for reference atoms
        base_mol_size = args.mol_size if args.mol_size is not None else 50
        min_mol_size = max(base_mol_size, ref_positions.shape[0] + 10)
        logger.info(f'Generating {n_graphs} molecules in this batch (minimum size: {min_mol_size})')

        batch_holder = make_data_placeholder(n_graphs=n_graphs, device=args.device, mean_size=min_mol_size)
        batch_node = batch_holder['batch_node']
        halfedge_index = batch_holder['halfedge_index']
        batch_halfedge = batch_holder['batch_halfedge']

        # Verify that size is sufficient
        if 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 = 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

        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)

        # Generate molecules
        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 memory error. Trying with smaller batch...")
            torch.cuda.empty_cache()
            n_graphs = 1
            batch_holder = make_data_placeholder(n_graphs=n_graphs, device=args.device, mean_size=args.mol_size)
            batch_node = batch_holder['batch_node']
            halfedge_index = batch_holder['halfedge_index']
            batch_halfedge = 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)

            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()}

        # Decode outputs to molecules
        batch_node = batch_node.cpu().numpy()
        halfedge_index = halfedge_index.cpu().numpy()
        batch_halfedge = batch_halfedge.cpu().numpy()
        try:
            output_list = seperate_outputs(outputs, n_graphs, batch_node, halfedge_index, batch_halfedge)
        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:  # Pass checks!
                logger.info('Success: %s' % smiles)
                gen_list.append(mol_info)

        # Save SDF molecules
        sdf_dir = log_dir + '_SDF'
        os.makedirs(sdf_dir, exist_ok=True)

        with open(os.path.join(log_dir, 'SMILES.txt'), 'a') as smiles_f:
            for i, data_finished in enumerate(gen_list):
                smiles_f.write(data_finished['smiles'] + '\n')
                rdmol_shifted = data_finished['rdmol_shifted']
                # Use global counter for SDF files
                sdf_filename = f"{global_sdf_counter}_shifted.sdf"
                Chem.MolToMolFile(rdmol_shifted, os.path.join(sdf_dir, sdf_filename))
                # Add SDF file information to molecule
                data_finished['sdf_filename'] = sdf_filename
                global_sdf_counter += 1

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

        # Free GPU memory after each batch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    # Save final information
    torch.save(pool, os.path.join(log_dir, 'samples_all.pt'))


if __name__ == "__main__":
    main()