import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import json
import plotly.graph_objects as go
from datetime import datetime
import pandas as pd

# Load class indices
with open("class_indices.json", "r") as f:
    class_indices = json.load(f)
    
# Reverse the mapping for predictions
class_names = {v: k for k, v in class_indices.items()}

# Load the TFLite model
@st.cache_resource
def load_model():
    interpreter = tf.lite.Interpreter(model_path="model.tflite")
    interpreter.allocate_tensors()
    return interpreter

interpreter = load_model()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Define the image preprocessing function
def preprocess_image(image, target_size=(224, 224)):
    image = image.resize(target_size)
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    return image.astype(np.float32)

# Define prediction function with detailed output
def predict(image):
    input_data = preprocess_image(image)
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])
    
    # Get top 3 predictions
    top_indices = np.argsort(output_data[0])[-3:][::-1]
    predictions = []
    for idx in top_indices:
        predictions.append({
            'class': class_names[idx],
            'confidence': float(output_data[0][idx])
        })
    return predictions

# Custom CSS
st.set_page_config(
    page_title="🌿 Smart Crop Disease Detective",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Custom CSS
st.markdown("""
    <style>
        .main {
            background-color: #f5f7f9;
        }
        .stButton>button {
            background-color: #2d6a4f;
            color: white;
            border-radius: 10px;
            padding: 0.5rem 1rem;
            border: none;
            width: 100%;
        }
        .stButton>button:hover {
            background-color: #40916c;
            border: none;
        }
        .prediction-box {
            background-color: white;
            padding: 20px;
            border-radius: 10px;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .info-box {
            background-color: #e9ecef;
            padding: 15px;
            border-radius: 5px;
            margin: 10px 0;
        }
        .status-box {
            padding: 10px;
            border-radius: 5px;
            margin: 10px 0;
            text-align: center;
        }
        .header-container {
            display: flex;
            align-items: center;
            justify-content: space-between;
            padding: 1rem;
            background-color: white;
            border-radius: 10px;
            margin-bottom: 2rem;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
    </style>
""", unsafe_allow_html=True)

# Session State initialization
if 'prediction_history' not in st.session_state:
    st.session_state.prediction_history = []

# App Header
col1, col2, col3 = st.columns([1,2,1])
with col2:
    st.markdown("""
        <div style='text-align: center'>
            <h1 style='color: #2d6a4f'>🌿 Smart Crop Disease Detective</h1>
            <p style='color: #40916c; font-size: 1.2em;'>
                Your AI-Powered Assistant for Crop Health Monitoring
            </p>
        </div>
    """, unsafe_allow_html=True)

# Sidebar
with st.sidebar:
    st.image("https://via.placeholder.com/250x150?text=Smart+Crop+AI", use_column_width=True)
    
    st.markdown("### 📊 Dashboard")
    total_scans = len(st.session_state.prediction_history)
    st.metric("Total Scans", total_scans)
    
    st.markdown("### 🎯 Features")
    st.markdown("""
        - 🔍 Real-time disease detection
        - 📊 Confidence scoring
        - 📈 Multiple disease possibilities
        - 💾 Scan history tracking
        - 🌱 Treatment recommendations
    """)
    
    st.markdown("### 💡 Tips for Best Results")
    st.info("""
        1. Ensure good lighting
        2. Focus on affected areas
        3. Avoid blurry images
        4. Include multiple angles
        5. Clean lens before capture
    """)
    
    if st.button("Clear History"):
        st.session_state.prediction_history = []
        st.success("History cleared!")

# Main Content
main_col1, main_col2 = st.columns([2,3])

with main_col1:
    st.markdown("### 📸 Upload Image")
    uploaded_file = st.file_uploader(
        "Choose a leaf image (JPG/PNG)",
        type=["jpg", "png", "jpeg"],
        help="Upload a clear image of the affected crop leaf"
    )
    
    if uploaded_file:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        analyze_btn = st.button("🔍 Analyze Image")
        if analyze_btn:
            with st.spinner("🔄 Analyzing image..."):
                predictions = predict(image)
                
                # Store prediction in history
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                st.session_state.prediction_history.append({
                    'timestamp': timestamp,
                    'predictions': predictions,
                    'filename': uploaded_file.name
                })

with main_col2:
    if uploaded_file and analyze_btn:
        st.markdown("### 🔍 Analysis Results")
        
        # Display confidence gauge for top prediction
        fig = go.Figure(go.Indicator(
            mode = "gauge+number",
            value = predictions[0]['confidence'] * 100,
            domain = {'x': [0, 1], 'y': [0, 1]},
            title = {'text': "Confidence Level"},
            gauge = {
                'axis': {'range': [None, 100]},
                'bar': {'color': "#2d6a4f"},
                'steps': [
                    {'range': [0, 50], 'color': "#ff9999"},
                    {'range': [50, 75], 'color': "#ffff99"},
                    {'range': [75, 100], 'color': "#99ff99"}
                ]
            }
        ))
        st.plotly_chart(fig)
        
        # Display predictions
        for i, pred in enumerate(predictions, 1):
            confidence_color = (
                "#ff0000" if pred['confidence'] < 0.5
                else "#ffa500" if pred['confidence'] < 0.7
                else "#008000"
            )
            
            st.markdown(f"""
                <div class="prediction-box">
                    <h4>Prediction {i}: {pred['class']}</h4>
                    <p style='color: {confidence_color}'>
                        Confidence: {pred['confidence']*100:.2f}%
                    </p>
                </div>
            """, unsafe_allow_html=True)
        
        # Treatment Recommendations (example)
        st.markdown("### 💊 Treatment Recommendations")
        st.markdown(f"""
            <div class="info-box">
                <h4>For {predictions[0]['class']}:</h4>
                <ul>
                    <li>Isolate affected plants</li>
                    <li>Apply appropriate fungicide/pesticide</li>
                    <li>Improve air circulation</li>
                    <li>Monitor moisture levels</li>
                </ul>
                <p><em>Consult with a local agricultural expert for specific treatment plans.</em></p>
            </div>
        """, unsafe_allow_html=True)

# History Section
if st.session_state.prediction_history:
    st.markdown("### 📜 Scan History")
    history_df = pd.DataFrame([
        {
            'Timestamp': h['timestamp'],
            'Filename': h['filename'],
            'Primary Prediction': h['predictions'][0]['class'],
            'Confidence': f"{h['predictions'][0]['confidence']*100:.2f}%"
        }
        for h in st.session_state.prediction_history
    ])
    st.dataframe(history_df, use_container_width=True)

# Footer
st.markdown("""
    <div style='text-align: center; color: gray; padding: 20px;'>
        <p>Developed with ❤️ for Final Yr Project</p>
        <p>Version 2.0 | Last Updated: 2024</p>
    </div>
""", unsafe_allow_html=True)