{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Virtual Screening Library Preparation\n\nThis notebook demonstrates how to prepare compound libraries for virtual screening campaigns using Auto3D.\n\n## CLI Quick Start (Recommended)\n\nFor most virtual screening workflows, the CLI is the fastest approach:\n\n```bash\n# Basic virtual screening - generate 5 conformers per molecule\nauto3d run compounds.smi --k=5 --gpu\n\n# With energy window selection\nauto3d run compounds.smi --k=5 --window=3.0 --gpu\n\n# Fast screening with ANI2xt\nauto3d run compounds.smi --k=3 --engine=ANI2xt --gpu\n\n# Large library with multi-GPU\nauto3d run large_library.smi --k=5 --gpu --gpu-idx=\"0,1,2,3\"\n\n# Validate input before processing\nauto3d validate compounds.smi\n\n# Use configuration preset for screening\nauto3d config init -p quick -o screening.yaml\nauto3d run compounds.smi -c screening.yaml --gpu\n```\n\nExample `screening.yaml`:\n```yaml\nk: 5\nwindow: 3.0\noptimizing_engine: ANI2xt\nthreshold: 0.3\nmax_confs: 100\nuse_gpu: true\n```\n\n---\n\n## Python Workflow\n\nThis notebook covers:\n\n1. **Library preprocessing** - SMILES validation, standardization, filtering\n2. **3D conformer generation** - optimal settings for docking\n3. **Property-based filtering** - Lipinski's Rule of Five, lead-likeness\n4. **Output formats** - docking-ready SDF files with proper metadata\n\n### Chemistry Background\n\nVirtual screening requires high-quality 3D conformers that represent the bioactive conformation. Key considerations:\n\n- **Conformational sampling**: Drug-like molecules typically have 1-10 rotatable bonds, generating 10s-1000s of conformers\n- **Energy window**: Bioactive conformations are usually within 3-5 kcal/mol of the global minimum\n- **Protonation state**: Most docking programs expect neutral or physiological pH forms"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "from pathlib import Path\n",
    "\n",
    "import Auto3D\n",
    "from Auto3D import Auto3DOptions, main\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import Descriptors, Lipinski, AllChem, PandasTools\n",
    "from rdkit.Chem.FilterCatalog import FilterCatalog, FilterCatalogParams\n",
    "import pandas as pd\n",
    "\n",
    "print(f\"Auto3D version: {Auto3D.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Library Preprocessing\n",
    "\n",
    "Before 3D generation, we validate and standardize the input SMILES. This catches common issues:\n",
    "\n",
    "- Invalid SMILES syntax\n",
    "- Disconnected fragments (salts, counterions)\n",
    "- Non-drug-like molecules (too large, reactive groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example compound library - common drug scaffolds\n",
    "# In practice, this would come from vendor catalogs or in-house databases\n",
    "compound_library = {\n",
    "    \"aspirin\": \"CC(=O)OC1=CC=CC=C1C(=O)O\",\n",
    "    \"ibuprofen\": \"CC(C)CC1=CC=C(C=C1)C(C)C(=O)O\",\n",
    "    \"acetaminophen\": \"CC(=O)NC1=CC=C(C=C1)O\",\n",
    "    \"caffeine\": \"CN1C=NC2=C1C(=O)N(C(=O)N2C)C\",\n",
    "    \"naproxen\": \"COC1=CC2=CC(C(C)C(=O)O)=CC=C2C=C1\",\n",
    "    \"celecoxib\": \"CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F\",\n",
    "    \"omeprazole\": \"COC1=CC2=NC(CS(=O)C3=NC4=CC=CC=C4N3)=NC2=C(OC)C1\",\n",
    "    \"atorvastatin\": \"CC(C)C1=C(C(=C(N1CCC(CC(CC(=O)O)O)O)C2=CC=C(C=C2)F)C3=CC=CC=C3)C(=O)NC4=CC=CC=C4\",\n",
    "    # Invalid/problematic entries for demonstration\n",
    "    \"invalid_smiles\": \"not_a_smiles\",\n",
    "    \"salt_form\": \"CC(=O)OC1=CC=CC=C1C(=O)[O-].[Na+]\",  # Sodium salt\n",
    "    \"too_large\": \"C\" * 100,  # Very long alkane\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate_and_standardize(smiles_dict, max_mw=600, max_rotatable=12):\n",
    "    \"\"\"\n",
    "    Validate and standardize a SMILES library for virtual screening.\n",
    "    \n",
    "    Filters applied:\n",
    "    - Valid SMILES parsing\n",
    "    - Remove salts/counterions (keep largest fragment)\n",
    "    - Molecular weight cutoff\n",
    "    - Rotatable bond limit (conformational complexity)\n",
    "    \n",
    "    Returns:\n",
    "        dict: Validated {name: canonical_smiles}\n",
    "        list: Rejected compounds with reasons\n",
    "    \"\"\"\n",
    "    valid = {}\n",
    "    rejected = []\n",
    "    \n",
    "    for name, smi in smiles_dict.items():\n",
    "        # Try to parse\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        if mol is None:\n",
    "            rejected.append((name, \"Invalid SMILES\"))\n",
    "            continue\n",
    "        \n",
    "        # Handle salts - keep largest fragment\n",
    "        fragments = Chem.GetMolFrags(mol, asMols=True, sanitizeFrags=True)\n",
    "        if len(fragments) > 1:\n",
    "            # Keep the largest fragment by heavy atom count\n",
    "            mol = max(fragments, key=lambda m: m.GetNumHeavyAtoms())\n",
    "        \n",
    "        # Check molecular weight\n",
    "        mw = Descriptors.MolWt(mol)\n",
    "        if mw > max_mw:\n",
    "            rejected.append((name, f\"MW {mw:.1f} > {max_mw}\"))\n",
    "            continue\n",
    "        \n",
    "        # Check rotatable bonds (conformational complexity)\n",
    "        n_rot = Lipinski.NumRotatableBonds(mol)\n",
    "        if n_rot > max_rotatable:\n",
    "            rejected.append((name, f\"Rotatable bonds {n_rot} > {max_rotatable}\"))\n",
    "            continue\n",
    "        \n",
    "        # Store canonical SMILES\n",
    "        valid[name] = Chem.MolToSmiles(mol, canonical=True)\n",
    "    \n",
    "    return valid, rejected\n",
    "\n",
    "\n",
    "valid_compounds, rejected = validate_and_standardize(compound_library)\n",
    "\n",
    "print(f\"Valid compounds: {len(valid_compounds)}\")\n",
    "print(f\"Rejected: {len(rejected)}\")\n",
    "print(\"\\nRejection reasons:\")\n",
    "for name, reason in rejected:\n",
    "    print(f\"  {name}: {reason}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Lipinski's Rule of Five Filtering\n",
    "\n",
    "For oral drug candidates, Lipinski's Rule of Five provides a useful filter:\n",
    "\n",
    "- **MW ≤ 500** - Molecular weight\n",
    "- **LogP ≤ 5** - Lipophilicity  \n",
    "- **HBD ≤ 5** - Hydrogen bond donors\n",
    "- **HBA ≤ 10** - Hydrogen bond acceptors\n",
    "\n",
    "Compounds violating ≥2 rules have poor oral bioavailability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_lipinski_properties(smiles_dict):\n",
    "    \"\"\"\n",
    "    Calculate Lipinski Rule of Five properties.\n",
    "    \n",
    "    Returns DataFrame with:\n",
    "    - MW, LogP, HBD, HBA\n",
    "    - Number of violations\n",
    "    - Pass/Fail status\n",
    "    \"\"\"\n",
    "    data = []\n",
    "    \n",
    "    for name, smi in smiles_dict.items():\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        \n",
    "        mw = Descriptors.MolWt(mol)\n",
    "        logp = Descriptors.MolLogP(mol)\n",
    "        hbd = Lipinski.NumHDonors(mol)\n",
    "        hba = Lipinski.NumHAcceptors(mol)\n",
    "        \n",
    "        # Count violations\n",
    "        violations = sum([\n",
    "            mw > 500,\n",
    "            logp > 5,\n",
    "            hbd > 5,\n",
    "            hba > 10\n",
    "        ])\n",
    "        \n",
    "        data.append({\n",
    "            \"Name\": name,\n",
    "            \"SMILES\": smi,\n",
    "            \"MW\": round(mw, 1),\n",
    "            \"LogP\": round(logp, 2),\n",
    "            \"HBD\": hbd,\n",
    "            \"HBA\": hba,\n",
    "            \"Violations\": violations,\n",
    "            \"Ro5_Pass\": violations < 2\n",
    "        })\n",
    "    \n",
    "    return pd.DataFrame(data)\n",
    "\n",
    "\n",
    "df_props = calculate_lipinski_properties(valid_compounds)\n",
    "print(df_props.to_string(index=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter to Ro5-compliant compounds\n",
    "df_druglike = df_props[df_props[\"Ro5_Pass\"]].copy()\n",
    "print(f\"Drug-like compounds (Ro5 pass): {len(df_druglike)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. PAINS Filtering\n",
    "\n",
    "**Pan-Assay Interference Compounds (PAINS)** are molecules that give false positives in many assays due to:\n",
    "\n",
    "- Aggregation\n",
    "- Redox cycling\n",
    "- Fluorescence interference\n",
    "- Covalent modification\n",
    "\n",
    "RDKit includes PAINS filters that should be applied before screening."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_pains(smiles_dict):\n",
    "    \"\"\"\n",
    "    Filter out PAINS (Pan-Assay Interference Compounds).\n",
    "    \n",
    "    These compounds cause false positives in biochemical assays.\n",
    "    \"\"\"\n",
    "    # Initialize PAINS filter catalog\n",
    "    params = FilterCatalogParams()\n",
    "    params.AddCatalog(FilterCatalogParams.FilterCatalogs.PAINS)\n",
    "    catalog = FilterCatalog(params)\n",
    "    \n",
    "    clean = {}\n",
    "    pains_hits = []\n",
    "    \n",
    "    for name, smi in smiles_dict.items():\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        entry = catalog.GetFirstMatch(mol)\n",
    "        \n",
    "        if entry is None:\n",
    "            clean[name] = smi\n",
    "        else:\n",
    "            pains_hits.append((name, entry.GetDescription()))\n",
    "    \n",
    "    return clean, pains_hits\n",
    "\n",
    "\n",
    "# Apply PAINS filter\n",
    "druglike_smiles = dict(zip(df_druglike[\"Name\"], df_druglike[\"SMILES\"]))\n",
    "clean_compounds, pains = filter_pains(druglike_smiles)\n",
    "\n",
    "print(f\"Compounds passing PAINS filter: {len(clean_compounds)}\")\n",
    "if pains:\n",
    "    print(f\"\\nPAINS hits removed:\")\n",
    "    for name, desc in pains:\n",
    "        print(f\"  {name}: {desc}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## 4. 3D Conformer Generation for Docking\n\nFor virtual screening, we need:\n\n- **Multiple conformers per molecule** (k=3-10) to sample conformational space\n- **Energy window** (typically 3-5 kcal/mol) to include bioactive conformations\n- **Reasonable optimization** - not too tight (slow) but well-minimized\n\nAuto3D's neural network potentials provide accurate geometries much faster than force fields.\n\n### CLI (Recommended)\n\n```bash\n# Basic virtual screening setup\nauto3d run filtered_compounds.smi --k=5 --window=3.0 --gpu\n\n# Fast screening with ANI2xt\nauto3d run filtered_compounds.smi --k=3 --engine=ANI2xt --gpu\n\n# Multi-GPU for large libraries\nauto3d run library.smi --k=5 --gpu --gpu-idx=\"0,1,2,3\"\n\n# With configuration file\nauto3d config init -p quick -o screening.yaml\nauto3d run filtered_compounds.smi -c screening.yaml\n```\n\nExample `screening.yaml` for virtual screening:\n\n```yaml\nk: 5\nwindow: 3.0\noptimizing_engine: ANI2xt\nthreshold: 0.3\nmax_confs: 100\nuse_gpu: true\n```\n\n### Python API"
  },
  {
   "cell_type": "markdown",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "Below is an example using the Python API for programmatic workflows:"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Auto3D for docking preparation\n",
    "# Key settings for virtual screening:\n",
    "# - k=5: Generate up to 5 diverse conformers per molecule\n",
    "# - window=3.0: Include conformers within 3 kcal/mol of minimum\n",
    "# - threshold=0.3: RMSD threshold for conformer diversity\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    config = Auto3DOptions(\n",
    "        path=input_file,\n",
    "        k=5,                        # Top 5 conformers\n",
    "        window=3.0,                 # Within 3 kcal/mol of minimum\n",
    "        optimizing_engine=\"AIMNET\", # Fast and accurate\n",
    "        threshold=0.3,              # RMSD diversity threshold (Angstroms)\n",
    "        max_confs=100,              # Initial conformer pool\n",
    "        use_gpu=True,\n",
    "        verbose=True,\n",
    "    )\n",
    "    \n",
    "    output_sdf = main(config)\n",
    "    print(f\"\\nOutput: {output_sdf}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Post-processing for Docking\n",
    "\n",
    "Most docking programs require specific preparation:\n",
    "\n",
    "- **AutoDock Vina**: PDBQT format with Gasteiger charges\n",
    "- **Glide**: Maestro format with proper protonation\n",
    "- **GOLD**: MOL2 format with proper atom types\n",
    "\n",
    "Here we add properties useful for post-docking analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_docking_properties(sdf_path, output_path=None):\n",
    "    \"\"\"\n",
    "    Add properties useful for docking analysis:\n",
    "    - Formal charge\n",
    "    - Number of rotatable bonds\n",
    "    - TPSA (topological polar surface area)\n",
    "    - Aromatic rings\n",
    "    \"\"\"\n",
    "    if output_path is None:\n",
    "        output_path = sdf_path.replace(\".sdf\", \"_docking_ready.sdf\")\n",
    "    \n",
    "    mols = list(Chem.SDMolSupplier(sdf_path, removeHs=False))\n",
    "    \n",
    "    with Chem.SDWriter(output_path) as writer:\n",
    "        for mol in mols:\n",
    "            if mol is None:\n",
    "                continue\n",
    "            \n",
    "            # Calculate useful properties\n",
    "            mol.SetProp(\"FormalCharge\", str(Chem.GetFormalCharge(mol)))\n",
    "            mol.SetProp(\"NumRotatableBonds\", str(Lipinski.NumRotatableBonds(mol)))\n",
    "            mol.SetProp(\"TPSA\", f\"{Descriptors.TPSA(mol):.1f}\")\n",
    "            mol.SetProp(\"NumAromaticRings\", str(Descriptors.NumAromaticRings(mol)))\n",
    "            mol.SetProp(\"NumHeavyAtoms\", str(mol.GetNumHeavyAtoms()))\n",
    "            \n",
    "            writer.write(mol)\n",
    "    \n",
    "    print(f\"Wrote {len(mols)} conformers to {output_path}\")\n",
    "    return output_path\n",
    "\n",
    "\n",
    "# Add docking-relevant properties\n",
    "if 'output_sdf' in dir():\n",
    "    docking_ready = add_docking_properties(output_sdf)\n",
    "    print(f\"Docking-ready file: {docking_ready}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Conformer Analysis\n",
    "\n",
    "Before docking, inspect the generated conformers to ensure quality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analyze_conformers(sdf_path):\n",
    "    \"\"\"\n",
    "    Analyze the generated conformer ensemble.\n",
    "    \"\"\"\n",
    "    mols = list(Chem.SDMolSupplier(sdf_path, removeHs=False))\n",
    "    \n",
    "    # Group by molecule name\n",
    "    from collections import defaultdict\n",
    "    conformers_by_mol = defaultdict(list)\n",
    "    \n",
    "    for mol in mols:\n",
    "        if mol is None:\n",
    "            continue\n",
    "        name = mol.GetProp(\"_Name\").split(\"_\")[0]  # Remove conformer suffix\n",
    "        energy = float(mol.GetProp(\"E_rel\")) if mol.HasProp(\"E_rel\") else 0.0\n",
    "        conformers_by_mol[name].append(energy)\n",
    "    \n",
    "    print(\"Conformer Summary:\")\n",
    "    print(\"-\" * 50)\n",
    "    print(f\"{'Molecule':<20} {'Count':>6} {'E_range (kcal/mol)':>20}\")\n",
    "    print(\"-\" * 50)\n",
    "    \n",
    "    for name, energies in conformers_by_mol.items():\n",
    "        e_range = max(energies) - min(energies) if len(energies) > 1 else 0.0\n",
    "        print(f\"{name:<20} {len(energies):>6} {e_range:>20.2f}\")\n",
    "    \n",
    "    print(\"-\" * 50)\n",
    "    print(f\"Total conformers: {len(mols)}\")\n",
    "\n",
    "\n",
    "if 'output_sdf' in dir():\n",
    "    analyze_conformers(output_sdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This workflow prepared a compound library for virtual screening:\n",
    "\n",
    "1. **Validated** SMILES and removed problematic entries\n",
    "2. **Filtered** by Lipinski's Rule of Five for oral bioavailability\n",
    "3. **Removed** PAINS compounds that cause assay interference\n",
    "4. **Generated** diverse 3D conformers with Auto3D\n",
    "5. **Annotated** with docking-relevant properties\n",
    "\n",
    "The output SDF file is ready for docking with programs like:\n",
    "- AutoDock Vina (convert to PDBQT with `prepare_ligand4.py`)\n",
    "- Glide (import into Maestro)\n",
    "- GOLD (convert to MOL2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cleanup\n",
    "if 'input_file' in dir():\n",
    "    os.unlink(input_file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}