import streamlit as st
from PIL import Image
from model_utils import BugClassifier, get_severity_prediction

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

# Initialize session state for model
@st.cache_resource
def load_model():
    try:
        print("Loading model...")
        model = BugClassifier()
        print("Model loaded successfully")
        return model
    except Exception as e:
        print(f"Error loading model: {str(e)}")
        return None

# Ensure model is loaded
if 'model' not in st.session_state:
    st.session_state.model = 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
    """)

    # Check if model loaded successfully
    if st.session_state.model is None:
        st.error("Error: Model failed to load. Please try refreshing the page.")
        return

    # 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:
            # Load and display image
            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)
                    print(f"Prediction: {prediction}, Confidence: {confidence}")
                    
                    st.success("Analysis Complete!")
                    st.markdown("### Identified Species")
                    st.markdown(f"**{prediction}**")
                    st.markdown(f"Confidence: {confidence:.2f}%")
                    
                    # Only show ecological impact for known insects
                    if prediction != "Unknown Insect" and prediction != "Error Processing Image":
                        severity = get_severity_prediction(prediction)
                        st.markdown("### Ecological Impact")
                        severity_color = {
                            "Low": "green",
                            "Medium": "orange",
                            "High": "red",
                            "Unknown": "gray"
                        }
                        st.markdown(
                            f"Severity: <span style='color: {severity_color[severity]}'>{severity}</span>", 
                            unsafe_allow_html=True
                        )

            # Display species information
            if prediction != "Unknown Insect" and prediction != "Error Processing Image":
                st.markdown("### About This Species")
                species_info = st.session_state.model.get_species_info(prediction)
                st.markdown(species_info)
                
                # Display 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
                pred1, conf1 = st.session_state.model.predict(image1)
                pred2, conf2 = st.session_state.model.predict(image2)
                
                if pred1 not in ["Unknown Insect", "Error Processing Image"] and \
                   pred2 not in ["Unknown Insect", "Error Processing Image"]:
                    
                    # 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}%")
                        gradcam1 = st.session_state.model.get_gradcam(image1)
                        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}%")
                        gradcam2 = st.session_state.model.get_gradcam(image2)
                        st.image(gradcam2, caption="Feature Highlights - Bug 2", use_container_width=True)
                    
                    # Display comparison
                    st.markdown("### Key Differences")
                    st.markdown(st.session_state.model.get_species_info(pred1))
                    st.markdown(st.session_state.model.get_species_info(pred2))
                else:
                    st.warning("Unable to generate meaningful comparison due to low confidence predictions.")
                
        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()