"""
Gradio Web Interface for Milk Spoilage Classification

This app provides an interactive web interface for predicting
milk spoilage type based on microbial count data.
"""

import gradio as gr
import joblib
import numpy as np


# Load the trained model
model = joblib.load("model.joblib")

# Feature information for the UI
FEATURE_INFO = {
    "SPC_D7": ("Standard Plate Count - Day 7", "log CFU/mL", 0.0, 10.0, 4.0),
    "SPC_D14": ("Standard Plate Count - Day 14", "log CFU/mL", 0.0, 10.0, 5.0),
    "SPC_D21": ("Standard Plate Count - Day 21", "log CFU/mL", 0.0, 10.0, 6.0),
    "TGN_D7": ("Total Gram-Negative - Day 7", "log CFU/mL", 0.0, 10.0, 3.0),
    "TGN_D14": ("Total Gram-Negative - Day 14", "log CFU/mL", 0.0, 10.0, 4.0),
    "TGN_D21": ("Total Gram-Negative - Day 21", "log CFU/mL", 0.0, 10.0, 5.0),
}

# Class descriptions
CLASS_DESCRIPTIONS = {
    "PPC": "Post-Pasteurization Contamination - Bacteria introduced after pasteurization",
    "no spoilage": "No significant spoilage detected in the sample",
    "spore spoilage": "Spoilage caused by spore-forming bacteria"
}


def predict_spoilage(spc_d7, spc_d14, spc_d21, tgn_d7, tgn_d14, tgn_d21):
    """
    Predict milk spoilage type based on microbial counts.
    
    Args:
        spc_d7: Standard Plate Count at Day 7
        spc_d14: Standard Plate Count at Day 14
        spc_d21: Standard Plate Count at Day 21
        tgn_d7: Total Gram-Negative count at Day 7
        tgn_d14: Total Gram-Negative count at Day 14
        tgn_d21: Total Gram-Negative count at Day 21
    
    Returns:
        Dictionary of class probabilities for Gradio Label component
    """
    # Prepare input features
    features = np.array([[spc_d7, spc_d14, spc_d21, tgn_d7, tgn_d14, tgn_d21]])
    
    # Get prediction and probabilities
    prediction = model.predict(features)[0]
    probabilities = model.predict_proba(features)[0]
    
    # Create probability dictionary for Gradio Label
    prob_dict = {
        cls: float(prob) 
        for cls, prob in zip(model.classes_, probabilities)
    }
    
    return prob_dict


def create_interface():
    """Create and configure the Gradio interface."""
    
    # Custom CSS for styling
    custom_css = """
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    }
    .feature-group {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        border-radius: 10px;
        padding: 15px;
        margin: 10px 0;
    }
    """
    
    with gr.Blocks(
        title="Milk Spoilage Classifier",
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="purple",
        ),
        css=custom_css
    ) as demo:
        
        # Header
        gr.Markdown(
            """
            # 🥛 Milk Spoilage Classification Model
            
            Predict milk spoilage type based on microbial count data measured at different time points.
            Enter the Standard Plate Count (SPC) and Total Gram-Negative (TGN) values below.
            """
        )
        
        with gr.Row():
            # Input Section
            with gr.Column(scale=1):
                gr.Markdown("### 📊 Standard Plate Count (SPC)")
                gr.Markdown("*Total bacterial count in log CFU/mL*")
                
                spc_d7 = gr.Number(
                    label="Day 7",
                    value=4.0,
                    minimum=0.0,
                    maximum=10.0,
                    info="SPC measurement at day 7"
                )
                spc_d14 = gr.Number(
                    label="Day 14",
                    value=5.0,
                    minimum=0.0,
                    maximum=10.0,
                    info="SPC measurement at day 14"
                )
                spc_d21 = gr.Number(
                    label="Day 21",
                    value=6.0,
                    minimum=0.0,
                    maximum=10.0,
                    info="SPC measurement at day 21"
                )
            
            with gr.Column(scale=1):
                gr.Markdown("### 🦠 Total Gram-Negative (TGN)")
                gr.Markdown("*Gram-negative bacterial count in log CFU/mL*")
                
                tgn_d7 = gr.Number(
                    label="Day 7",
                    value=3.0,
                    minimum=0.0,
                    maximum=10.0,
                    info="TGN measurement at day 7"
                )
                tgn_d14 = gr.Number(
                    label="Day 14",
                    value=4.0,
                    minimum=0.0,
                    maximum=10.0,
                    info="TGN measurement at day 14"
                )
                tgn_d21 = gr.Number(
                    label="Day 21",
                    value=5.0,
                    minimum=0.0,
                    maximum=10.0,
                    info="TGN measurement at day 21"
                )
        
        # Predict button
        predict_btn = gr.Button("🔬 Classify Spoilage Type", variant="primary", size="lg")
        
        # Output Section
        gr.Markdown("### 📋 Prediction Results")
        
        output_label = gr.Label(
            label="Spoilage Classification",
            num_top_classes=3
        )
        
        # Connect the prediction function
        predict_btn.click(
            fn=predict_spoilage,
            inputs=[spc_d7, spc_d14, spc_d21, tgn_d7, tgn_d14, tgn_d21],
            outputs=output_label
        )
        
        # Also trigger on any input change
        for input_component in [spc_d7, spc_d14, spc_d21, tgn_d7, tgn_d14, tgn_d21]:
            input_component.change(
                fn=predict_spoilage,
                inputs=[spc_d7, spc_d14, spc_d21, tgn_d7, tgn_d14, tgn_d21],
                outputs=output_label
            )
        
        # Information Section
        gr.Markdown(
            """
            ---
            ### ℹ️ About the Classes
            
            | Class | Description |
            |-------|-------------|
            | **PPC** | Post-Pasteurization Contamination - Bacteria introduced after pasteurization process |
            | **no spoilage** | No significant spoilage detected in the sample |
            | **spore spoilage** | Spoilage caused by spore-forming bacteria that survive pasteurization |
            
            ---
            ### 📖 How to Use
            
            1. Enter the microbial count values (in log CFU/mL) for each time point
            2. Click "Classify Spoilage Type" or wait for automatic prediction
            3. View the predicted spoilage category and confidence scores
            
            ---
            *Model: Random Forest Classifier trained on milk quality data*
            """
        )
    
    return demo


# Create and launch the interface
if __name__ == "__main__":
    demo = create_interface()
    demo.launch()