src/ui/components.py

"""UI components for MelanoScope AI."""
import os
import logging
from typing import List
import gradio as gr
from ..config.settings import UIConfig
from ..core.utils import create_empty_dataframe
from .styles import get_custom_css, create_theme, get_header_html, get_footer_html, get_model_info_html

logger = logging.getLogger(__name__)

class MelanoScopeUI:
    """Handles UI components and layout."""
    
    def __init__(self, model_instance, classes: List[str]):
        self.model = model_instance
        self.classes = classes
        self.theme = create_theme()
        self.css = get_custom_css()
        logger.info("UI initialized")
    
    def create_interface(self) -> gr.Blocks:
        """Create complete Gradio interface."""
        with gr.Blocks(theme=self.theme, css=self.css) as interface:
            self._create_header()
            
            with gr.Row(equal_height=True):
                self._create_input_column()
                self._create_results_column()
            
            self._create_footer()
            self._setup_event_handlers()
        
        return interface
    
    def _create_header(self) -> None:
        """Create header section."""
        with gr.Row():
            with gr.Column(scale=6):
                gr.Markdown(get_header_html())
            with gr.Column(scale=1, min_width=UIConfig.THEME_TOGGLE_MIN_WIDTH):
                try:
                    self.dark_toggle = gr.ThemeMode(label="Theme", value="system")
                except Exception:
                    gr.Markdown("")
    
    def _create_input_column(self) -> None:
        """Create input column with image upload and controls."""
        with gr.Column(scale=UIConfig.LEFT_COLUMN_SCALE):
            self.image_input = gr.Image(
                type="numpy",
                label="Lesion Image",
                height=UIConfig.IMAGE_HEIGHT,
                sources=["upload", "clipboard"]
            )
            
            with gr.Row():
                self.analyze_btn = gr.Button("Analyze", variant="primary")
                self.clear_btn = gr.Button("Clear")
            
            self._create_examples_section()
            self.latency_output = gr.Label(label="Inference Time")
    
    def _create_examples_section(self) -> None:
        """Create examples section if files exist."""
        example_files = [
            "melanoma.jpg", "vascular_lesion.jpg"
        ]
        
        existing_examples = [f for f in example_files if os.path.exists(f)]
        if existing_examples:
            gr.Examples(
                examples=existing_examples,
                inputs=self.image_input,
                label="Quick Examples",
                cache_examples=False
            )

    
    def _create_results_column(self) -> None:
        """Create results column with predictions and info."""
        with gr.Column(scale=UIConfig.RIGHT_COLUMN_SCALE):
            self._create_prediction_results()
            self._create_information_tabs()
    
    def _create_prediction_results(self) -> None:
        """Create prediction results section."""
        with gr.Group():
            with gr.Row():
                self.prediction_output = gr.Label(label="Primary Prediction", elem_classes=["pred-card"])
                self.confidence_output = gr.Label(label="Confidence")
            
            self.probability_plot = gr.BarPlot(
                value=create_empty_dataframe(self.classes),
                x="item", y="probability",
                title="Probability Distribution (Top-k)",
                x_title="Class", y_title="Probability",
                vertical=False, tooltip=["item", "probability"],
                width=UIConfig.PLOT_WIDTH, height=UIConfig.PLOT_HEIGHT,
            )
    
    def _create_information_tabs(self) -> None:
        """Create information tabs."""
        with gr.Tabs():
            with gr.TabItem("Medical Details"):
                self._create_medical_details()
            with gr.TabItem("About Model"):
                gr.Markdown(get_model_info_html())
    
    def _create_medical_details(self) -> None:
        """Create medical details accordions."""
        with gr.Accordion("Description", open=True):
            self.description_output = gr.Textbox(lines=UIConfig.TEXTBOX_LINES, interactive=False)
        with gr.Accordion("Symptoms", open=False):
            self.symptoms_output = gr.Textbox(lines=UIConfig.TEXTBOX_LINES, interactive=False)
        with gr.Accordion("Causes", open=False):
            self.causes_output = gr.Textbox(lines=UIConfig.TEXTBOX_LINES, interactive=False)
        with gr.Accordion("Treatment", open=False):
            self.treatment_output = gr.Textbox(lines=UIConfig.TEXTBOX_LINES, interactive=False)
    
    def _create_footer(self) -> None:
        """Create footer section."""
        gr.Markdown(get_footer_html())
    
    def _setup_event_handlers(self) -> None:
        """Set up event handlers."""
        outputs = [
            self.prediction_output, self.confidence_output, self.description_output,
            self.symptoms_output, self.causes_output, self.treatment_output,
            self.probability_plot, self.latency_output
        ]
        
        self.analyze_btn.click(
            fn=self.model.predict, inputs=[self.image_input], 
            outputs=outputs, show_progress="full"
        )
        
        self.clear_btn.click(
            fn=self._clear_all, inputs=[], 
            outputs=[self.image_input] + outputs
        )
    
    def _clear_all(self) -> tuple:
        """Clear all inputs and outputs."""
        empty_df = create_empty_dataframe(self.classes)
        return (None, "", "", "", "", "", "", empty_df, "")