import streamlit as st
from PIL import Image
import numpy as np
from model_utils import BugClassifier, get_severity_prediction
from transformers import AutoFeatureExtractor

# Page configuration
st.set_page_config(
    page_title="Bug-O-Scope 🔍🐞",
    page_icon="🔍",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
@st.cache_resource
def load_model():
    try:
        return BugClassifier(), AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None

if 'model' not in st.session_state:
    st.session_state.model, st.session_state.feature_extractor = load_model()

def main():
    # Header
    st.title("Bug-O-Scope 🔍🐞")
    st.markdown("""
    Welcome to Bug-O-Scope! Upload a picture of an insect to learn more about it.
    This educational tool helps you identify bugs and understand their role in our ecosystem.
    """)

    # Sidebar
    st.sidebar.header("About Bug-O-Scope")
    st.sidebar.markdown("""
    Bug-O-Scope is an AI-powered tool that helps you:
    * 🔍 Identify insects from photos
    * 📚 Learn about different species
    * 🌍 Understand their ecological impact
    * 🔬 Compare different insects
    """)

    # Main content
    tab1, tab2 = st.tabs(["Single Bug Analysis", "Bug Comparison"])

    with tab1:
        single_bug_analysis()

    with tab2:
        compare_bugs()

def single_bug_analysis():
    """Handle single bug analysis"""
    uploaded_file = st.file_uploader("Upload a bug photo", type=['png', 'jpg', 'jpeg'], key="single")
    
    if uploaded_file:
        try:
            image = Image.open(uploaded_file)
            col1, col2 = st.columns(2)
            
            with col1:
                st.image(image, caption="Uploaded Image", use_container_width=True)
            
            with col2:
                with st.spinner("Analyzing your bug..."):
                    # Get predictions
                    prediction, confidence = st.session_state.model.predict(image)
                    severity = get_severity_prediction(prediction)
                    
                    st.success("Analysis Complete!")
                    st.markdown(f"### Identified Species")
                    st.markdown(f"**{prediction}**")
                    st.markdown(f"Confidence: {confidence:.2f}%")
                    
                    st.markdown("### Ecological Impact")
                    severity_color = {
                        "Low": "green",
                        "Medium": "orange",
                        "High": "red"
                    }
                    st.markdown(
                        f"Severity: <span style='color: {severity_color[severity]}'>{severity}</span>", 
                        unsafe_allow_html=True
                    )

            # Generate and display species information
            st.markdown("### About This Species")
            species_info = st.session_state.model.get_species_info(prediction)
            st.markdown(species_info)
            
            # Display Grad-CAM visualization
            st.markdown("### Feature Highlights")
            gradcam = st.session_state.model.get_gradcam(image)
            st.image(gradcam, caption="Important Features", use_container_width=True)
            
        except Exception as e:
            st.error(f"Error processing image: {str(e)}")
            st.info("Please try uploading a different image.")

def compare_bugs():
    """Handle bug comparison"""
    col1, col2 = st.columns(2)
    
    with col1:
        file1 = st.file_uploader("Upload first bug photo", type=['png', 'jpg', 'jpeg'], key="compare1")
        if file1:
            try:
                image1 = Image.open(file1)
                st.image(image1, caption="First Bug", use_container_width=True)
            except Exception as e:
                st.error(f"Error loading first image: {str(e)}")
                return
            
    with col2:
        file2 = st.file_uploader("Upload second bug photo", type=['png', 'jpg', 'jpeg'], key="compare2")
        if file2:
            try:
                image2 = Image.open(file2)
                st.image(image2, caption="Second Bug", use_container_width=True)
            except Exception as e:
                st.error(f"Error loading second image: {str(e)}")
                return
    
    if file1 and file2:
        try:
            with st.spinner("Generating comparison..."):
                # Get predictions for both images
                pred1, conf1 = st.session_state.model.predict(image1)
                pred2, conf2 = st.session_state.model.predict(image2)
                
                # Generate Grad-CAM visualizations
                gradcam1 = st.session_state.model.get_gradcam(image1)
                gradcam2 = st.session_state.model.get_gradcam(image2)
                
                # Display results
                st.markdown("### Comparison Results")
                comp_col1, comp_col2 = st.columns(2)
                
                with comp_col1:
                    st.markdown(f"**Species 1**: {pred1}")
                    st.markdown(f"Confidence: {conf1:.2f}%")
                    st.image(gradcam1, caption="Feature Highlights - Bug 1", use_container_width=True)
                    
                with comp_col2:
                    st.markdown(f"**Species 2**: {pred2}")
                    st.markdown(f"Confidence: {conf2:.2f}%")
                    st.image(gradcam2, caption="Feature Highlights - Bug 2", use_container_width=True)
                
                # Display comparison information
                st.markdown("### Key Differences")
                differences = st.session_state.model.compare_species(pred1, pred2)
                st.markdown(differences)
                
        except Exception as e:
            st.error(f"Error comparing images: {str(e)}")
            st.info("Please try uploading different images or try again.")

if __name__ == "__main__":
    main()