{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum Chemistry Refinement\n",
    "\n",
    "This notebook demonstrates using Auto3D structures as starting points for DFT calculations:\n",
    "\n",
    "1. **Input generation** - Gaussian, ORCA, Psi4 formats\n",
    "2. **Basis set selection** - appropriate for your accuracy needs\n",
    "3. **Property calculations** - energies, frequencies, NMR\n",
    "4. **Result analysis** - parsing output files\n",
    "\n",
    "## Why Use Auto3D for QM?\n",
    "\n",
    "- **Better starting geometries** → fewer optimization steps\n",
    "- **Pre-filtered conformers** → avoid redundant calculations  \n",
    "- **NNP energies** → initial ranking before expensive QM"
   ]
  },
  {
   "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 AllChem\n",
    "\n",
    "print(f\"Auto3D version: {Auto3D.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Generate Optimized Structures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Molecules for QM calculation\n",
    "molecules = {\n",
    "    \"caffeine\": \"CN1C=NC2=C1C(=O)N(C(=O)N2C)C\",\n",
    "    \"aspirin\": \"CC(=O)OC1=CC=CC=C1C(=O)O\",\n",
    "}\n",
    "\n",
    "with tempfile.NamedTemporaryFile(mode='w', suffix='.smi', delete=False) as f:\n",
    "    for name, smi in molecules.items():\n",
    "        f.write(f\"{smi} {name}\\n\")\n",
    "    input_file = f.name\n",
    "\n",
    "# Generate tight geometry\n",
    "if __name__ == \"__main__\":\n",
    "    config = Auto3DOptions(\n",
    "        path=input_file,\n",
    "        k=1,\n",
    "        optimizing_engine=\"AIMNET\",\n",
    "        use_gpu=True,\n",
    "        convergence_threshold=0.003,  # Tight for QM starting point\n",
    "    )\n",
    "    \n",
    "    sdf_output = main(config)\n",
    "    print(f\"Output: {sdf_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Gaussian Input Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mol_to_gaussian(mol, job_type=\"opt freq\", method=\"B3LYP\", basis=\"6-31G(d)\",\n",
    "                    charge=0, multiplicity=1, nproc=8, mem=\"16GB\"):\n",
    "    \"\"\"\n",
    "    Generate Gaussian input file from RDKit molecule.\n",
    "    \n",
    "    Args:\n",
    "        mol: RDKit molecule with 3D coordinates\n",
    "        job_type: Gaussian job type (opt, freq, nmr, etc.)\n",
    "        method: DFT functional or method\n",
    "        basis: Basis set\n",
    "        charge: Molecular charge\n",
    "        multiplicity: Spin multiplicity\n",
    "        nproc: Number of processors\n",
    "        mem: Memory allocation\n",
    "    \"\"\"\n",
    "    name = mol.GetProp(\"_Name\") if mol.HasProp(\"_Name\") else \"molecule\"\n",
    "    conf = mol.GetConformer()\n",
    "    \n",
    "    # Header\n",
    "    lines = [\n",
    "        f\"%NProcShared={nproc}\",\n",
    "        f\"%Mem={mem}\",\n",
    "        f\"%Chk={name}.chk\",\n",
    "        f\"# {method}/{basis} {job_type}\",\n",
    "        \"\",\n",
    "        f\"{name} - Auto3D optimized geometry\",\n",
    "        \"\",\n",
    "        f\"{charge} {multiplicity}\",\n",
    "    ]\n",
    "    \n",
    "    # Coordinates\n",
    "    for atom in mol.GetAtoms():\n",
    "        pos = conf.GetAtomPosition(atom.GetIdx())\n",
    "        symbol = atom.GetSymbol()\n",
    "        lines.append(f\"{symbol:2s}  {pos.x:12.6f}  {pos.y:12.6f}  {pos.z:12.6f}\")\n",
    "    \n",
    "    lines.append(\"\")  # Blank line at end\n",
    "    \n",
    "    return \"\\n\".join(lines)\n",
    "\n",
    "\n",
    "# Generate Gaussian inputs\n",
    "if 'sdf_output' in dir():\n",
    "    mols = list(Chem.SDMolSupplier(sdf_output, removeHs=False))\n",
    "    \n",
    "    for mol in mols:\n",
    "        if mol is None:\n",
    "            continue\n",
    "        \n",
    "        name = mol.GetProp(\"_Name\").split(\"_\")[0]\n",
    "        charge = Chem.GetFormalCharge(mol)\n",
    "        \n",
    "        # Geometry optimization + frequency\n",
    "        gjf_opt = mol_to_gaussian(mol, \"opt freq\", \"B3LYP\", \"6-311+G(2d,p)\", charge)\n",
    "        \n",
    "        print(f\"\\n=== {name}.gjf ===\")\n",
    "        print(gjf_opt[:500] + \"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. ORCA Input Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mol_to_orca(mol, job_type=\"Opt Freq\", method=\"B3LYP\", basis=\"def2-TZVP\",\n",
    "                charge=0, multiplicity=1, nproc=8):\n",
    "    \"\"\"\n",
    "    Generate ORCA input file from RDKit molecule.\n",
    "    \"\"\"\n",
    "    name = mol.GetProp(\"_Name\") if mol.HasProp(\"_Name\") else \"molecule\"\n",
    "    conf = mol.GetConformer()\n",
    "    \n",
    "    # Header\n",
    "    lines = [\n",
    "        f\"# ORCA input - {name}\",\n",
    "        f\"! {method} {basis} {job_type} D3BJ\",  # Include dispersion\n",
    "        f\"%pal nprocs {nproc} end\",\n",
    "        \"\",\n",
    "        f\"* xyz {charge} {multiplicity}\",\n",
    "    ]\n",
    "    \n",
    "    # Coordinates\n",
    "    for atom in mol.GetAtoms():\n",
    "        pos = conf.GetAtomPosition(atom.GetIdx())\n",
    "        symbol = atom.GetSymbol()\n",
    "        lines.append(f\"  {symbol:2s}  {pos.x:12.6f}  {pos.y:12.6f}  {pos.z:12.6f}\")\n",
    "    \n",
    "    lines.append(\"*\")\n",
    "    \n",
    "    return \"\\n\".join(lines)\n",
    "\n",
    "\n",
    "# Generate ORCA input\n",
    "if 'sdf_output' in dir():\n",
    "    mol = next(Chem.SDMolSupplier(sdf_output, removeHs=False))\n",
    "    if mol:\n",
    "        orca_input = mol_to_orca(mol)\n",
    "        print(\"=== ORCA Input ===\")\n",
    "        print(orca_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Psi4 Input (Python Script)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mol_to_psi4(mol, method=\"B3LYP\", basis=\"cc-pVDZ\", charge=0, multiplicity=1):\n",
    "    \"\"\"\n",
    "    Generate Psi4 Python script from RDKit molecule.\n",
    "    \"\"\"\n",
    "    name = mol.GetProp(\"_Name\") if mol.HasProp(\"_Name\") else \"molecule\"\n",
    "    conf = mol.GetConformer()\n",
    "    \n",
    "    # Build geometry string\n",
    "    geom_lines = []\n",
    "    for atom in mol.GetAtoms():\n",
    "        pos = conf.GetAtomPosition(atom.GetIdx())\n",
    "        symbol = atom.GetSymbol()\n",
    "        geom_lines.append(f\"    {symbol:2s}  {pos.x:12.6f}  {pos.y:12.6f}  {pos.z:12.6f}\")\n",
    "    \n",
    "    geometry = \"\\n\".join(geom_lines)\n",
    "    \n",
    "    script = f'''\n",
    "\"\"\"Psi4 calculation for {name}\"\"\"\n",
    "import psi4\n",
    "\n",
    "psi4.set_memory(\"16 GB\")\n",
    "psi4.set_num_threads(8)\n",
    "\n",
    "mol = psi4.geometry(\"\"\"\n",
    "{charge} {multiplicity}\n",
    "{geometry}\n",
    "\"\"\")\n",
    "\n",
    "# Geometry optimization\n",
    "psi4.set_options({{\n",
    "    \"basis\": \"{basis}\",\n",
    "    \"reference\": \"rhf\" if {multiplicity} == 1 else \"uhf\",\n",
    "    \"scf_type\": \"df\",\n",
    "    \"d_convergence\": 1e-8,\n",
    "}})\n",
    "\n",
    "E_opt = psi4.optimize(\"{method}\", molecule=mol)\n",
    "print(f\"Optimized energy: {{E_opt:.6f}} Hartree\")\n",
    "\n",
    "# Frequency calculation\n",
    "E_freq, wfn = psi4.frequency(\"{method}\", molecule=mol, return_wfn=True)\n",
    "freqs = wfn.frequencies().to_array()\n",
    "print(f\"Frequencies (cm-1): {{freqs}}\")\n",
    "'''\n",
    "    return script.strip()\n",
    "\n",
    "\n",
    "if 'sdf_output' in dir():\n",
    "    mol = next(Chem.SDMolSupplier(sdf_output, removeHs=False))\n",
    "    if mol:\n",
    "        psi4_script = mol_to_psi4(mol)\n",
    "        print(\"=== Psi4 Script ===\")\n",
    "        print(psi4_script[:800] + \"...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Method/Basis Set Recommendations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\"\"\n",
    "QM Method Recommendations for Drug-Like Molecules\n",
    "=================================================\n",
    "\n",
    "GEOMETRY OPTIMIZATION:\n",
    "  Fast:     B3LYP/6-31G(d)           (~1 min/mol)\n",
    "  Standard: B3LYP-D3/6-311G(d,p)     (~5 min/mol)\n",
    "  Accurate: ωB97X-D/def2-TZVP        (~30 min/mol)\n",
    "\n",
    "SINGLE-POINT ENERGY:\n",
    "  Standard: B3LYP-D3/def2-TZVP\n",
    "  Accurate: DLPNO-CCSD(T)/cc-pVTZ    (ORCA)\n",
    "  Gold std: CCSD(T)/CBS extrapolation\n",
    "\n",
    "THERMOCHEMISTRY:\n",
    "  Method:   B3LYP-D3/6-311+G(2d,2p) with frequency\n",
    "  Scale:    0.967 for B3LYP frequencies\n",
    "\n",
    "NMR CALCULATIONS:\n",
    "  1H/13C:   mPW1PW91/6-311+G(2d,p) with GIAO\n",
    "  or:       B3LYP/6-311+G(2d,p) with GIAO\n",
    "\n",
    "DISPERSION:\n",
    "  Always include for drug-like molecules:\n",
    "  - D3(BJ) for Gaussian/ORCA\n",
    "  - DFT-D3 for Psi4\n",
    "  - Built-in for ωB97X-D, ωB97M-V\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Batch Input Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_batch_inputs(sdf_path, output_dir, software=\"gaussian\"):\n",
    "    \"\"\"\n",
    "    Generate QM inputs for all molecules in SDF.\n",
    "    \"\"\"\n",
    "    output_dir = Path(output_dir)\n",
    "    output_dir.mkdir(exist_ok=True)\n",
    "    \n",
    "    mols = list(Chem.SDMolSupplier(sdf_path, removeHs=False))\n",
    "    \n",
    "    generated = []\n",
    "    for mol in mols:\n",
    "        if mol is None:\n",
    "            continue\n",
    "        \n",
    "        name = mol.GetProp(\"_Name\").replace(\" \", \"_\")\n",
    "        charge = Chem.GetFormalCharge(mol)\n",
    "        \n",
    "        if software.lower() == \"gaussian\":\n",
    "            content = mol_to_gaussian(mol, charge=charge)\n",
    "            ext = \".gjf\"\n",
    "        elif software.lower() == \"orca\":\n",
    "            content = mol_to_orca(mol, charge=charge)\n",
    "            ext = \".inp\"\n",
    "        elif software.lower() == \"psi4\":\n",
    "            content = mol_to_psi4(mol, charge=charge)\n",
    "            ext = \".py\"\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown software: {software}\")\n",
    "        \n",
    "        output_file = output_dir / f\"{name}{ext}\"\n",
    "        with open(output_file, 'w') as f:\n",
    "            f.write(content)\n",
    "        \n",
    "        generated.append(str(output_file))\n",
    "    \n",
    "    return generated\n",
    "\n",
    "\n",
    "print(\"Batch generation function available.\")\n",
    "print(\"Usage: generate_batch_inputs('structures.sdf', 'qm_inputs/', 'gaussian')\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Output Parsing (Example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_gaussian_log(log_path):\n",
    "    \"\"\"\n",
    "    Parse Gaussian log file for key results.\n",
    "    \"\"\"\n",
    "    results = {\n",
    "        \"converged\": False,\n",
    "        \"energy\": None,\n",
    "        \"frequencies\": [],\n",
    "        \"imaginary_freqs\": 0,\n",
    "    }\n",
    "    \n",
    "    with open(log_path) as f:\n",
    "        for line in f:\n",
    "            if \"Normal termination\" in line:\n",
    "                results[\"converged\"] = True\n",
    "            elif \"SCF Done\" in line:\n",
    "                results[\"energy\"] = float(line.split()[4])\n",
    "            elif \"Frequencies --\" in line:\n",
    "                freqs = [float(x) for x in line.split()[2:]]\n",
    "                results[\"frequencies\"].extend(freqs)\n",
    "                results[\"imaginary_freqs\"] += sum(1 for f in freqs if f < 0)\n",
    "    \n",
    "    return results\n",
    "\n",
    "\n",
    "print(\"Output parsing functions available.\")\n",
    "print(\"\\nFor comprehensive parsing, consider:\")\n",
    "print(\"  - cclib: pip install cclib\")\n",
    "print(\"  - ASE: ase.io.read()\")\n",
    "print(\"  - GaussView or Avogadro for visualization\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This tutorial covered:\n",
    "\n",
    "1. **Auto3D → QM workflow** - NNP-optimized starting geometries\n",
    "2. **Gaussian** input file generation\n",
    "3. **ORCA** input file generation\n",
    "4. **Psi4** Python script generation\n",
    "5. **Method selection** - appropriate basis sets and functionals\n",
    "6. **Batch processing** - generating inputs for multiple molecules\n",
    "\n",
    "Key benefits:\n",
    "- **Fewer SCF cycles** with better starting geometry\n",
    "- **Pre-ranking** with NNP energies\n",
    "- **Conformer filtering** before expensive QM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cleanup\n",
    "if 'input_file' in dir() and os.path.exists(input_file):\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
}