"""
Molecular Variations Module

This module handles generation of multiple visualizations
and variations of chemical structures.
"""

from rdkit import Chem
from rdkit.Chem import Draw


def generate_molecule_variations(base_smiles, num_variations=12):
    """Generate multiple variations of a chemical structure for the grid display."""
    mol = Chem.MolFromSmiles(base_smiles)
    if not mol:
        return []
    
    variations = []
    
    # Generate different rendering styles and sizes
    sizes = [(150, 150), (180, 180), (200, 200), (160, 160)]
    styles = [True, False]  # kekulize vs non-kekulize
    
    for i in range(num_variations):
        size = sizes[i % len(sizes)]
        kekulize = styles[i % len(styles)]
        
        try:
            # Create variation with different rendering parameters
            img = Draw.MolToImage(mol, size=size, kekulize=kekulize)
            variations.append({
                'image': img,
                'smiles': base_smiles,
                'variation_id': i + 1,
                'size': size,
                'kekulize': kekulize
            })
        except:
            # Fallback to basic rendering
            try:
                img = Draw.MolToImage(mol, size=(150, 150), kekulize=True)
                variations.append({
                    'image': img,
                    'smiles': base_smiles,
                    'variation_id': i + 1,
                    'size': (150, 150),
                    'kekulize': True
                })
            except:
                continue
    
    return variations


def generate_chemical_series_variations(base_smiles):
    """Generate a series of related chemical structures for drug discovery."""
    mol = Chem.MolFromSmiles(base_smiles)
    if not mol:
        return []
    
    variations = []
    
    # Create different visualization styles
    styles = [
        {'size': (200, 200), 'kekulize': True, 'style': 'Standard'},
        {'size': (180, 180), 'kekulize': False, 'style': 'Kekulé'},
        {'size': (220, 220), 'kekulize': True, 'style': 'Large'},
        {'size': (160, 160), 'kekulize': False, 'style': 'Compact'},
        {'size': (200, 200), 'kekulize': True, 'style': 'Detailed'},
        {'size': (190, 190), 'kekulize': False, 'style': 'Minimal'},
        {'size': (210, 210), 'kekulize': True, 'style': 'Enhanced'},
        {'size': (170, 170), 'kekulize': False, 'style': 'Focused'},
        {'size': (200, 200), 'kekulize': True, 'style': 'Classic'},
        {'size': (185, 185), 'kekulize': False, 'style': 'Modern'},
        {'size': (195, 195), 'kekulize': True, 'style': 'Scientific'},
        {'size': (175, 175), 'kekulize': False, 'style': 'Clean'}
    ]
    
    for i, style_config in enumerate(styles):
        try:
            print(f"Generating variation {i+1}: {style_config['style']}")
            img = Draw.MolToImage(mol, size=style_config['size'], kekulize=style_config['kekulize'])
            print(f"Successfully generated image for {style_config['style']}, type: {type(img)}")
            variations.append({
                'image': img,
                'smiles': base_smiles,
                'variation_id': i + 1,
                'style': style_config['style'],
                'size': style_config['size'],
                'kekulize': style_config['kekulize']
            })
        except Exception as e:
            print(f"Failed to generate image for {style_config['style']}: {e}")
            continue
    
    return variations


def generate_molecule_images(smiles_list):
    """Generate images for a list of SMILES strings."""
    images = []
    for smiles in smiles_list:
        mol = Chem.MolFromSmiles(smiles)
        if mol:
            # Create a high-quality image with better rendering
            img = Draw.MolToImage(mol, size=(200, 200), kekulize=True)
            images.append(img)
        else:
            # Add a placeholder if SMILES is invalid
            try:
                from PIL import Image
                images.append(Image.new('RGB', (200, 200), color='white'))
            except ImportError:
                # Fallback if PIL not available
                images.append(None)
    return images