import gradio as gr
import torch
from transformers import AutoProcessor, AutoModel

# Load MetaCLIP
model_name = "facebook/metaclip-2-worldwide-s16"
processor = AutoProcessor.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)


def classify(image, labels_text):
    if image is None or not labels_text.strip():
        return (
            "<div style='font-size:22px; color:#b00000;'>กรุณาอัปโหลดภาพและกรอก labels ก่อน</div>",
            ""
        )

    labels = [l.strip() for l in labels_text.split("\n") if l.strip()]
    inputs = processor(text=labels, images=image, return_tensors="pt", padding=True)

    with torch.no_grad():
        out = model(**inputs)
        probs = out.logits_per_image.softmax(dim=1)[0]

    best_idx = torch.argmax(probs)
    best_label = labels[best_idx]
    best_prob = float(probs[best_idx])

    # ===== HTML การ์ดไฮไลท์ตรงกลาง =====
    highlight_html = f"""
    <div style="
        text-align:center;
        padding:30px;
        border-radius:20px;
        background: linear-gradient(135deg, #ffe7e7, #fff4e6);
        border:1px solid #f2d0d0;
        box-shadow:0 4px 20px rgba(0,0,0,0.08);
        margin-bottom: 20px;
    ">
        <div style='font-size:40px; font-weight:900; color:#b80000; margin-bottom:10px;'>
            ผลลัพธ์ที่ใกล้ที่สุด
        </div>

        <div style='font-size:48px; font-weight:900; color:#d00000;'>
            🎯 {best_label} 🎯
        </div>

        <div style='font-size:26px; margin-top:14px; color:#444;'>
            ความมั่นใจ: <b>{best_prob:.4f}</b>
        </div>
    </div>
    """

    # ===== Markdown รายละเอียดทั้งหมด =====
    details_md = "### รายละเอียดทั้งหมด\n"
    for lbl, p in zip(labels, probs):
        details_md += f"- **{lbl}**: {float(p):.4f}\n"

    return highlight_html, details_md


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎨 ClipSense Mini  
        ทดลองว่า MetaCLIP 2 คิดว่าภาพของคุณ “ใกล้คำไหนที่สุด”
        """
    )

    with gr.Row(equal_height=True):
        # ===== LEFT : Image + Labels =====
        with gr.Column(scale=1):
            image = gr.Image(type="pil", height=350, label="อัปโหลดรูปภาพ")

            labels_box = gr.Textbox(
                label="Labels (ใส่ทีละบรรทัด / ใช้ภาษาไทยได้)",
                lines=7,
                placeholder="ผู้หญิง\nผู้ชาย\nรูปโป๊\nนางตานี\nแมว\nหมา"
            )

            btn = gr.Button("วิเคราะห์ภาพด้วย ClipSense")

        # ===== RIGHT : Result =====
        with gr.Column(scale=1):
            highlight = gr.HTML("<div>ผลลัพธ์จะปรากฏที่นี่</div>")
            details = gr.Markdown("")

    btn.click(classify, [image, labels_box], [highlight, details])

demo.launch()