import streamlit as st
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np
import time
import os
import json

# Set page configuration
st.set_page_config(
    page_title="Spider Mite Detection",
    page_icon="🍃",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'Get Help': 'https://www.github.com/your-repo',
        'Report a bug': 'https://www.github.com/your-repo/issues',
        'About': 'Advanced AI system for detecting spider mite infestations on plant leaves'
    }
)

# Define model architectures
MODEL_MAP = {
    'mobilenetv3': {
        'model_fn': models.mobilenet_v3_small,
        'classifier_update': lambda model, num_classes: setattr(
            model, 'classifier', nn.Sequential(
                *list(model.classifier.children())[:-1], 
                nn.Linear(model.classifier[-1].in_features, num_classes)
            )
        )
    },
    'efficientnet': {
        'model_fn': models.efficientnet_b0,
        'classifier_update': lambda model, num_classes: setattr(
            model, 'classifier', nn.Sequential(
                *list(model.classifier.children())[:-1], 
                nn.Linear(model.classifier[-1].in_features, num_classes)
            )
        )
    },
    'resnet18': {
        'model_fn': models.resnet18,
        'classifier_update': lambda model, num_classes: setattr(
            model, 'fc', nn.Linear(model.fc.in_features, num_classes)
        )
    }
}

# Define image transformation
def get_transform():
    return transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

# Define function to load model
@st.cache_resource
def load_model(model_path="models/model.pth"):
    # Try to find model_info.json to determine which model architecture to use
    model_dir = os.path.dirname(model_path)
    model_info_path = os.path.join(model_dir, "model_info.json")
    
    # Default model type if info file doesn't exist
    model_type = "mobilenetv3"
    
    if os.path.exists(model_info_path):
        try:
            with open(model_info_path, 'r') as f:
                model_info = json.load(f)
                model_type = model_info.get('model_name', model_type)
        except:
            st.warning("Couldn't read model info file. Using default model type.")
    
    # Initialize the model
    if model_type not in MODEL_MAP:
        st.error(f"Unknown model type: {model_type}")
        return None
        
    model = MODEL_MAP[model_type]['model_fn'](weights=None)
    MODEL_MAP[model_type]['classifier_update'](model, 2)  # 2 classes: healthy, infested
    
    # Load weights
    if os.path.exists(model_path):
        try:
            model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
            model.eval()
            return model
        except Exception as e:
            st.error(f"Error loading model: {e}")
            return None
    else:
        st.error(f"Model file not found: {model_path}")
        return None

# Predict function
def predict(model, image):
    transform = get_transform()
    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
    
    # Make prediction
    with torch.no_grad():
        start_time = time.time()
        outputs = model(image_tensor)
        inference_time = time.time() - start_time
        
        # Get probabilities
        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
        
    return probabilities.numpy(), inference_time

# Main app UI
def main():
    # Custom CSS
    st.markdown("""
        <style>
        .main-header {
            font-size: 3rem;
            background: linear-gradient(90deg, #4CAF50, #2196F3);
            -webkit-background-clip: text;
            background-clip: text;
            color: transparent;
            text-align: center;
            margin-bottom: 2rem;
            padding: 20px 0;
            font-weight: 800;
            letter-spacing: 1px;
            text-shadow: 0px 2px 4px rgba(0,0,0,0.1);
        }
        .sub-header {
            font-size: 1.8rem;
            color: #1E88E5;
            margin-bottom: 1.5rem;
            border-bottom: 2px solid #E0E0E0;
            padding-bottom: 10px;
        }
        .result-header {
            font-size: 2rem;
            font-weight: bold;
            margin-bottom: 1.5rem;
            color: #333;
        }
        .healthy {
            color: #2E7D32;
            font-weight: bold;
            font-size: 1.5rem;
            display: flex;
            align-items: center;
            background-color: rgba(46, 125, 50, 0.1);
            padding: 10px 15px;
            border-radius: 8px;
        }
        .infested {
            color: #C62828;
            font-weight: bold;
            font-size: 1.5rem;
            display: flex;
            align-items: center;
            background-color: rgba(198, 40, 40, 0.1);
            padding: 10px 15px;
            border-radius: 8px;
        }
        .sidebar-content {
            font-size: 1.1rem;
            padding: 15px 0;
        }
        .sidebar-title {
            background: linear-gradient(90deg, #4CAF50, #2196F3);
            -webkit-background-clip: text;
            background-clip: text;
            color: transparent;
            font-weight: 700;
            margin-bottom: 15px;
        }
        .stProgress > div > div {
            background-color: #4CAF50;
        }
        div[data-testid="stFileUploadDropzone"] {
            border: 2px dashed #4CAF50;
            border-radius: 8px;
            padding: 30px 20px;
            background-color: rgba(76, 175, 80, 0.05);
            margin-bottom: 25px;
            transition: all 0.3s ease;
        }
        div[data-testid="stFileUploadDropzone"]:hover {
            background-color: rgba(76, 175, 80, 0.1);
            border-color: #2E7D32;
        }
        .stButton>button {
            background-color: #2196F3;
            color: white;
            border-radius: 5px;
            border: none;
            padding: 10px 20px;
            font-weight: 600;
            transition: all 0.3s ease;
        }
        .stButton>button:hover {
            background-color: #1976D2;
            transform: translateY(-2px);
            box-shadow: 0 5px 15px rgba(0,0,0,0.1);
        }
        .info-box {
            background-color: #E3F2FD;
            border-left: 5px solid #2196F3;
            padding: 15px;
            border-radius: 5px;
            margin-bottom: 20px;
        }
        .metrics-container {
            display: grid;
            grid-template-columns: 1fr 1fr;
            gap: 15px;
            margin-bottom: 20px;
        }
        .metric-card {
            background-color: #f8f9fa;
            border-radius: 8px;
            padding: 15px;
            box-shadow: 0 2px 5px rgba(0,0,0,0.1);
        }
        .metric-value {
            font-size: 1.8rem;
            font-weight: bold;
            color: #1976D2;
        }
        .metric-label {
            color: #5f6368;
            font-size: 0.9rem;
            margin-top: 5px;
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Sidebar information
    st.sidebar.markdown('<div class="sidebar-content">', unsafe_allow_html=True)
    st.sidebar.image("https://www.gardeningknowhow.com/wp-content/uploads/2019/08/spider-mite-1.jpg", width=280)
    st.sidebar.markdown('<h2 class="sidebar-title">About This Tool</h2>', unsafe_allow_html=True)
    st.sidebar.markdown("""
    <div class="info-box">
    This application uses deep learning to detect spider mite infestations on plant leaves with high accuracy.
    </div>
    
    ### What are Spider Mites?
    - Tiny arachnids (0.5mm) that damage crops
    - Feed on plant tissues, causing yellowing and spotting
    - Can rapidly reproduce and spread throughout plants
    - Often go unnoticed until significant damage occurs
    
    ### How to Use This Tool
    1. Upload a high-quality image of a plant leaf
    2. Our AI will analyze the image for infestation signs
    3. Review the detection results and follow recommendations
    
    ### Model Information
    - Technology: Deep Learning with Transfer Learning
    - Architecture: MobileNetV3
    - Accuracy: ~95%+ on validation data
    - Training: Custom dataset of healthy and infested leaves
    """, unsafe_allow_html=True)
    st.sidebar.markdown('</div>', unsafe_allow_html=True)
    
    # Main area
    st.markdown('<h1 class="main-header">Spider Mite Detection System</h1>', unsafe_allow_html=True)
    
    # Tabs for different functionalities
    tab1, tab2 = st.tabs(["🔍 Detect Spider Mites", "ℹ️ About the Project"])
    
    with tab1:
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            st.markdown('<h2 class="sub-header">Upload a leaf image for analysis</h2>', unsafe_allow_html=True)
        
        # Load model
        with st.spinner("Loading AI model..."):
            model = load_model()
        
        if model is None:
            st.error("⚠️ Failed to load model. Please check if the model file exists.")
            return
        else:
            st.success("✅ AI model loaded successfully")
        
        # Image upload with enhanced UI
        upload_container = st.container()
        with upload_container:
            uploaded_file = st.file_uploader("Choose a leaf image to analyze", type=["jpg", "jpeg", "png"])
            
            # Add test images from test_images folder
            st.markdown("<h3>Or select a test image:</h3>", unsafe_allow_html=True)
            test_cols = st.columns(4)
            
            # Get test images from test_images folder
            test_images = []
            for root, dirs, files in os.walk("test_images"):
                for file in files:
                    if file.lower().endswith((".jpg", ".jpeg", ".png")):
                        test_images.append({"name": file, "path": os.path.join(root, file)})
            
            # Display test images
            use_test = None
            for i, test_img in enumerate(test_images):
                with test_cols[i % 4]:
                    folder_type = "Healthy" if "healthy" in test_img["path"].lower() else "Spider Mite Infested"
                    st.image(test_img["path"], caption=folder_type, width=150)
                    if st.button(f"Use {folder_type}", key=f"test_button_{i}"):
                        use_test = test_img["path"]
                        uploaded_file = use_test
        
        # Process image if uploaded
        if uploaded_file is not None:
            # Handle test image path
            if isinstance(uploaded_file, str):
                try:
                    image = Image.open(uploaded_file).convert('RGB')
                except:
                    st.error(f"Could not open test image: {uploaded_file}")
                    return
            st.markdown("---")
            
            analysis_container = st.container()
            with analysis_container:
                col1, col2 = st.columns([1, 1])
                
                with col1:
                    st.markdown('<h3 class="sub-header">Uploaded Image</h3>', unsafe_allow_html=True)
                    image = Image.open(uploaded_file).convert('RGB')
                    st.image(image, caption="", use_container_width=True)
                
                with col2:
                    st.markdown('<h3 class="result-header">Analysis Result</h3>', unsafe_allow_html=True)
                    
                    # Run prediction
                    with st.spinner("🔍 Analyzing leaf image..."):
                        probabilities, inference_time = predict(model, image)
                        
                        # Class names
                        class_names = ['Healthy', 'Spider Mite Infested']
                        
                        # Get prediction
                        predicted_class = np.argmax(probabilities)
                        confidence = probabilities[predicted_class] * 100
                        
                        # Display result with improved UI
                        if predicted_class == 0:
                            st.markdown(f'<div class="healthy">✅ {class_names[predicted_class]}</div>', unsafe_allow_html=True)
                            emoji = "🌿"
                            result_color = "#2E7D32"
                        else:
                            st.markdown(f'<div class="infested">⚠️ {class_names[predicted_class]}</div>', unsafe_allow_html=True)
                            emoji = "🕷️"
                            result_color = "#C62828"
                        
                        # Result metrics in a nice grid
                        st.markdown('<div class="metrics-container">', unsafe_allow_html=True)
                        st.markdown(f"""
                        <div class="metric-card">
                            <div class="metric-value">{confidence:.1f}%</div>
                            <div class="metric-label">Confidence</div>
                        </div>
                        <div class="metric-card">
                            <div class="metric-value">{inference_time*1000:.0f}ms</div>
                            <div class="metric-label">Processing Time</div>
                        </div>
                        """, unsafe_allow_html=True)
                        st.markdown('</div>', unsafe_allow_html=True)
                        
                        # Progress bars for confidence scores
                        st.markdown("### Detection Confidence")
                        for i, class_name in enumerate(class_names):
                            # Convert float32 to Python float to fix the error
                            prob_value = float(probabilities[i])
                            prob_pct = prob_value * 100
                            emoji_prefix = "🌿" if i == 0 else "🕷️"
                            st.progress(prob_value)
                            st.markdown(f"{emoji_prefix} **{class_name}:** {prob_pct:.1f}%")
                        
                        # Recommendation with enhanced styling
                        st.markdown("### Recommendation")
                        if predicted_class == 0:
                            st.success("✅ No action needed. Your plant appears healthy!")
                        else:
                            if confidence > 90:
                                st.error("""
                                🚨 **Immediate treatment recommended!**
                                
                                High probability of spider mite infestation detected.
                                """)
                                
                                st.markdown("""
                                <div style="background-color: #fff3e0; padding: 15px; border-radius: 5px; border-left: 5px solid #c62828;">
                                <h4 style="color: #c62828; margin-top: 0;">Treatment options:</h4>
                                <ul style="color: #333;">
                                <li><strong>Natural remedies:</strong> Apply neem oil or insecticidal soap</li>
                                <li><strong>Biological control:</strong> Introduce predatory mites</li>
                                <li><strong>Management:</strong> Prune heavily infested leaves</li>
                                <li><strong>Prevention:</strong> Increase humidity and monitor regularly</li>
                                </ul>
                                </div>
                                """, unsafe_allow_html=True)
                            else:
                                st.warning("""
                                ⚠️ **Potential infestation detected.**
                                
                                Monitor your plant closely and consider preventative treatment.
                                """)
        else:
            # Display placeholder when no image is uploaded
            st.info("👆 Upload a leaf image to get started with the analysis.")
            
            # Add a placeholder image with instructions
            col1, col2, col3 = st.columns([1, 2, 1])
            with col2:
                st.markdown("""
                <div style="text-align: center; padding: 40px; background-color: #f5f5f5; border-radius: 10px; margin: 20px 0;">
                    <img src="https://www.planetnatural.com/wp-content/uploads/2013/01/spider-mite-control.jpg" style="max-width: 80%; border-radius: 8px; margin-bottom: 20px;" />
                    <p style="font-size: 1.2rem; color: #555;">Upload a clear image of your plant leaf to detect spider mite infestations</p>
                </div>
                """, unsafe_allow_html=True)
    
    with tab2:
        col1, col2 = st.columns([3, 2])
        
        with col1:
            st.markdown('<h2 class="sub-header">About This Project</h2>', unsafe_allow_html=True)
            st.markdown("""
            ## Spider Mite Detection Using AI
            
            Spider mites are tiny pests that cause significant damage to crops worldwide. Early detection is crucial to 
            prevent severe crop damage and yield loss. Our AI-powered detection system helps farmers and gardeners
            identify infestations before they become severe.
            
            ### Project Goals
            1. Develop an AI model capable of classifying leaves as infested or healthy
            2. Achieve high accuracy (>90%) in detection
            3. Create an accessible application for farmers to use
            4. Help reduce crop losses due to spider mite damage
            
            ### Technology Used
            - **Deep Learning**: Transfer learning with modern CNN architectures
            - **Model Architectures**: MobileNetV3, EfficientNet, ResNet18
            - **Training Data**: Curated dataset of healthy and infested plant leaves
            - **Web Application**: Built with Streamlit for accessibility
            
            ### Team Members
            - Nitesh Kumar Datha Vemanapall
            - Jithin Garapati
            - Pavan Sai Korlapati
            
            ### Future Improvements
            - Mobile application for in-field use
            - Multi-class detection for various plant diseases
            - Integration with automated spraying systems
            - Expanded dataset for more plant species
            """)
        
        with col2:
            st.image("https://www.cropscience.bayer.co.uk/-/media/bcs-inter/ws_uk/images/article-images/pest-encyclopedia/two-spotted-spider-mite.jpg", caption="Spider mite damage on leaves", use_column_width=True)
            st.image("https://www.planetnatural.com/wp-content/uploads/2013/01/spider-mite-control.jpg", caption="Spider mite close-up", use_column_width=True)
            
            # Add some statistics
            st.markdown("""
            <div style="background-color: #e8f5e9; padding: 20px; border-radius: 10px; margin-top: 20px;">
            <h3 style="color: #2e7d32; margin-top: 0;">Spider Mite Impact</h3>
            <ul>
            <li><strong>Up to 60%</strong> crop yield reduction in severe infestations</li>
            <li><strong>$1+ billion</strong> in annual agricultural losses worldwide</li>
            <li><strong>200+ plant species</strong> are vulnerable to spider mite attacks</li>
            <li><strong>95% accuracy</strong> achieved by our detection system</li>
            </ul>
            </div>
            """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()