# Plant Species Classifier - Hugging Face Space
# Developed with PyTorch + Gradio

import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import wikipediaapi
import random
from io import BytesIO
import requests

# Load model and mappings
print("Loading model...")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

checkpoint = torch.load('best_model.pt', map_location=device)
num_classes = checkpoint['num_classes']
idx_to_label = checkpoint['idx_to_label']

# Build model
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, num_classes)
model.load_state_dict(checkpoint['model_state_dict'])
model = model.to(device)
model.eval()

# Transforms
val_test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Wikipedia API
wiki = wikipediaapi.Wikipedia(
    language='en',
    user_agent='PlantClassifier/1.0 (https://huggingface.co/spaces)'
)

def get_species_info(species_name):
    try:
        page = wiki.page(species_name)
        if page.exists():
            summary = page.summary[:500]
            if len(page.summary) > 500:
                summary += "..."
            return {
                'title': page.title,
                'summary': summary,
                'url': page.fullurl,
                'found': True
            }
        else:
            genus = species_name.split()[0]
            page = wiki.page(genus)
            if page.exists():
                summary = page.summary[:500]
                if len(page.summary) > 500:
                    summary += "..."
                return {
                    'title': f"{page.title} (genus)",
                    'summary': summary,
                    'url': page.fullurl,
                    'found': True
                }
            return {
                'title': species_name,
                'summary': f"No Wikipedia information found for {species_name}.",
                'url': '',
                'found': False
            }
    except Exception as e:
        return {
            'title': species_name,
            'summary': f"Error fetching information: {str(e)}",
            'url': '',
            'found': False
        }

def predict_and_display(input_image, image_url):
    img = None

    if image_url and image_url.strip():
        try:
            response = requests.get(image_url.strip(), timeout=10, 
                                   headers={"User-Agent": "Mozilla/5.0"})
            response.raise_for_status()
            img = Image.open(BytesIO(response.content)).convert('RGB')
        except Exception as e:
            return f"❌ Error downloading image from URL: {str(e)}", None, []
    elif input_image is not None:
        img = input_image
    else:
        return "⚠️ Please upload an image or provide a URL", None, []

    model.eval()
    img_tensor = val_test_transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(img_tensor)
        probs = torch.nn.functional.softmax(outputs, dim=1)

    topk_probs, topk_indices = probs.topk(5, dim=1)

    predictions = []
    for prob, idx in zip(topk_probs[0], topk_indices[0]):
        label = idx_to_label[str(idx.item())]
        predictions.append((label, prob.item()))

    html = "<h2>🎯 Top 5 Predictions</h2>"
    html += "<div style='margin-top: 20px;'>"

    for i, (label, prob) in enumerate(predictions, 1):
        if prob > 0.5:
            color = "#2ecc71"
        elif prob > 0.2:
            color = "#f39c12"
        else:
            color = "#95a5a6"

        bar_width = int(prob * 100)

        html += f"""
        <div style='margin-bottom: 15px; padding: 15px; background: #f8f9fa; border-radius: 8px; border-left: 4px solid {color};'>
            <div style='display: flex; justify-content: space-between; align-items: center; margin-bottom: 10px;'>
                <span style='font-size: 18px; font-weight: bold; color: #2c3e50;'>#{i}</span>
                <span style='font-size: 18px; font-weight: 700; color: {color};'>{prob*100:.2f}%</span>
            </div>
            <div style='font-size: 16px; font-weight: 600; color: #2c3e50; margin-bottom: 10px; line-height: 1.4;'>{label}</div>
            <div style='width: 100%; background: #e0e0e0; border-radius: 10px; height: 10px;'>
                <div style='width: {bar_width}%; background: {color}; height: 10px; border-radius: 10px;'></div>
            </div>
        </div>
        """

    html += "</div>"
    return html, img, predictions

def show_species_details(predictions, selected_index):
    if not predictions or selected_index < 1 or selected_index > len(predictions):
        return "⚠️ Please make a prediction first"

    selected_species = predictions[int(selected_index) - 1][0]
    confidence = predictions[int(selected_index) - 1][1]

    wiki_info = get_species_info(selected_species)

    html = f"""
    <div style='padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                border-radius: 12px; color: white; margin-bottom: 20px;'>
        <h1 style='margin: 0 0 10px 0; font-size: 28px;'>{selected_species}</h1>
        <p style='margin: 0; font-size: 18px; opacity: 0.9;'>Confidence: {confidence*100:.2f}%</p>
    </div>

    <div style='background: #f8f9fa; padding: 20px; border-radius: 12px; margin-bottom: 20px;'>
        <h2 style='color: #2c3e50; margin-top: 0;'>📚 {wiki_info['title']}</h2>
        <p style='line-height: 1.6; color: #34495e; font-size: 15px;'>{wiki_info['summary']}</p>
        {f"<a href='{wiki_info['url']}' target='_blank' style='color: #3498db; text-decoration: none; font-weight: 600;'>📖 Read more on Wikipedia →</a>" if wiki_info['url'] else ""}
    </div>
    """

    return html

# Create Gradio Interface
with gr.Blocks(title="🌿 Plant Species Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🌿 Plant Species Classifier
    ### AI-powered plant species identification from GBIF TOP100 dataset

    **How to use:**
    1. Upload an image or paste an image URL
    2. Click "🔍 Predict" to see top 5 species predictions
    3. Click "📊 Show Details" to explore detailed information

    **Model:** ResNet18 trained on 100 plant species from GBIF
    """)

    predictions_state = gr.State([])

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Input")

            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                height=300
            )

            gr.Markdown("**OR**")

            image_url = gr.Textbox(
                label="Image URL",
                placeholder="https://example.com/plant.jpg",
                lines=1
            )

            with gr.Row():
                predict_btn = gr.Button("🔍 Predict", variant="primary", size="lg")
                clear_btn = gr.ClearButton([input_image, image_url], value="🗑️ Clear")

        with gr.Column(scale=1):
            gr.Markdown("### 🎯 Predictions")
            predictions_display = gr.HTML(
                "<p style='text-align: center; color: #95a5a6; padding: 50px;'>Predictions will appear here...</p>"
            )

    gr.Markdown("---")
    gr.Markdown("### 🔬 Species Details")
    gr.Markdown("Select a prediction number (1-5) to see detailed information:")

    with gr.Row():
        species_selector = gr.Slider(
            minimum=1,
            maximum=5,
            step=1,
            value=1,
            label="Select Prediction #",
            info="Choose which prediction to explore"
        )
        details_btn = gr.Button("📊 Show Details", variant="secondary", size="lg")

    species_info = gr.HTML()

    def predict_wrapper(img, url):
        html, display_img, preds = predict_and_display(img, url)
        return html, preds

    predict_btn.click(
        fn=predict_wrapper,
        inputs=[input_image, image_url],
        outputs=[predictions_display, predictions_state]
    )

    details_btn.click(
        fn=show_species_details,
        inputs=[predictions_state, species_selector],
        outputs=[species_info]
    )

    gr.Markdown("""
    ---
    <div style='text-align: center; color: #95a5a6;'>
        <p>🌱 Built with PyTorch + Gradio | Data from GBIF</p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch()