src/ui/components.py

"""
UI Components Module

This module contains reusable UI components and layouts
for the drug discovery application.
"""

import gradio as gr
from .handlers import VariationHandlers


def create_molecular_analysis_tab():
    """Create the molecular analysis tab components."""
    with gr.Column(scale=1):
        gr.Markdown("## 🔬 Molecular Analysis")
        smiles_input = gr.Textbox(
            label="Enter SMILES string",
            placeholder="e.g., C[C@H](N)C(=O)O",
            info="Enter a valid SMILES string for molecular analysis"
        )
        
        # Real-time validation status
        validation_status = gr.Markdown("✅ Ready to analyze")
        preview_image = gr.Image(
            label="Preview",
            show_download_button=False,
            width=200,
            height=200,
            visible=False
        )
        
        # Bookmark functionality
        with gr.Row():
            bookmark_name = gr.Textbox(
                placeholder="Enter molecule name (optional)",
                label="Molecule Name",
                scale=2
            )
            bookmark_btn = gr.Button("🔖 Bookmark", variant="secondary", size="sm", scale=1)
        
        bookmark_status = gr.Markdown("")
        
        gr.Markdown("### 📊 Analysis Results")
        analysis_output = gr.Markdown()
    
    with gr.Column(scale=1):
        gr.Markdown("### 🧪 Molecular Structure")
        molecule_image = gr.Image(
            label="Chemical Structure",
            show_download_button=False,
            width=400,
            height=400
        )
    
    return {
        'smiles_input': smiles_input,
        'validation_status': validation_status,
        'preview_image': preview_image,
        'bookmark_name': bookmark_name,
        'bookmark_btn': bookmark_btn,
        'bookmark_status': bookmark_status,
        'analysis_output': analysis_output,
        'molecule_image': molecule_image
    }


def create_chemical_variations_tab():
    """Create the consolidated drug discovery tab with molecular analysis and variations."""
    # Add notification container
    gr.HTML("""
    <div id="notification-container"></div>
    <script>
    function showNotification(message, type = 'success', duration = 2000) {
        // Remove existing notifications
        const container = document.getElementById('notification-container');
        container.innerHTML = '';
        
        // Create notification element
        const notification = document.createElement('div');
        notification.className = `temp-notification ${type}`;
        notification.textContent = message;
        
        // Add to container
        container.appendChild(notification);
        
        // Auto-remove after duration
        setTimeout(() => {
            if (notification.parentNode) {
                notification.style.opacity = '0';
                notification.style.transform = 'translateX(100%)';
                setTimeout(() => {
                    if (notification.parentNode) {
                        notification.remove();
                    }
                }, 300);
            }
        }, duration);
    }
    
    // Listen for validation status changes
    document.addEventListener('DOMContentLoaded', function() {
        const validationStatus = document.getElementById('validation_status');
        if (validationStatus) {
            const observer = new MutationObserver(function(mutations) {
                mutations.forEach(function(mutation) {
                    if (mutation.type === 'childList' || mutation.type === 'characterData') {
                        const text = validationStatus.textContent || validationStatus.innerText;
                        if (text && text.trim()) {
                            if (text.includes('✅') && text.includes('Ready to analyze')) {
                                showNotification('✅ Ready to analyze', 'success');
                            } else if (text.includes('✅') && text.includes('Valid SMILES')) {
                                showNotification('✅ Valid SMILES', 'success');
                            } else if (text.includes('❌')) {
                                showNotification('❌ Invalid SMILES', 'error');
                            }
                        }
                    }
                });
            });
            observer.observe(validationStatus, { childList: true, characterData: true, subtree: true });
        }
    });
    </script>
    """, visible=False)
    
    with gr.Row():
        with gr.Column(scale=3):
            # Left side: Molecular Analysis
            variation_smiles_input = gr.Textbox(
                label="Enter SMILES string",
                placeholder="e.g., C[C@H](N)C(=O)O",
                info="Enter a valid SMILES string for molecular analysis"
            )
            
            # Hidden validation status for popup notifications
            validation_status = gr.Markdown(
                value="",
                visible=False,
                elem_id="validation_status"
            )
            preview_image = gr.Image(
                label="Preview",
                show_download_button=False,
                width=200,
                height=200,
                visible=False
            )
            
            # Generate variations controls - moved up for better UX
            with gr.Row():
                generate_variations_btn = gr.Button("🔄 Generate Variations", variant="primary")
                clear_variations_btn = gr.Button("🗑️ Clear", variant="secondary")
            
            # Bookmark functionality
            with gr.Row():
                bookmark_name = gr.Textbox(
                    placeholder="Enter molecule name (optional)",
                    label="Molecule Name",
                    scale=2
                )
                bookmark_btn = gr.Button("🔖 Bookmark", variant="secondary", size="sm", scale=1)
            
            bookmark_status = gr.Markdown("")
            
            # Hidden components for internal use (not displayed to user)
            selected_smiles_display = gr.Textbox(
                visible=False,
                elem_id="selected_smiles"
            )
            selected_style_display = gr.Textbox(
                visible=False,
                elem_id="selected_style"
            )
            
            # Enhanced AI Chat Interface
            gr.Markdown("#### 🤖 AI Drug Discovery Assistant")
            gr.Markdown("Ask the AI about molecular properties, generate new structures, or get drug discovery insights")
            
            # AI Settings
            with gr.Row():
                hf_token_input = gr.Textbox(
                    label="Hugging Face Token",
                    type="password",
                    placeholder="Enter your Hugging Face token for AI features",
                    scale=2
                )
                ai_temperature = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.7,
                    step=0.1,
                    label="AI Temperature",
                    scale=1
                )
            
            # Main chat interface
            ai_chatbot = gr.Chatbot(
                label="AI Drug Discovery Chat",
                height=400,
                elem_id="ai_chatbot",
                type="messages"
            )
            
            with gr.Row():
                ai_chat_input = gr.Textbox(
                    placeholder="e.g., 'Explain the drug-likeness of this molecule' or 'Generate a more soluble derivative'",
                    label="Your message",
                    scale=4
                )
                ai_send_btn = gr.Button("Send", variant="primary", scale=1)
            
            # AI Generated Structures display
            gr.Markdown("#### 🧬 AI Generated Structures")
            ai_generated_grid = gr.Gallery(
                label="AI Generated Structures",
                show_label=False,
                elem_id="ai_generated_grid",
                columns=3,
                rows=2,
                height=200,
                object_fit="contain",
                allow_preview=True
            )
        
        with gr.Column(scale=4):
            # Right side: Main structure display + variations grid
            main_structure_display = gr.Image(
                show_download_button=False,
                elem_id="main_structure"
            )
            
            # Seamless properties display
            properties_display = gr.Markdown(
                value="*Click 'Analyze Molecule' or select a variation to see properties*",
                elem_id="properties_display",
                elem_classes="seamless-properties"
            )
            
            # Seamless variations section
            with gr.Column(elem_id="variations_container"):
                variations_grid = gr.Gallery(
                    show_label=False,
                    elem_id="variations_gallery",
                    columns=3,
                    rows=4,
                    height='auto',
                    object_fit="contain",
                    allow_preview=True,
                    show_share_button=False,
                    show_download_button=False
                )
            
            # Navigation controls
            with gr.Row():
                prev_page_btn = gr.Button("⬅️ Previous", size="sm")
                page_info = gr.Markdown("Page 1 of 1", elem_classes="page-info")
                next_page_btn = gr.Button("➡️ Next", size="sm")
            
            # Grid controls
            with gr.Row(elem_classes="grid-controls"):
                grid_size_slider = gr.Slider(
                    minimum=4,
                    maximum=8,
                    value=4,
                    step=1,
                    label="Grid Columns",
                    elem_id="grid_size_slider"
                )
                variation_count_slider = gr.Slider(
                    minimum=6,
                    maximum=24,
                    value=12,
                    step=6,
                    label="Number of Variations",
                    elem_id="variation_count_slider"
                )
    
    return {
        # Molecular Analysis components
        'variation_smiles_input': variation_smiles_input,
        'validation_status': validation_status,
        'preview_image': preview_image,
        'bookmark_name': bookmark_name,
        'bookmark_btn': bookmark_btn,
        'bookmark_status': bookmark_status,
        'main_structure_display': main_structure_display,
        'properties_display': properties_display,
        
        # Variations components
        'generate_variations_btn': generate_variations_btn,
        'clear_variations_btn': clear_variations_btn,
        'selected_smiles_display': selected_smiles_display,
        'selected_style_display': selected_style_display,
        'variations_grid': variations_grid,
        'prev_page_btn': prev_page_btn,
        'page_info': page_info,
        'next_page_btn': next_page_btn,
        'grid_size_slider': grid_size_slider,
        'variation_count_slider': variation_count_slider,
        
        # AI components
        'hf_token_input': hf_token_input,
        'ai_temperature': ai_temperature,
        'ai_chatbot': ai_chatbot,
        'ai_chat_input': ai_chat_input,
        'ai_send_btn': ai_send_btn,
        'ai_generated_grid': ai_generated_grid
    }


def create_molecular_gallery_tab():
    """Create the molecular gallery tab components."""
    gr.Markdown("### Common Drug Discovery Molecules")
    image_components = []
    
    # Display images in a 3x3 grid
    with gr.Row():
        for row in range(3):
            with gr.Row():
                for col in range(3):
                    idx = row * 3 + col
                    img = gr.Image(
                        show_download_button=False,
                        width=200,
                        height=200,
                        label=f"Molecule {idx + 1}",
                    )
                    image_components.append(img)
    
    # Bookmarked molecules section
    gr.Markdown("### 🔖 Your Bookmarked Molecules")
    bookmarked_gallery = gr.Gallery(
        label="Bookmarked Structures",
        show_label=False,
        elem_id="bookmarked_gallery",
        columns=4,
        rows=1,
        height=200,
        object_fit="contain"
    )
    
    return {
        'image_components': image_components,
        'bookmarked_gallery': bookmarked_gallery
    }


def create_drug_library_tab():
    """Create the drug discovery library tab components."""
    gr.Markdown("### NSAID Drug Series")
    drug_images = []
    
    # Display drug molecules
    with gr.Row():
        for i in range(3):  # First row
            img = gr.Image(
                show_download_button=False,
                width=200,
                height=200,
                label=f"Drug {i + 1}",
            )
            drug_images.append(img)
    
    with gr.Row():
        for i in range(2):  # Second row
            img = gr.Image(
                show_download_button=False,
                width=200,
                height=200,
                label=f"Drug {i + 4}",
            )
            drug_images.append(img)
    
    return {
        'drug_images': drug_images
    }


def create_new_experiment_tab():
    """Create a blank workspace tab for new experiments."""
    with gr.Column():
        gr.Markdown("## 🧪 New Experiment Workspace")
        gr.Markdown(
            "This tab is intentionally left blank so you can prototype new "
            "ideas without affecting the existing lab, gallery, or library views."
        )
        gr.Markdown(
            "- Add new components here as you explore ideas\n"
            "- Copy elements from other tabs if needed\n"
            "- Wire up handlers in `src/app.py` once ready"
        )
        gr.HTML(
            "<div style='border:1px dashed #bbb; padding:2rem; text-align:center;'>"
            "Your custom UI goes here"
            "</div>"
        )
    return {}