import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import os
import google.generativeai as genai
import pandas as pd # 1. Impor library pandas

# ==============================================================================
# 1. KONFIGURASI DAN MEMUAT SEMUA MODEL
# ==============================================================================
print("Memuat semua model dan konfigurasi...")
device = torch.device('cpu') 

# --- A. Muat data referensi link dari CSV ---
try:
    df_referensi = pd.read_csv('link_referensi.csv')
    print("File 'link_referensi.csv' berhasil dimuat.")
    referensi_loaded = True
except FileNotFoundError:
    print("PERINGATAN: File 'link_referensi.csv' tidak ditemukan. Fitur referensi resep lokal akan dinonaktifkan.")
    referensi_loaded = False

# --- B. Konfigurasi Gemini API ---
try:
    genai.configure(api_key=os.environ['GEMINI_API_KEY'])
    gemini_model = genai.GenerativeModel("gemini-1.5-flash")
    gemini_enabled = True
    print("Koneksi ke Gemini API (gemini-1.5-flash) berhasil.")
except KeyError:
    gemini_enabled = False
    print("PERINGATAN: GEMINI_API_KEY tidak ditemukan. Fitur deskripsi dinonaktifkan.")

# --- C. Muat Model Klasifikasi Biner ---
model_biner = models.mobilenet_v2(pretrained=False)
model_biner.classifier[1] = nn.Linear(model_biner.last_channel, 1)
model_biner.load_state_dict(torch.load('plant_detector_binary.pth', map_location=device))
model_biner.eval()
print("Model deteksi biner (plant_detector_binary.pth) berhasil dimuat.")

# --- D. Muat Model Klasifikasi Herbal ---
with open('class_names.txt', 'r') as f:
    class_names = [line.strip() for line in f.readlines()]
num_classes = len(class_names)

model_herbal = models.vgg11_bn(pretrained=False)
num_ftrs = model_herbal.classifier[6].in_features
model_herbal.classifier[6] = nn.Linear(num_ftrs, num_classes)
model_herbal.load_state_dict(torch.load('herbal_vgg11_optimized.pth', map_location=device))
model_herbal.eval()
print(f"Model deteksi herbal (herbal_vgg11_optimized.pth) dengan {num_classes} kelas berhasil dimuat.")

# --- E. Definisikan Transformasi Gambar ---
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# ==============================================================================
# 2. DEFINISIKAN FUNGSI-FUNGSI BANTUAN
# ==============================================================================
def get_plant_info_from_gemini(plant_name):
    """Meminta informasi tanaman dari Gemini (tanpa resep)."""
    if not gemini_enabled:
        return "Fitur deskripsi tidak aktif karena GEMINI_API_KEY tidak dikonfigurasi."
    try:
        # Prompt diubah untuk tidak secara eksplisit meminta link resep lagi
        prompt = f"""
        Berikan informasi mengenai tanaman herbal '{plant_name}' dalam Bahasa Indonesia.
        Tolong susun jawabannya dalam format yang jelas:

        **Deskripsi:**
        [Tulis deskripsi singkat tentang apa itu tanaman {plant_name}, asalnya, dan ciri-ciri utamanya di sini]

        **Manfaat:**
        [Sajikan dalam bentuk daftar poin (bullet points)]
        - [Manfaat 1]
        - [Manfaat 2]
        - [dan seterusnya...]

        **Cara Pengolahan Umum:**
        [Jelaskan cara pengolahan secara umum untuk penggunaan luar dan konsumsi.]
        """
        response = gemini_model.generate_content(prompt)
        return response.text
    except Exception as e:
        return f"Gagal menghubungi Gemini: {e}"

# 2. Fungsi baru untuk mencari link di file CSV
def get_links_from_csv(plant_name):
    """Mencari link referensi dari file CSV yang sudah dimuat."""
    if not referensi_loaded:
        return "File referensi lokal tidak dimuat."
    
    # Cari baris yang cocok dengan plant_name
    # Menggunakan .str.contains untuk pencarian yang lebih fleksibel, case=False agar tidak sensitif huruf besar/kecil
    hasil = df_referensi[df_referensi['Nama Tanaman'].str.contains(plant_name, case=False, na=False)]
    
    if hasil.empty:
        return f"Tidak ditemukan referensi resep lokal untuk '{plant_name}'."
    
    # Format output menjadi markdown list
    links_markdown = "### Referensi Resep (dari file lokal):\n"
    for index, row in hasil.iterrows():
        links_markdown += f"- [{row['URL']}]({row['URL']})\n"
        
    return links_markdown

# ==============================================================================
# 3. DEFINISIKAN FUNGSI PREDIKSI UTAMA
# ==============================================================================
def predict_cascaded(image):
    """Fungsi utama yang menjalankan semua tahap prediksi."""
    if image is None:
        return None, "Tidak ada gambar.", "Silakan unggah gambar.", "Tidak ada data."
            
    image_pil = Image.fromarray(image.astype('uint8'), 'RGB')
    image_transformed = transform(image_pil).unsqueeze(0).to(device)

    # --- Tahap 1: Deteksi Biner ---
    with torch.no_grad():
        output_biner = model_biner(image_transformed)
        prob_is_plant = torch.sigmoid(output_biner).item()

    DETECTION_THRESHOLD = 0.5
    if prob_is_plant < DETECTION_THRESHOLD:
        return {"Bukan Tanaman": 1 - prob_is_plant}, "Gambar tidak terdeteksi sebagai tanaman.", "Silakan unggah gambar tanaman.", "Tidak ada data."

    # --- Tahap 2: Identifikasi Jenis Herbal ---
    with torch.no_grad():
        outputs_herbal = model_herbal(image_transformed)
        probabilities = torch.nn.functional.softmax(outputs_herbal[0], dim=0)
        confidences = {class_names[i]: float(probabilities[i]) for i in range(num_classes)}
        
    # Mendapatkan nama prediksi teratas tanpa nama latinnya
    top_prediction_full_name = max(confidences, key=confidences.get)
    # Membersihkan nama tanaman, mengambil bagian sebelum tanda kurung buka
    top_prediction_name = top_prediction_full_name.split('(')[0].strip()

    # --- Tahap 3: Dapatkan Informasi dari Gemini dan File CSV ---
    plant_info_gemini = get_plant_info_from_gemini(top_prediction_name)
    plant_links_csv = get_links_from_csv(top_prediction_name) # Panggil fungsi baru

    # 3. Kembalikan 4 nilai output
    return confidences, top_prediction_full_name, plant_info_gemini, plant_links_csv

# ==============================================================================
# 4. BUAT DAN LUNCURKAN ANTARMUKA GRADIO
# ==============================================================================
interface = gr.Interface(
    fn=predict_cascaded,
    inputs=gr.Image(label="Unggah Gambar Tanaman"),
    # 4. Tambahkan output baru untuk menampilkan link dari CSV
    outputs=[
        gr.Label(num_top_classes=5, label="Hasil Prediksi (Probabilitas)"),
        gr.Textbox(label="Objek Terdeteksi"),
        gr.Markdown(label="Informasi Tanaman (dari Gemini)"),
        gr.Markdown(label="Referensi Resep (dari File Lokal)") # Output baru
    ],
    title="🌿 Daunesia: Deteksi Cerdas Tanaman Herbal",
    description="Sistem deteksi dua-tahap. Unggah gambar untuk mengidentifikasi jenis tanaman, mendapatkan deskripsi dari AI, dan melihat resep dari file lokal.",
    examples=[
        # Anda mungkin perlu mengganti contoh ini dengan file gambar yang Anda miliki
        # ["nangkacempedak.jpg"], 
        # ["aloevera.jpeg"]
    ],
    allow_flagging="never",
    theme=gr.themes.Soft()
)

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