utils.py

"""
Utility functions for SMARTSViewer application.
"""

import yaml
from typing import List, Tuple, Optional, Union
from io import BytesIO
import streamlit as st
from rdkit import Chem
from rdkit.Chem.rdchem import Mol
from rdkit.Chem import AllChem, Draw, rdDepictor
from PIL import Image


def load_config() -> dict:
    """Load configuration from config.yaml file."""
    try:
        with open('config.yaml', 'r') as f:
            return yaml.safe_load(f)
    except FileNotFoundError:
        st.error("Configuration file not found. Please ensure config.yaml exists.")
        return {}


def validate_smarts(smarts: str) -> bool:
    """
    Validate if a SMARTS pattern is correctly formatted.
    
    Args:
        smarts: SMARTS pattern string to validate
        
    Returns:
        True if valid, False otherwise
    """
    try:
        if not smarts or not smarts.strip():
            return False
        pattern = Chem.MolFromSmarts(smarts)
        return pattern is not None
    except Exception:
        return False


def validate_reaction_smarts(reaction_smarts: str) -> bool:
    """
    Validate if a reaction SMARTS pattern is correctly formatted.
    
    Args:
        reaction_smarts: Reaction SMARTS pattern string to validate
        
    Returns:
        True if valid, False otherwise
    """
    try:
        if not reaction_smarts or not reaction_smarts.strip():
            return False
        rxn = AllChem.ReactionFromSmarts(reaction_smarts)
        return rxn is not None
    except Exception:
        return False


def smiles_to_mol(smiles: str) -> Optional[Mol]:
    """
    Convert SMILES string to RDKit Mol object.
    
    Args:
        smiles: SMILES string
        
    Returns:
        RDKit Mol object or None if invalid
    """
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        Chem.SanitizeMol(mol)
        Chem.SetAromaticity(mol)
        mol.RemoveAllConformers()
        rdDepictor.Compute2DCoords(mol)
        return mol
    except Exception:
        return None


def get_substructure_matches(mol: Mol, smarts: str) -> Tuple[List[int], List[int]]:
    """
    Find atoms and bonds matching a SMARTS pattern in a molecule.
    
    Args:
        mol: RDKit Mol object
        smarts: SMARTS pattern string
        
    Returns:
        Tuple of (matching_atoms, matching_bonds) lists
    """
    try:
        pattern = Chem.MolFromSmarts(smarts)
        if pattern is None:
            return [], []
        
        matches = mol.GetSubstructMatches(pattern)
        
        atoms = []
        bonds = []
        
        for match in matches:
            atoms.extend(match)
        
        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
    except Exception:
        return [], []


def create_highlighted_mol_image(mol: Mol, atoms: List[int], bonds: List[int], 
                                size: Tuple[int, int] = (600, 400)) -> Image.Image:
    """
    Create a high-quality image of a molecule with highlighted atoms and bonds.
    
    Args:
        mol: RDKit Mol object
        atoms: List of atom indices to highlight
        bonds: List of bond indices to highlight
        size: Image size tuple (width, height)
        
    Returns:
        PIL Image object
    """
    try:
        # Use high-quality drawer for highlighting
        from rdkit.Chem.Draw import rdMolDraw2D
        from io import BytesIO
        
        # Create drawer with high resolution
        drawer = rdMolDraw2D.MolDraw2DCairo(size[0], size[1])
        
        # Set drawing options for better appearance
        opts = drawer.drawOptions()
        opts.atomLabelFontSize = 22
        opts.bondLineWidth = 2.5
        opts.highlightBondWidthMultiplier = 15
        opts.addAtomIndices = False
        opts.addStereoAnnotation = True
        opts.continuousHighlight = True
        opts.circleAtoms = True
        opts.fillHighlights = True
        
        # Set highlight colors
        highlight_colors = {}
        highlight_bond_colors = {}
        
        # Color highlighted atoms (red)
        for atom in atoms:
            highlight_colors[atom] = (1.0, 0.2, 0.2, 0.5)  # Light red with transparency
            
        # Color highlighted bonds (blue)
        for bond in bonds:
            highlight_bond_colors[bond] = (0.2, 0.2, 1.0, 0.5)  # Light blue with transparency
        
        # Draw the molecule with highlights
        drawer.DrawMolecule(mol, highlightAtoms=atoms, highlightBonds=bonds,
                           highlightAtomColors=highlight_colors,
                           highlightBondColors=highlight_bond_colors)
        drawer.FinishDrawing()
        
        # Get the image data
        img_data = drawer.GetDrawingText()
        
        # Convert to PIL Image
        img = Image.open(BytesIO(img_data))
        return img
        
    except Exception:
        # Fallback to basic RDKit drawing
        try:
            img = Draw.MolToImage(
                mol, 
                size=size, 
                highlightAtoms=atoms, 
                highlightBonds=bonds,
                kekulize=True,
                imageType='png'
            )
            return img
        except Exception:
            return None



def sort_reactants(reactants:List[Mol], reaction_smarts:str):
    reactants_smarts = reaction_smarts[:reaction_smarts.find(">>")].split(".")
    reactants_smarts = [Chem.MolFromSmarts(s) for s in reactants_smarts]
    result = [None] * len(reactants_smarts)
    used_mols = set()
    
    # Create a mapping of which molecules contain which patterns
    pattern_to_mols = {}
    for i, pattern in enumerate(reactants_smarts):
        pattern_to_mols[i] = [j for j, mol in enumerate(reactants) if mol.HasSubstructMatch(pattern)]
    
    # Greedily assign molecules to patterns, starting with patterns that have fewer options
    pattern_indices = sorted(range(len(reactants_smarts)), key=lambda i: len(pattern_to_mols[i]))
    
    for pattern_idx in pattern_indices:
        available_mols = [mol_idx for mol_idx in pattern_to_mols[pattern_idx] 
                         if mol_idx not in used_mols]
        
        if available_mols:
            chosen_mol_idx = available_mols[0]
            result[pattern_idx] = reactants[chosen_mol_idx]
            used_mols.add(chosen_mol_idx)
        else:
            raise ValueError(f"No available molecule for pattern '{reactants_smarts[pattern_idx]}'")
    
    return result



def run_reaction(reaction_smarts: str, reactants: List[Mol]) -> List[Tuple[Mol, ...]]:
    """
    Run a chemical reaction using reaction SMARTS pattern and process products.
    
    Args:
        reaction_smarts: Reaction SMARTS pattern
        reactants: List of reactant molecules
        
    Returns:
        List of product tuples with properly sanitized molecules
    """
    try:
        rxn = AllChem.ReactionFromSmarts(reaction_smarts)
        if rxn is None:
            return []
        reactants = sort_reactants(reactants, reaction_smarts)
        print(Chem.MolToSmiles(r) for r in reactants)
        products = rxn.RunReactants(reactants)
        
        # Process and sanitize product molecules
        processed_products = []
        for product_set in products:
            processed_set = []
            for product in product_set:
                # Sanitize and prepare each product molecule
                processed_product = _prepare_molecule(product)
                if processed_product is not None:
                    processed_set.append(processed_product)
            if processed_set:  # Only add non-empty product sets
                processed_products.append(tuple(processed_set))
        
        return processed_products
    except Exception as e:
        print(e)
        return []


def mol_to_image(mol: Mol, size: Tuple[int, int] = (400, 300)) -> Image.Image:
    """
    Convert RDKit Mol object to PIL Image with basic rendering.
    
    Args:
        mol: RDKit Mol object
        size: Image size tuple (width, height)
        
    Returns:
        PIL Image object
    """
    try:
        # Ensure molecule is properly sanitized and has 2D coordinates
        mol = _prepare_molecule(mol)
        if mol is None:
            return None
            
        # Use basic RDKit drawing
        return Draw.MolToImage(
            mol, 
            size=size,
            kekulize=True,
            wedgeBonds=True,
            imageType='png',
            fitImage=True
        )
        
    except Exception:
        return None


def _prepare_molecule(mol: Mol) -> Optional[Mol]:
    """
    Prepare molecule for rendering by ensuring proper sanitization and 2D coordinates.
    
    Args:
        mol: RDKit Mol object
        
    Returns:
        Prepared molecule or None if preparation fails
    """
    try:
        if mol is None:
            return None
            
        # Create a copy to avoid modifying the original
        mol_copy = Chem.MolFromSmiles(Chem.MolToSmiles(mol))
        
        # Sanitize the molecule
        try:
            Chem.SanitizeMol(mol_copy)
        except:
            # If sanitization fails, try to fix common issues
            try:
                Chem.SanitizeMol(mol_copy, sanitizeOps=Chem.SANITIZE_ALL^Chem.SANITIZE_KEKULIZE)
                Chem.Kekulize(mol_copy)
            except:
                pass
        
        # Ensure aromaticity is set
        try:
            Chem.SetAromaticity(mol_copy)
        except:
            pass
        
        try:
            mol_copy.RemoveAllConformers()
            rdDepictor.Compute2DCoords(mol_copy)
        except:
            pass
        
        return mol_copy
        
    except Exception:
        return mol





def display_success_message(message: str) -> None:
    """Display success message in Streamlit."""
    st.success(message)


def display_error_message(message: str) -> None:
    """Display error message in Streamlit."""
    st.error(message)


def display_warning_message(message: str) -> None:
    """Display warning message in Streamlit."""
    st.warning(message)


def display_info_message(message: str) -> None:
    """Display info message in Streamlit."""
    st.info(message)


def _safe_sanitize_mol(mol: Mol) -> Optional[Mol]:
    """
    Safely sanitize a molecule, handling potential errors.
    
    Args:
        mol: RDKit Mol object
        
    Returns:
        Sanitized molecule or None if sanitization fails
    """
    try:
        if mol is None:
            return None
        Chem.SanitizeMol(mol)
        return mol
    except Exception:
        return None


def _generate_2d_coords(mol: Mol) -> Mol:
    """
    Generate 2D coordinates for a molecule.
    
    Args:
        mol: RDKit Mol object
        
    Returns:
        Molecule with 2D coordinates
    """
    try:
        mol.RemoveAllConformers()
        rdDepictor.Compute2DCoords(mol)
        return mol
    except Exception:
        return mol