#!/usr/bin/env python
# coding: utf-8

# # Conformer Search

# In this notebook, we use the ORCA Python Interface (OPI) to perform a conformer search with a force field. Since the resulting energy ranking may lack accuracy, we re-rank the conformers using DFT single point energy calculations. We will first demonstrate how this can be done with RDKit and afterwards wih GOAT.
# 
# Workflow:
# 1. Import required Python dependencies.
# 2. Define a working directory.
# 3. Define an input structure via SMILES
# 4. Generate conformers with RDKit.
# 5. perform DFT single point calculations for re-ranking.
# 6. Visualize the resulting conformer ensemble.
# 7. Generate conformers using GOAT via OPI and re-rank them.

# ## Step 1: Import Dependencies
# 
# We start by importing the modules needed for:
# - Interfacing with ORCA input/output
# - RDKit
# - Numerical calculations and data handling
# - Plotting results
# 
# > **Note:** We additionally import modules for visualization/plotting like `py3Dmol`. For this, it might be necessary to install `py3Dmol` into your OPI `venv` (e.g., by activating the `.venv` and using `uv pip install py3Dmol`).

# In[1]:


import re
import copy
import shutil
from pathlib import Path

# RDKit as conformer generator
from rdkit import Chem
from rdkit.Chem import AllChem, rdMolAlign
from rdkit.Chem.rdchem import Mol

# OPI imports for performing ORCA calculations and reading output
from opi.core import Calculator
from opi.input.structures.structure import Structure
from opi.input.simple_keywords import BasisSet, Dft, ForceField, Scf
from opi.input.simple_keywords.goat import Goat

# for plotting results and visualization of molecules
import py3Dmol
import pandas as pd
import matplotlib.pyplot as plt


# ## Step 2: Working Directory and Conversion Factor
# 
# We define a subfolder `RUN` in which the actual ORCA calculations will take place. Also, we define a conversion factor, since we want the resulting interaction energies in kcal/mol for better interpretability.

# In[ ]:


# > Calculation is performed in `RUN`
working_dir = Path("RUN")
# > The `working_dir`is automatically (re-)created
shutil.rmtree(working_dir, ignore_errors=True)
working_dir.mkdir()
# > Conversion factor for atomic units to kcal/mol
unit_conversion = 627.509 


# ## Step 3: Define an Input Structure via SMILES
# 
# We define the input structure as a smiles string and load the smiles string to an RDKit Mol. Then we visualize the 2D structure. After that we add hydrogen atoms and identify the chiral center. 

# In[3]:


input_smiles='C1=CC(=CC=C1C[C@@H](C(=O)O)N)O'
# Convert SMILES to RDKit Mol and label atoms for visualization
mol = Chem.MolFromSmiles(input_smiles)
display(mol)
for atom in mol.GetAtoms():
    atom.SetProp('atomLabel', str(atom.GetIdx()))
display(mol)

# Add hydrogens and identify chiral centers
mol = Chem.AddHs(mol)
chiral_centers = Chem.FindMolChiralCenters(mol)
if chiral_centers:
    print('Chiral centers identified:', chiral_centers)


# ## Step 4: Conformer Generation with RDKit
# Conformer ensembles do not have to be generated with OPI. They can come from various sources. Here, we use the conformers generated by RDKit:

# In[4]:


def confGenerator_RDKit(smiles: str, e_thresh: float, confs: int, 
                        rms: float, method: str) -> tuple[list[tuple[float, int]], Mol, str]:
    """Generate and optimize conformers using RDKit"""

    # Generate 3D conformers
    conformer_ids = AllChem.EmbedMultipleConfs(
        mol,
        numConfs=confs,
        pruneRmsThresh=rms,
        randomSeed=1,
        useExpTorsionAnglePrefs=True,
        useBasicKnowledge=True
    )
    print('Number of raw conformers:', len(conformer_ids))

    # Optimization for each conformer
    mmff_props = AllChem.MMFFGetMoleculeProperties(mol, mmffVariant=method)
    conformer_energies = []
    for conf_id in conformer_ids:
        ff = AllChem.MMFFGetMoleculeForceField(mol, mmff_props, confId=conf_id)
        ff.Minimize()
        energy_value = float(ff.CalcEnergy())
        conformer_energies.append((energy_value, conf_id))

    # Sort by energy and filter below energy threshold
    conformer_energies.sort()
    min_energy = conformer_energies[0][0]

    filtered_mol = copy.deepcopy(mol)
    filtered_mol.RemoveAllConformers()
    selected_conf_ids = []
    filtered_rel_e = []

    data = []
    for idx, (energy_value, conf_id) in enumerate(conformer_energies, start=1):
        delta_e = energy_value - min_energy
        status = "below" if delta_e <= e_thresh else "above"
        data.append({
            "ConfID": conf_id,
            "Energy (kcal/mol)": f"{energy_value:.2f}",
            "ΔE (kcal/mol)": f"{delta_e:.2f}",
            "Status": status
        })
        if delta_e <= e_thresh:
            conf = mol.GetConformer(conf_id)
            filtered_mol.AddConformer(conf)
            selected_conf_ids.append(conf_id)
            filtered_rel_e.append((delta_e, conf_id))

    df = pd.DataFrame(data)
    print(f"\nNumber of conformers below {e_thresh} kcal/mol ({method}): {len(selected_conf_ids)}")
    display(df.style.hide(axis="index"))
    return filtered_rel_e, filtered_mol, method

# > Run RDKit conformer search
filtered_rel_e_rdkit, filtered_mol_rdkit, method_rdkit = confGenerator_RDKit(smiles=input_smiles, e_thresh=1.5, confs=60, rms=0.02, method='MMFF94s')


# ## Step 5: DFT Single Point Energy Calculations
# 
# To re-rank the generated conformers we can employ OPI for DFT single&ndash;point energy calculations. In the functions below this is done with the composite DFT method r²SCAN-3c. The re-ranking is directly visualized with matplotlib. The energy window for performing DFT calculations is set smaller than necessary for prdouction runs to keep the computational costs of this notebook low.

# In[5]:


def dft_calculations(smiles: str, e_thresh: float, functional: Dft, basis: BasisSet | None, 
                     rel_e: list[tuple[float, int]], mol: Mol, method: str, working_dir: Path = Path("RUN")) -> tuple[Mol, list[int]]:
    """Perform DFT single point energy calculations using ORCA"""
    smiles = re.sub(r'[^A-Za-z0-9]', '_', smiles)
    smiles_dir = working_dir / f"{smiles}_{method}"
    smiles_dir.mkdir(exist_ok=True)

    dft_energies = []
    comparison = []
    print("\nRunning DFT calculations with ORCA...")

    for delta_e, cid in rel_e:

        xyz_file = smiles_dir / f"conf_{cid}.xyz"
        xyz_block = Chem.MolToXYZBlock(mol, confId=cid)

        with open(xyz_file, 'w') as f:
            f.write(xyz_block)

        # > Set up the structure
        xyz_file = Path(xyz_file)
        mol_name = xyz_file.stem
        structure = Structure.from_xyz(xyz_file)

        # > Define a directory for calculation files
        calc_dir = xyz_file.parent / mol_name
        calc_dir.mkdir(exist_ok=True)

        # > Set up the calculator
        calc = Calculator(basename=mol_name, working_dir=calc_dir)
        calc.structure = structure
        # > Neutral structure
        calc.structure.charge = 0
        # > Spin multiplicity of 1 (closed-shell)
        calc.structure.multiplicity = 1
        # Some methods like r²SCAN-3c have their own predefined basis set
        if basis:
            calc.input.add_simple_keywords(functional, basis, Scf.NOAUTOSTART)
        else:
            calc.input.add_simple_keywords(functional, Scf.NOAUTOSTART)

        # > Run the calculation on 4 cores
        calc.input.ncores = 4

        # > Write ORCA input file
        calc.write_input()

        # > Run the calculation
        calc.run()

        # > Get and check the output
        output = calc.get_output()
        # > Check for proper termination of ORCA
        status = output.terminated_normally()
        if not status:
            # > ORCA did not terminate normally
            raise RuntimeError(f"ORCA did not terminate normally, see output file: {output.get_outfile()}")
        else:
            # > ORCA did terminate normally so we can parse the output
            output.parse()

        # > Now check for convergence of the SCF
        if not output.results_properties.geometries[0].single_point_data.converged:
            raise RuntimeError("SCF DID NOT CONVERGE")

        # > Obtain the energy from the DFT calculation
        e_dft = output.results_properties.geometries[0].single_point_data.finalenergy

        dft_energies.append((e_dft, cid))
        comparison.append({
            'Conformer': cid,
            f'{method} (kcal/mol)': delta_e,                
            'DFT (kcal/mol)': 0
        })
        print(f"Conformer {cid}: DFT Energy = {e_dft:.6f} Eh")

    # Calculate relative DFT energies
    dft_energies.sort()
    min_dft = dft_energies[0][0]

    for entry in comparison:
        cid = entry['Conformer']
        for e_dft, dft_id in dft_energies:
            if cid == dft_id:
                entry['DFT (kcal/mol)'] = (e_dft - min_dft) * unit_conversion
                break

    df_dft = pd.DataFrame(comparison).sort_values(f'{method} (kcal/mol)')

    # Plot energies differences
    x_vals = range(len(df_dft))
    conformer_ids = df_dft['Conformer']

    plt.figure(figsize=(10, 6))
    plt.plot(x_vals, df_dft[f'{method} (kcal/mol)'], 'bo', label=f'{method}')
    plt.plot(x_vals, df_dft['DFT (kcal/mol)'], 'ro', label='DFT')
    plt.xticks(x_vals, conformer_ids)
    plt.xlabel('Conformer ID')
    plt.ylabel('Relative Energy (kcal/mol)')
    plt.title(f'{method} vs DFT Relative Energies')
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()

    # Compare conformer ranks by FF vs DFT energy
    print("\nEnergy Comparison Table:")
    method_rank = df_dft.sort_values(f'{method} (kcal/mol)')['Conformer'].tolist()
    dft_rank = df_dft.sort_values('DFT (kcal/mol)')['Conformer'].tolist()

    comparison_table = []
    for _, row in df_dft.iterrows():
        cid = row['Conformer']
        method_r = method_rank.index(cid) + 1
        dft_r = dft_rank.index(cid) + 1
        delta = method_r - dft_r
        mark = f"{'↑' if delta > 0 else '↓' if delta < 0 else '='}{abs(delta)}" if delta else "-"

        comparison_table.append({
            "ConfID": int(cid),
            f"{method} (kcal/mol)": f"{row[f'{method} (kcal/mol)']:.2f}",
            "DFT (kcal/mol)": f"{row['DFT (kcal/mol)']:.2f}",
            "Rank Change": mark
        })

    df_rank = pd.DataFrame(comparison_table)
    display(df_rank.style.hide(axis="index"))

    # Final filter based on DFT energy threshold
    filtered_mol_dft = copy.deepcopy(mol)
    filtered_mol_dft.RemoveAllConformers()
    final_ids = []

    data_dft = []
    for idx, (e_dft, conf_id) in enumerate(dft_energies, start=1):
        delta_dft = (e_dft - min_dft) * 627.509
        status = "below" if delta_dft <= e_thresh else "above"
        data_dft.append({
            "ConfID": conf_id,
            "DFT Energy (Eh)": f"{e_dft:.6f}",
            "ΔE (kcal/mol)": f"{delta_dft:.2f}",
            "Status": status
        })
        if delta_dft <= e_thresh:
            filtered_mol_dft.AddConformer(mol.GetConformer(conf_id))
            final_ids.append(conf_id)

    df_dft_filtered = pd.DataFrame(data_dft)
    print(f"\nNumber of conformers below {e_thresh} kcal/mol (DFT): {len(final_ids)}")
    display(df_dft_filtered.style.hide(axis="index"))
    return filtered_mol_dft, final_ids

# > Run DFT single point energy calculations with r²SCAN-3c with OPI
filtered_mol_dft_rdkit, final_ids_rdkit = dft_calculations(smiles=input_smiles, e_thresh=1.0, functional= Dft.R2SCAN_3C, basis=None, rel_e=filtered_rel_e_rdkit, mol=filtered_mol_rdkit, method=method_rdkit, working_dir=working_dir)


# ## Step 6: Visualize Selected Conformers Using py3Dmol
# 
# The final conformer ensemble can be visualized using py3Dmol:

# In[6]:


def visualization(align: list[int], highlight: list[int] | None, mol: Mol, ids: list[int]) -> None:
    """Visualize selected conformers using py3Dmol with optional highlighting"""
    if not isinstance(highlight, list):
        highlight = [highlight]
    # Visualize final conformers with py3Dmol
    view = py3Dmol.view(width=500, height=500)
    mol = Chem.RemoveHs(mol)

    rdMolAlign.AlignMolConformers(mol, atomIds=align)

    for idx, conf_id in enumerate(ids, start=1):
        view.addModel(Chem.MolToMolBlock(mol, confId=conf_id))
        model = view.getModel()
        if not highlight or idx in highlight:
            model.setStyle({'stick': {'opacity': 1, 'radius': 0.2}})
        else:
            model.setStyle({'stick': {'opacity': 0.8, 'radius': 0.1}})

    view.setBackgroundColor('white')
    view.zoomTo()
    view.show()

visualization(align=[0,1,2,3], highlight=1, mol=filtered_mol_dft_rdkit, ids=final_ids_rdkit)


# ## Step 7: Elaborate Conformer Generation with GOAT
# A more elaborate (but also computationally more expensive) way to generate a conformer ensemble is using [GOAT](https://doi.org/10.1002/anie.202500393). Here, we use GOAT with GFN-FF, but in principle, every method available in ORCA could be used. After that we re-rank with r²SCAN-3c and visualize the results as before.

# In[7]:


def confGenerator_GOAT(smiles: str, e_thresh: float, method: ForceField, working_dir: Path = Path("RUN"))-> tuple[list[tuple[float, int]], Mol, str]:
    """Generate and optimize conformers using ORCA GOAT"""

    # Generate initial 3D structure 
    smiles = re.sub(r'[^A-Za-z0-9]', '_', smiles)
    smiles_dir = working_dir / Path(f"{smiles}_{method}")
    smiles_dir.mkdir(exist_ok=True)

    xyz_file = smiles_dir / f"{smiles}.xyz"
    AllChem.EmbedMolecule(mol, AllChem.ETKDG())
    xyz_block = Chem.MolToXYZBlock(mol)

    with open(xyz_file, 'w') as f:
        f.write(xyz_block)

    # > Set up the GOAT run

    # > Prepare the structures
    xyz_file = Path(xyz_file)
    mol_name = xyz_file.stem
    structure = Structure.from_xyz(xyz_file)

    # > Set up the calcualtor
    calc_dir = xyz_file.parent / mol_name
    calc_dir.mkdir(exist_ok=True)
    calc = Calculator(basename=mol_name, working_dir=calc_dir)
    calc.structure = structure
    # > Neutral molecule
    calc.structure.charge = 0
    # > Spin multiplicity of 1 (close-shell)
    calc.structure.multiplicity = 1

    calc.input.add_simple_keywords(
        method,
        Goat.GOAT
    )

    # > Use 8 cores for the calculation
    calc.input.ncores = 8

    # > Write the ORCA input file
    calc.write_input()

    # > Run the calculation
    calc.run()

    # The results of the conformer run can be found in the basename.finalensemble.xyz file
    result_file = calc_dir / f"{mol_name}.finalensemble.xyz"

    with open(result_file, 'r') as f:
        lines = f.readlines()

    i = 0
    conformer_energies = []
    mol.RemoveAllConformers()
    conf_counter = 0

    while i < len(lines):
        atom_count = int(lines[i].strip())
        energy_line = lines[i + 1].strip()
        atoms_block = lines[i + 2:i + 2 + atom_count]

        match = re.match(r"([-+]?[0-9]*\.?[0-9]+)", energy_line)
        energy = float(match.group(1)) * unit_conversion

        conf = Chem.Conformer(atom_count)
        for idx, line in enumerate(atoms_block):
            _, x, y, z = line.split()
            conf.SetAtomPosition(idx, (float(x), float(y), float(z)))

        conf_id = mol.AddConformer(conf, assignId=True)
        conformer_energies.append((energy, conf_id))

        i += 2 + atom_count
        conf_counter += 1
    print('Number of raw conformers:', len(conformer_energies))

    # Sort by energy and filter below energy threshold
    conformer_energies.sort()
    min_energy = conformer_energies[0][0]

    filtered_mol = copy.deepcopy(mol)
    filtered_mol.RemoveAllConformers()
    selected_conf_ids = []
    filtered_rel_e = []

    data = []
    for idx, (energy_value, conf_id) in enumerate(conformer_energies, start=1):
        delta_e = energy_value - min_energy
        status = "below" if delta_e <= e_thresh else "above"
        data.append({
            "ConfID": conf_id,
            "Energy (kcal/mol)": f"{energy_value:.2f}",
            "ΔE (kcal/mol)": f"{delta_e:.2f}",
            "Status": status
        })
        if delta_e <= e_thresh:
            conf = mol.GetConformer(conf_id)
            filtered_mol.AddConformer(conf)
            selected_conf_ids.append(conf_id)
            filtered_rel_e.append((delta_e, conf_id))

    df = pd.DataFrame(data)
    print(f"\nNumber of conformers below {e_thresh} kcal/mol ({method}): {len(selected_conf_ids)}")
    display(df.style.hide(axis="index"))
    return filtered_rel_e, filtered_mol, method

# > Run conformer generation with GOAT with 1.5 kcal/mol energy window and the GFN-FF
filtered_rel_e_goat, filtered_mol_goat, method_goat = confGenerator_GOAT(smiles=input_smiles, e_thresh=1.5, method=ForceField.GFN_FF, working_dir=working_dir)
# > Perform DFT single point energy calculations
filtered_mol_dft_goat, final_ids_goat = dft_calculations(smiles=input_smiles, e_thresh=1.0, functional= Dft.R2SCAN_3C, basis= None, rel_e=filtered_rel_e_goat, mol=filtered_mol_goat, method=method_goat, working_dir=working_dir)
# > Visualize the results
visualization(align=[0,1,2,3], highlight=1, mol=filtered_mol_dft_goat, ids=final_ids_goat)


# ## Summary
# 
# In this notebook, we demonstrated handling conformers. We utilized RDKit and GOAT to generate the conformers and re-ranked the conformers with DFT. Both can be done directly with OPI, streamlining common conformer workflows. The results were visualized within this notebook.  
