benchmark_formats.py

"""
Comprehensive benchmark comparing MolBox vs SDF vs OEB formats.

Tests all input/output combinations:
- MolBox with RDKit molecules
- MolBox with OpenEye molecules
- SDF with RDKit molecules (all conformers)
- SDF with OpenEye molecules (all conformers)
- OEB with OpenEye molecules (all conformers)

Measures: write time, read time, file size, conformer preservation
"""

import os
import sys
import time
import tempfile
from pathlib import Path
import pandas as pd
from rdkit import Chem
from rdkit.Chem import AllChem

from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing


sys.path.insert(0, str(Path(__file__).parent))
from molbox.database import MoleculeDatabase

try:
    from openeye import oechem
    from openeye.oeomega import OEOmega
    HAS_OPENEYE = True
except ImportError:
    HAS_OPENEYE = False
    print("⚠ Warning: OpenEye not available - skipping OEB tests")
from rdkit import Chem

def write_sdf_with_all_conformers_rdkit(molecules, filepath):
    """
    Write RDKit molecules with ALL conformers to SDF.
    Each conformer is written as a separate SD record with metadata to group them.
    """
    writer = Chem.SDWriter(filepath)
    writer.SetForceV3000(False)

    for mol_idx, mol in enumerate(molecules):
        # Get or create a molecule name for grouping
        if mol.HasProp("_Name"):
            mol_name = mol.GetProp("_Name")
        else:
            mol_name = f"mol_{mol_idx}"
            mol.SetProp("_Name", mol_name)
        
        nconf = mol.GetNumConformers()
        if nconf == 0:
            # Write single record without conformer ID
            writer.write(mol)
        else:
            # Write each conformer as a separate record with grouping metadata
            for conf_id in range(nconf):
                # Add conformer metadata to help with reconstruction
                mol.SetProp("_ConfId", str(conf_id))
                mol.SetProp("_MolName", mol_name)  # Explicit grouping key
                writer.write(mol, confId=conf_id)
    
    writer.close()


def read_sdf_with_all_conformers_rdkit(filepath, verbose=False):
    """
    Read an SDF that was written with write_sdf_with_all_conformers_rdkit
    and reconstruct RDKit molecules with all conformers grouped by _MolName.
    """
    suppl = Chem.SDMolSupplier(filepath, removeHs=False)
    mol_dict = {}
    problems = 0
    records = 0

    for mol in suppl:
        records += 1
        if mol is None:
            if verbose:
                print(f"[WARN] Record {records}: RDKit returned None (bad record).")
            problems += 1
            continue

        # Get the grouping key - use _MolName if available, fallback to _Name
        mol_name = None
        if mol.HasProp("_MolName"):
            mol_name = mol.GetProp("_MolName")
        elif mol.HasProp("_Name"):
            mol_name = mol.GetProp("_Name")
        else:
            # Fallback: treat as new molecule
            mol_name = f"mol_record_{records}"

        if mol_name not in mol_dict:
            # First conformer for this molecule - store it
            mol_dict[mol_name] = mol
        else:
            # Additional conformer - add to existing molecule
            existing_mol = mol_dict[mol_name]

            try:
                conf = mol.GetConformer()
            except Exception as e:
                if verbose:
                    print(f"[WARN] Could not get conformer for record {records} name={mol_name}: {e}")
                problems += 1
                continue

            # Verify atom count matches
            if conf.GetNumAtoms() != existing_mol.GetNumAtoms():
                problems += 1
                if verbose:
                    print(f"[WARN] Atom count mismatch for '{mol_name}': "
                          f"conf={conf.GetNumAtoms()} != mol={existing_mol.GetNumAtoms()}. "
                          f"Skipping conformer.")
                continue

            # Copy conformer and add to existing molecule
            conf_copy = Chem.Conformer(conf)
            existing_mol.AddConformer(conf_copy, assignId=True)

    if verbose:
        total_confs = sum(m.GetNumConformers() for m in mol_dict.values())
        print(f"[INFO] Read {records} records; assembled {len(mol_dict)} molecules "
              f"with {total_confs} total conformers; {problems} problems.")

    return list(mol_dict.values())


def write_sdf_with_all_conformers_openeye(molecules, filepath):
    """
    Write OpenEye molecules with all conformers to SDF.
    OpenEye naturally handles multi-conformer molecules in SDF.
    """
    ofs = oechem.oemolostream(filepath)
    ofs.SetFormat(oechem.OEFormat_SDF)
    # Enable multi-conformer writing
    ofs.SetFlavor(oechem.OEFormat_SDF, oechem.OEOFlavor_SDF_Default)

    for mol in molecules:
        oechem.OEWriteMolecule(ofs, mol)

    ofs.close()


def read_sdf_with_all_conformers_openeye(filepath):
    """
    Read SDF with all conformers using OpenEye.
    Configure to read all conformers into each molecule object.
    """
    molecules = []
    
    ifs = oechem.oemolistream(filepath)
    ifs.SetFormat(oechem.OEFormat_SDF)
    # Set flavor to read all conformers
    ifs.SetConfTest(oechem.OEAbsoluteConfTest())
    
    for mol in ifs.GetOEMols():
        molecules.append(oechem.OEMol(mol))
    
    ifs.close()
    
    return molecules


def write_oeb_openeye(molecules, filepath):
    """Write OpenEye molecules to OEB format."""
    ofs = oechem.oemolostream(filepath)

    for mol in molecules:
        oechem.OEWriteMolecule(ofs, mol)

    ofs.close()


def read_oeb_openeye(filepath):
    """Read OpenEye molecules from OEB format."""
    molecules = []

    ifs = oechem.oemolistream(filepath)

    for mol in ifs.GetOEMols():
        molecules.append(oechem.OEMol(mol))

    ifs.close()

    return molecules


def load_smiles_from_csv(csv_path, smiles_column='SMILES', n_molecules=None):
    """Load SMILES from CSV file."""
    print(f"\n[SETUP] Loading SMILES from {csv_path}...")

    df = pd.read_csv(csv_path)

    if smiles_column not in df.columns:
        raise ValueError(f"Column '{smiles_column}' not found. Available: {list(df.columns)}")

    smiles_list = df[smiles_column].tolist()

    # Limit to n_molecules if specified
    if n_molecules is not None:
        smiles_list = smiles_list[:n_molecules]

    print(f"  Loaded {len(smiles_list)} SMILES from CSV")

    return smiles_list

def _generate_mol_with_confs(smi, n_conformers, idx):
    """Helper function for multiprocessing — generates a single molecule with conformers."""
    try:
        mol = Chem.MolFromSmiles(smi)
        if mol is None:
            return None

        mol = Chem.AddHs(mol)
        mol.SetProp("_Name", f"rdkit_mol_{idx}")

        AllChem.EmbedMultipleConfs(
            mol,
            numConfs=n_conformers,
            randomSeed=42,
            useRandomCoords=True,
            clearConfs=True
        )

        if mol.GetNumConformers() == 0:
            return None
        return mol

    except Exception:
        return None


def generate_rdkit_molecules_with_conformers_parallel(smiles_list, n_conformers=200, n_jobs=None):
    """Generate RDKit molecules with conformers from SMILES in parallel."""
    if n_jobs is None:
        n_jobs = max(1, multiprocessing.cpu_count() - 1)

    print(f"\n[SETUP] Generating RDKit molecules with {n_conformers} conformers each using {n_jobs} cores...")

    molecules = []
    failed = 0

    with ProcessPoolExecutor(max_workers=n_jobs) as executor:
        futures = {
            executor.submit(_generate_mol_with_confs, smi, n_conformers, i): i
            for i, smi in enumerate(smiles_list)
        }

        for future in as_completed(futures):
            mol = future.result()
            if mol is not None:
                molecules.append(mol)
            else:
                failed += 1

    total_confs = sum(mol.GetNumConformers() for mol in molecules)
    print(f"  Generated {len(molecules)} molecules with {total_confs} total conformers")
    if len(molecules) > 0:
        print(f"  Average: {total_confs/len(molecules):.1f} conformers/molecule")
    print(f"  Failed: {failed} molecules")

    return molecules


def generate_openeye_molecules_with_conformers(smiles_list, n_conformers=200):
    """Generate OpenEye molecules with conformers using Omega."""
    if not HAS_OPENEYE:
        return []

    print(f"\n[SETUP] Generating OpenEye molecules with Omega (~{n_conformers} conformers each)...")

    omega = OEOmega()
    omega.SetMaxConfs(n_conformers)
    omega.SetStrictStereo(False)

    molecules = []
    failed = 0

    for i, smi in enumerate(smiles_list):
        mol = oechem.OEMol()

        if oechem.OESmilesToMol(mol, smi):
            mol.SetTitle(f"oe_mol_{i}")

            if omega(mol):
                molecules.append(oechem.OEMol(mol))
            else:
                failed += 1
        else:
            failed += 1

    total_confs = sum(mol.NumConfs() for mol in molecules)
    print(f"  Generated {len(molecules)} molecules with {total_confs} total conformers")
    print(f"  Average: {total_confs/len(molecules):.1f} conformers/molecule")
    if failed > 0:
        print(f"  Failed: {failed} molecules")

    return molecules


def benchmark_format(name, molecules, write_func, read_func, filepath, verify_conformers=True):
    """Benchmark a specific format."""
    results = {}

    # Write benchmark
    start = time.time()
    write_func(molecules, filepath)
    write_time = time.time() - start

    # File size
    file_size_mb = os.path.getsize(filepath) / (1024 * 1024)

    # Read benchmark
    start = time.time()
    loaded_molecules = read_func(filepath)
    read_time = time.time() - start

    # Verify
    n_mols = len(loaded_molecules)

    if verify_conformers:
        # Count conformers
        if hasattr(loaded_molecules[0], 'GetNumConformers'):
            # RDKit
            total_confs = sum(mol.GetNumConformers() for mol in loaded_molecules)
        else:
            # OpenEye
            total_confs = sum(mol.NumConfs() for mol in loaded_molecules)

        avg_confs = total_confs / n_mols if n_mols > 0 else 0
    else:
        total_confs = 0
        avg_confs = 0

    results = {
        'format': name,
        'n_molecules': n_mols,
        'total_conformers': total_confs,
        'avg_conformers': avg_confs,
        'write_time_s': write_time,
        'read_time_s': read_time,
        'file_size_mb': file_size_mb,
        'write_rate_mol_s': n_mols / write_time if write_time > 0 else 0,
        'read_rate_mol_s': n_mols / read_time if read_time > 0 else 0,
    }

    return results


def main():
    print("=" * 80)
    print("MOLBOX FORMAT BENCHMARK - REAL DRUG-LIKE MOLECULES")
    print("=" * 80)

    # Configuration
    # TODO: Replace with your own CSV file containing SMILES data
    CSV_PATH = "path/to/your/molecules.csv"
    SMILES_COLUMN = "SMILES"
    N_MOLECULES = 3000  # Subset for benchmark
    N_CONFORMERS = 100

    # Load SMILES from CSV
    smiles_list = load_smiles_from_csv(CSV_PATH, SMILES_COLUMN, N_MOLECULES)

    # Generate test data with real molecules
    rdkit_mols = generate_rdkit_molecules_with_conformers_parallel(smiles_list, N_CONFORMERS, n_jobs=16)

    if HAS_OPENEYE:
        oe_mols = generate_openeye_molecules_with_conformers(smiles_list, N_CONFORMERS)
    else:
        oe_mols = []

    all_results = []

    print("\n" + "=" * 80)
    print("BENCHMARKING FORMATS")
    print("=" * 80)

    with tempfile.TemporaryDirectory() as tmpdir:

        # ─────────────────────────────────────────────────────────────
        # Test 1: MolBox (RDKit molecules)
        # ─────────────────────────────────────────────────────────────
        print("\n[1/5] MolBox (RDKit molecules)...")
        rdkit_molbox_path = os.path.join(tmpdir, "molbox_rdkit.parquet")
        result = benchmark_format(
            name="MolBox (RDKit)",
            molecules=rdkit_mols,
            write_func=lambda mols, path: MoleculeDatabase.save_molecules(mols, path, show_progress=False),
            read_func=lambda path: MoleculeDatabase.load_molecules(path, show_progress=False),
            filepath=rdkit_molbox_path
        )
        all_results.append(result)
        print(f"  ✓ Write: {result['write_time_s']:.2f}s ({result['write_rate_mol_s']:.0f} mol/s)")
        print(f"  ✓ Read:  {result['read_time_s']:.2f}s ({result['read_rate_mol_s']:.0f} mol/s)")
        print(f"  ✓ Size:  {result['file_size_mb']:.2f} MB")
        print(f"  ✓ Conformers preserved: {result['total_conformers']:.0f} ({result['avg_conformers']:.1f} avg)")

        # ─────────────────────────────────────────────────────────────
        # Test 2: SDF (RDKit molecules, all conformers)
        # ─────────────────────────────────────────────────────────────
        print("\n[2/5] SDF (RDKit molecules, all conformers)...")
        result = benchmark_format(
            name="SDF (RDKit)",
            molecules=rdkit_mols,
            write_func=write_sdf_with_all_conformers_rdkit,
            read_func=read_sdf_with_all_conformers_rdkit,
            filepath=os.path.join(tmpdir, "rdkit.sdf")
        )
        all_results.append(result)
        print(f"  ✓ Write: {result['write_time_s']:.2f}s ({result['write_rate_mol_s']:.0f} mol/s)")
        print(f"  ✓ Read:  {result['read_time_s']:.2f}s ({result['read_rate_mol_s']:.0f} mol/s)")
        print(f"  ✓ Size:  {result['file_size_mb']:.2f} MB")
        print(f"  ✓ Conformers preserved: {result['total_conformers']:.0f} ({result['avg_conformers']:.1f} avg)")

        # ─────────────────────────────────────────────────────────────
        # Tests 3–5: OpenEye benchmarks (if available)
        # ─────────────────────────────────────────────────────────────
        if HAS_OPENEYE:
            print("\n[3/5] MolBox (OpenEye molecules)...")
            oe_molbox_path = os.path.join(tmpdir, "molbox_oe.parquet")
            result = benchmark_format(
                name="MolBox (OpenEye)",
                molecules=oe_mols,
                write_func=lambda mols, path: MoleculeDatabase.save_molecules(mols, path, show_progress=False),
                read_func=lambda path: MoleculeDatabase.load_molecules(path, out_type='openeye', show_progress=False),
                filepath=oe_molbox_path
            )
            all_results.append(result)
            print(f"  ✓ Write: {result['write_time_s']:.2f}s ({result['write_rate_mol_s']:.0f} mol/s)")
            print(f"  ✓ Read:  {result['read_time_s']:.2f}s ({result['read_rate_mol_s']:.0f} mol/s)")
            print(f"  ✓ Size:  {result['file_size_mb']:.2f} MB")
            print(f"  ✓ Conformers preserved: {result['total_conformers']:.0f} ({result['avg_conformers']:.1f} avg)")

            print("\n[4/5] SDF (OpenEye molecules, all conformers)...")
            result = benchmark_format(
                name="SDF (OpenEye)",
                molecules=oe_mols,
                write_func=write_sdf_with_all_conformers_openeye,
                read_func=read_sdf_with_all_conformers_openeye,
                filepath=os.path.join(tmpdir, "oe.sdf")
            )
            all_results.append(result)
            print(f"  ✓ Write: {result['write_time_s']:.2f}s ({result['write_rate_mol_s']:.0f} mol/s)")
            print(f"  ✓ Read:  {result['read_time_s']:.2f}s ({result['read_rate_mol_s']:.0f} mol/s)")
            print(f"  ✓ Size:  {result['file_size_mb']:.2f} MB")
            print(f"  ✓ Conformers preserved: {result['total_conformers']:.0f} ({result['avg_conformers']:.1f} avg)")

            print("\n[5/5] OEB (OpenEye molecules)...")
            result = benchmark_format(
                name="OEB (OpenEye)",
                molecules=oe_mols,
                write_func=write_oeb_openeye,
                read_func=read_oeb_openeye,
                filepath=os.path.join(tmpdir, "oe.oeb")
            )
            all_results.append(result)
            print(f"  ✓ Write: {result['write_time_s']:.2f}s ({result['write_rate_mol_s']:.0f} mol/s)")
            print(f"  ✓ Read:  {result['read_time_s']:.2f}s ({result['read_rate_mol_s']:.0f} mol/s)")
            print(f"  ✓ Size:  {result['file_size_mb']:.2f} MB")
            print(f"  ✓ Conformers preserved: {result['total_conformers']:.0f} ({result['avg_conformers']:.1f} avg)")

            # ─────────────────────────────────────────────────────────────
            # Test 6: Cross-load OpenEye molecules from RDKit MolBox file
            # ─────────────────────────────────────────────────────────────
            print("\n[6/7] MolBox (RDKit → OpenEye cross-load)...")
            result = benchmark_format(
                name="MolBox (RDKit→OpenEye)",
                molecules=rdkit_mols,
                write_func=lambda mols, path: MoleculeDatabase.save_molecules(mols, path, show_progress=False),
                read_func=lambda path: MoleculeDatabase.load_molecules(path, out_type='openeye', show_progress=False),
                filepath=os.path.join(tmpdir, "molbox_rdkit_to_oe.parquet")
            )
            all_results.append(result)
            print(f"  ✓ Cross-load: {result['total_conformers']:.0f} confs ({result['avg_conformers']:.1f} avg)")

            # ─────────────────────────────────────────────────────────────
            # Test 7: Cross-load RDKit molecules from OpenEye MolBox file
            # ─────────────────────────────────────────────────────────────
            print("\n[7/7] MolBox (OpenEye → RDKit cross-load)...")
            result = benchmark_format(
                name="MolBox (OpenEye→RDKit)",
                molecules=oe_mols,
                write_func=lambda mols, path: MoleculeDatabase.save_molecules(mols, path, show_progress=False),
                read_func=lambda path: MoleculeDatabase.load_molecules(path, out_type='rdkit', show_progress=False),
                filepath=os.path.join(tmpdir, "molbox_oe_to_rdkit.parquet")
            )
            all_results.append(result)
            print(f"  ✓ Cross-load: {result['total_conformers']:.0f} confs ({result['avg_conformers']:.1f} avg)")

        else:
            print("\n[3-7] SKIPPED: OpenEye not available")


    # Create results table
    print("\n" + "=" * 80)
    print("RESULTS SUMMARY")
    print("=" * 80)

    df = pd.DataFrame(all_results)

    # Format for display
    print("\nPerformance Comparison:")
    print("-" * 80)
    display_df = df[['format', 'write_rate_mol_s', 'read_rate_mol_s', 'file_size_mb', 'avg_conformers']].copy()
    display_df.columns = ['Format', 'Write (mol/s)', 'Read (mol/s)', 'Size (MB)', 'Avg Confs']
    print(display_df.to_string(index=False))

    print("\n" + "=" * 80)
    print("KEY FINDINGS")
    print("=" * 80)

    # Calculate speedups vs SDF
    if len(all_results) >= 2:
        molbox_rdkit = all_results[0]
        sdf_rdkit = all_results[1]

        write_speedup = molbox_rdkit['write_rate_mol_s'] / sdf_rdkit['write_rate_mol_s']
        read_speedup = molbox_rdkit['read_rate_mol_s'] / sdf_rdkit['read_rate_mol_s']
        size_ratio = molbox_rdkit['file_size_mb'] / sdf_rdkit['file_size_mb']

        print(f"\nMolBox vs SDF (RDKit):")
        print(f"  Write speed: {write_speedup:.1f}x faster")
        print(f"  Read speed:  {read_speedup:.1f}x faster")
        print(f"  File size:   {size_ratio:.2f}x (smaller is better)")
        print(f"  Conformers:  All preserved ✓")

    if HAS_OPENEYE and len(all_results) >= 5:
        molbox_oe = all_results[2]
        oeb_oe = all_results[4]

        write_speedup = molbox_oe['write_rate_mol_s'] / oeb_oe['write_rate_mol_s']
        read_speedup = molbox_oe['read_rate_mol_s'] / oeb_oe['read_rate_mol_s']
        size_ratio = molbox_oe['file_size_mb'] / oeb_oe['file_size_mb']

        print(f"\nMolBox vs OEB (OpenEye):")
        print(f"  Write speed: {write_speedup:.1f}x {'faster' if write_speedup > 1 else 'slower'}")
        print(f"  Read speed:  {read_speedup:.1f}x {'faster' if read_speedup > 1 else 'slower'}")
        print(f"  File size:   {size_ratio:.2f}x (smaller is better)")
        print(f"  Conformers:  All preserved ✓")

    print("\n" + "=" * 80)


if __name__ == "__main__":
    main()