import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import tempfile
import os
import functools

# ------------------- Load Model -------------------
@functools.lru_cache(maxsize=1)
def load_model():
    model_path = "best.pt"  # file must exist in repo root or /app directory
    if not os.path.exists(model_path):
        raise FileNotFoundError("Model file 'best.pt' not found in the Space directory.")
    return YOLO(model_path)

# Load once globally
model = load_model()

# ------------------- Wound Descriptions -------------------DO NOT chanage class name (It matched with data.yalm data set)
wound_descriptions = {
    "wound_hesitation": "บาดแผลลังเล (Hesitation): มักพบในผู้พยายามทำร้ายตนเอง มีลักษณะเป็นแผลตื้นหลายแผลขนานหรือเกือบขนานกัน อยู่ใกล้แผลหลัก หรือบริเวณที่ทำตนเองได้ง่าย เช่น ข้อมือด้านใน",
    "wound_laceration": "บาดแผลฉีกขาดขอบไม่เรียบ (Laceration): เกิดจากวัตถุแข็งไม่มีคมกระแทก มีลักษณะขอบไม่เรียบ มีถลอกและฟกช้ำที่ขอบแผล และมี tissue bridging/undermining",
    "wound_open_fracture": "บาดแผลกระดูกหักแบบเปิด (Open fracture): เกิดจากกระดูกหักทิ่มออกมานอกผิวหนัง ในกรณีถูกรถชน สามารถประมวลเหตุการณ์ชนได้ (reconstruction)",
    "wound_burn": "บาดแผลไหม้ (Fire burn): บาดแผลที่เกิดจากการไหม้ ให้ดูความลึกและการกระจายของบาดแผล (pattern) ว่าสอดคล้องกับเหตุการณ์หรือไม่",
    "wound_hanging": "บาดแผลกดรัดบริเวณลำคอ (Hanging): เกิดจากน้ำหนักร่างกายทำให้เกิดบาดแผลกดรัดที่คอ โดยทั่วไปหากพบลักษณะการกดรัดเฉียงขึ้น จะเป็นลักษณะของ hanging ซึ่งต้องดูประกอบกับหลักฐานอื่น",
    "wound_strangulation": "บาดแผลกดรัดบริเวณลำคอ (Strangulation): เกิดจากแรงกระทำที่ไม่ใช่น้ำหนักร่างกาย ทำให้เกิดกดรัดที่คอ มีลักษณะการกดรัดแนวขวาง ซึ่งหากพบลักษณะบาดแผลกดรัดสองรูปแบบให้นึกถึงการฆาตกรรมอำพราง",
    "gsw_entrance": "บาดแผลทางเข้ากระสุนปืน (gunshot wound entrance): ลักษณะบาดแผลกระสุนปืนจะมีลักษณะเฉพาะ คือ punch-out lesion ซึ่งทางเข้าอาจพบองค์ประกอบการยิง เช่น เขม่าดินปืนดังภาพ (soot/gun powder tatooing)",
    "gsw_exit": "บาดแผลทางออกกระสุนปืน (gunshot wound exit): ลักษณะบาดแผลกระสุนปืนจะมีลักษณะเฉพาะ คือ punch-out lesion โดยทางออกจะไม่พบองค์ประกบการยิอง และโดยทั่วไปจะขนาดใหญ่กว่าทางเข้า อาจะมีรูปร่างแฉกคล้ายบาดแผลฉีกขาดขอบไม่เรียบ"
}

# ------------------- Detection Function -------------------
def detect(image, conf):
    if image is None:
        return None, "No image uploaded.", None

    img = np.array(image)
    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # Convert RGB → BGR for YOLO

    results = model(img_bgr, conf=conf)
    annotated_bgr = results[0].plot()  # returns BGR
    annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB)

    detected_classes = set()
    for r in results[0].boxes.cls.cpu().numpy():
        cls_name = results[0].names[int(r)]
        detected_classes.add(cls_name)

    # Add wound descriptions
    if detected_classes:
        desc_texts = [f"**{cls}**: {wound_descriptions.get(cls, '(No description available)')}" for cls in detected_classes]
        desc_str = "\n\n".join(desc_texts)
    else:
        desc_str = "No wounds detected."

    # Save annotated image for download
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    Image.fromarray(annotated_rgb).save(temp_file.name)

    return annotated_rgb, desc_str, temp_file.name

# ------------------- Footer -------------------
footer = """
---
<div style='text-align:center; font-size:14px; color:gray; line-height:1.6;'>
    <b>Forensic Education Version 2.0.0</b> | © 2025 BH <br>
    <a href="https://docs.google.com/document/d/18KlYv7Xbp3Y4Snatfez_jff0OW7DWKPoYP3HA3fx2cQ/edit?usp=sharing" target="_blank">📄 User Manual</a> |
    <a href="https://forms.gle/WgGnkcUQPafyhmng8" target="_blank">👍 Feedback Please</a> |
    <a href="https://drive.google.com/file/d/1Naa1ev1oMZbzFnVl7HW6yJazV371sOUo/view" target="_blank">📚 more about forensic wounds</a>
</div>
"""

# Create main Gradio Interface
demo = gr.Interface(
    fn=detect,
    inputs=[
        gr.Image(type="pil", label="📸 Upload wound image"),
        gr.Slider(0, 1, 0.25, step=0.05, label="Confidence threshold"),
    ],
    outputs=[
        gr.Image(label="Detection result"),
        gr.Markdown(label="Detected wound details"),
        gr.File(label="Download annotated image"),
    ],
    title="🤕 Forensic Wound Detection (YOLOv8n)",
    description="Upload an image to detect and describe forensic wound types using YOLOv8n. Model trained for educational forensic simulation.",
    flagging_mode="never"
)

# Add footer below the app
with gr.Blocks() as app:
    demo.render()
    gr.Markdown(footer)

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