#!/usr/bin/env python3
"""
Generate a 3D MolBlock from a SMILES string using RDKit.
Usage: python3 generate3d.py "<SMILES>"
Outputs MolBlock to stdout, error message to stderr with exit code 1.
"""
import sys
from rdkit import Chem
from rdkit.Chem import AllChem

def generate(smiles):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        print('Invalid SMILES', file=sys.stderr)
        sys.exit(1)
    mol = Chem.AddHs(mol)

    # Support older RDKit builds on Ubuntu that may not include ETKDGv3.
    embed_factories = [
        getattr(AllChem, 'ETKDGv3', None),
        getattr(AllChem, 'ETKDGv2', None),
        getattr(AllChem, 'ETKDG', None),
    ]

    result = -1
    for factory in embed_factories:
        if factory is None:
            continue
        result = AllChem.EmbedMolecule(mol, factory())
        if result != -1:
            break
    if result == -1:
        print('3D embedding failed', file=sys.stderr)
        sys.exit(1)
    AllChem.UFFOptimizeMolecule(mol)
    print(Chem.MolToMolBlock(mol))

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print('Usage: generate3d.py "<SMILES>"', file=sys.stderr)
        sys.exit(1)
    generate(sys.argv[1])