main.py

"""
SMARTSViewer - Molecular Structure Analysis Tool
Main entry point for the Streamlit application.
"""

from rdkit import Chem
from rdkit.Chem.rdchem import Mol
from rdkit.Chem import AllChem, Draw, rdDepictor


def get_mol(smiles: str) -> Mol:
    """
    Create a sanitized RDKit molecule from SMILES string.
    
    Args:
        smiles: SMILES string
        
    Returns:
        RDKit Mol object with 2D coordinates
    """
    mol = Chem.MolFromSmiles(smiles)
    Chem.SanitizeMol(mol)
    Chem.SetAromaticity(mol)
    mol.RemoveAllConformers()
    rdDepictor.Compute2DCoords(mol)
    return mol


def get_matched_atoms_and_bonds(smiles: str, smarts: str):
    """
    Find atoms and bonds matching a SMARTS pattern.
    
    Args:
        smiles: SMILES string of the molecule
        smarts: SMARTS pattern to match
        
    Returns:
        Tuple of (atoms, bonds) lists
    """
    mol = get_mol(smiles)
    pattern = Chem.MolFromSmarts(smarts)
    # Get matching atoms
    matches = mol.GetSubstructMatches(pattern)

    atoms = []
    bonds = []

    for match in matches:
        for i in match:
            atoms.append(i)
    for bond in mol.GetBonds():
        a1 = bond.GetBeginAtomIdx()
        a2 = bond.GetEndAtomIdx()
        for match in matches:
            if a1 in match and a2 in match:
                bonds.append(bond.GetIdx())
    return atoms, bonds


def _get_highlighted_mol_image(mol: Mol, atoms: list[int], bonds: list[int]):
    """
    Create a highlighted molecule image.
    
    Args:
        mol: RDKit Mol object
        atoms: List of atom indices to highlight
        bonds: List of bond indices to highlight
        
    Returns:
        PIL Image object
    """
    img = Draw.MolToImage(mol, highlightAtoms=atoms, highlightBonds=bonds)
    return img


def perform_reaction(reaction_smarts: str, *reactants: Mol):
    """
    Perform a chemical reaction using reaction SMARTS.
    
    Args:
        reaction_smarts: Reaction SMARTS pattern
        reactants: Variable number of reactant molecules
        
    Returns:
        Tuple of product tuples
    """
    rxn = AllChem.ReactionFromSmarts(reaction_smarts)
    prods = rxn.RunReactants(reactants)
    return prods


def main():
    """Main entry point - runs the Streamlit application."""
    print("Starting SMARTSViewer Streamlit application...")
    print("Please use 'streamlit run app.py' to run the web application.")
    print("Or use 'streamlit run main.py' for direct execution.")
    
    # Import and run the Streamlit app
    from app import main as app_main
    app_main()


if __name__ == "__main__":
    main()