from PIL import Image
from transformers import pipeline
import gradio as gr

# 🌐 Load pre-trained image classification model
classifier = pipeline("image-classification", model="microsoft/resnet-50")

# 🔍 Define bilingual label mapping
label_map = {
    "agaric": ("Edible", "กินได้"),
    "bolete": ("Edible", "กินได้"),
    "gyromitra": ("Poisonous", "พิษ"),
    "amanita": ("Poisonous", "พิษ"),
    "earthstar": ("Edible", "กินได้"),
    "hen-of-the-woods": ("Edible", "กินได้"),
    "mushroom": ("Unknown", "ไม่ทราบ"),
    "coral fungus": ("Poisonous", "พิษ"),
    "Amanita muscaria":("Poisonous","พิษ")
    # Add more if needed
}

# 🧠 Classification function
def classify_mushroom(image: Image.Image):
    print("✅ classify_mushroom: NEW VERSION")

    try:
        image = image.convert("RGB")
        result = classifier(image)
        print("🔍 Raw result from model:", result)

        result = result[0]
        label = result['label'].lower()
        score = round(result['score'] * 100, 2)

        prediction_en, prediction_th = label_map.get(label, ("Unknown", "ไม่ทราบ"))
        print("✅ RETURNING:", prediction_en, prediction_th, f"{score:.2f}%", label)

        return prediction_en, prediction_th, f"{score:.2f}%", label

    except Exception as e:
        print(f"❌ Error: {e}")
        return "Error", "ผิดพลาด", "N/A", "N/A"

# 🎨 Gradio UI
if __name__ == "__main__":
    with gr.Blocks() as demo:
        gr.Markdown("## 🍄 Mushroom Safety Classifier")
        gr.Markdown("Upload a mushroom photo to check if it’s edible or poisonous.\nอัปโหลดรูปเห็ดเพื่อทำนายว่าเห็ดกินได้หรือมีพิษ")

        with gr.Row():
            image_input = gr.Image(type="pil", label="📷 Upload Mushroom Image")
            with gr.Column():
                label_en = gr.Textbox(label="🧠 Prediction (English)")
                label_th = gr.Textbox(label="🗣️ คำทำนาย (ภาษาไทย)")
                confidence = gr.Textbox(label="📶 Confidence Score")
                label_raw = gr.Textbox(label="🏷️ Predicted Mushroom Name")

        classify_btn = gr.Button("🔍 Classify")

        classify_btn.click(
            fn=classify_mushroom,
            inputs=image_input,
            outputs=[label_en, label_th, confidence, label_raw]
        )

    demo.launch()