import gradio as gr
from ultralytics import YOLO
import google.generativeai as genai
import os
from PIL import Image
import warnings

import torch
from ultralytics.nn.tasks import ClassificationModel 
from ultralytics.nn.modules import conv, block, head
from torch.nn import Sequential, Conv2d, Module, BatchNorm2d, SiLU, ModuleList, AdaptiveAvgPool2d, Dropout, Linear
from torchvision.transforms import transforms
from torchvision.transforms.functional import InterpolationMode

torch.serialization.add_safe_globals([
    ClassificationModel,
    Sequential,
    conv.Conv,
    Conv2d,
    block.C2f,
    block.Bottleneck,
    block.SPPF,
    Module,
    BatchNorm2d,
    SiLU,
    ModuleList,
    head.Classify,
    AdaptiveAvgPool2d,
    Dropout,
    Linear,
    transforms.Compose,
    transforms.Resize,
    transforms.CenterCrop,
    transforms.ToTensor,
    transforms.Normalize,
    InterpolationMode 
])

warnings.filterwarnings("ignore")

yolo_model = None
gemini_model = None
try:
    yolo_model = YOLO('./weights/best.pt')
except Exception as e:
    print(f"Error memuat model YOLO: {e}")

GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY')
if GEMINI_API_KEY:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
        gemini_model = genai.GenerativeModel('gemini-2.0-flash')
    except Exception as e:
        print(f"Error konfigurasi Gemini: {e}")
else:
    print("Peringatan: Secret 'GEMINI_API_KEY' tidak ditemukan.")

kamus_hama = {
    'aphids': 'Kutu Daun', 'armyworm': 'Ulat Grayak', 'beetle': 'Kumbang',
    'bollworm': 'Ulat Penggerek Buah', 'grasshopper': 'Belalang', 'mites': 'Tungau',
    'mosquito': 'Nyamuk', 'sawfly': 'Lalat Gergaji', 'stem_borer': 'Penggerek Batang',
    'ants': 'Semut', 'bees': 'Lebah', 'catterpillar': 'Ulat Bulu',
    'earthworms': 'Cacing Tanah', 'earwig': 'Gegat', 'moth': 'Ngengat', 
    'slug': 'Siput Telanjang', 'snail': 'Siput', 'wasp': 'Tawon', 'weevil': 'Kumbang Bubuk'
}

def dapatkan_panduan_hama(nama_hama_indonesia):
    if not gemini_model:
        return "⚠️ **Model Gemini tidak terkonfigurasi.**\nHarap atur secret `GEMINI_API_KEY` di pengaturan Hugging Face Space untuk mendapatkan panduan penanganan."
    prompt = f"""
    Sebagai seorang ahli pertanian, berikan panduan penanganan yang jelas dan praktis untuk hama '{nama_hama_indonesia}'.
    Struktur jawaban Anda harus mencakup poin-poin berikut dengan penjelasan dengan jelas:
    - Deskripsi : Penampilan dan siklus hidup hama. (buat dalam bentuk paragraf)
    - Gejala Serangan: Tanda-tanda pada tanaman yang diserang. (buat dalam bentuk poin dan subpoin penting)
    - Metode Pencegahan: Cara-cara mencegah hama datang. (buat dalam bentuk poin dan subpoin penting)
    - Penanganan Non-Kimiawi: Contoh musuh alami atau pestisida nabati. (buat dalam bentuk poin dan subpoin penting)
    - Penanganan Kimiawi : Contoh bahan aktif insektisida yang efektif. (buat dalam bentuk poin dan subpoin penting)
    """
    try:
        response = gemini_model.generate_content(prompt)
        return response.text
    except Exception as e:
        return f"❌ Gagal mendapatkan panduan dari Gemini. Error: {e}"

def identifikasi_hama_dan_panduan(gambar_input):
    if not yolo_model:
        return None, "Model YOLO tidak berhasil dimuat."
    results = yolo_model.predict(gambar_input, verbose=False)
    result = results[0]
    probabilities = result.probs
    nama_inggris = yolo_model.names[probabilities.top1]
    keyakinan = probabilities.top1conf
    nama_indonesia = kamus_hama.get(nama_inggris, nama_inggris)
    hasil_prediksi = f"### 🎯 Hasil Identifikasi\n- **Jenis Hama:** {nama_indonesia} (*{nama_inggris}*)\n- **Tingkat Akurasi:** {keyakinan.item():.2%}"
    panduan_penanganan = dapatkan_panduan_hama(nama_indonesia)
    return gambar_input, hasil_prediksi, panduan_penanganan

with gr.Blocks(theme=gr.themes.Soft(primary_hue="lime", secondary_hue="green")) as app:
    gr.Markdown("# 🌿 Eleaf : Pest Identification AI 🌿")
    gr.Markdown("Unggah gambar hama pada tanaman Anda. AI akan mengidentifikasi jenisnya dan memberikan panduan penanganan lengkap.")
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="pil", label="Unggah Gambar di Sini")
            submit_button = gr.Button("🔍 Identifikasi Hama", variant="primary")
        with gr.Column(scale=2):
            output_image = gr.Image(label="Gambar yang Dianalisis", interactive=False)
            output_prediksi = gr.Markdown(label="Hasil Analisis")
            output_panduan = gr.Markdown(label="Panduan Penanganan")
    submit_button.click(fn=identifikasi_hama_dan_panduan, inputs=input_image, outputs=[output_image, output_prediksi, output_panduan])

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