import gradio as gr
import torch
import numpy as np
from collections import defaultdict
from transformers import AutoImageProcessor, AutoModel
from sklearn.metrics.pairwise import cosine_similarity

# Load model pretrained (visual embedding)
processor = AutoImageProcessor.from_pretrained("facebook/dino-vits16")
model = AutoModel.from_pretrained("facebook/dino-vits16")
model.eval()

# Database: label -> list of features
feature_db = defaultdict(list)

def extract_feature(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    return outputs.last_hidden_state.mean(dim=1).numpy()

def tambah_sampling(image, label):
    if image is None or label.strip() == "":
        return "❌ Gambar dan label wajib diisi"

    feat = extract_feature(image)
    feature_db[label].append(feat)

    return f"✅ Sampling ditambahkan untuk label '{label}' (total: {len(feature_db[label])})"

def tebak_gambar(image):
    if image is None:
        return "❌ Upload gambar terlebih dahulu"

    if len(feature_db) == 0:
        return "❌ Belum ada data sampling"

    test_feat = extract_feature(image)

    label_scores = {}
    for label, feats in feature_db.items():
        sims = cosine_similarity(test_feat, np.vstack(feats))
        label_scores[label] = sims.mean()  # rata-rata kemiripan

    # Ambil label dengan skor tertinggi
    pred_label = max(label_scores, key=label_scores.get)
    skor = label_scores[pred_label]

    hasil = "📌 HASIL PREDIKSI (BERDASARKAN KEMIRIPAN)\n\n"
    for lbl, sc in sorted(label_scores.items(), key=lambda x: x[1], reverse=True):
        hasil += f"- {lbl}: {sc:.2f}\n"

    hasil += f"\n✅ Gambar PALING MIRIP dengan: **{pred_label}**"

    return hasil

with gr.Blocks() as app:
    gr.Markdown("## 🧠 Tebak Gambar Supervised (Similarity-Based)")
    gr.Markdown(
        "Upload gambar sampling dengan label, lalu upload gambar baru. "
        "AI menentukan label berdasarkan kemiripan visual, bukan gambar yang sama persis."
    )

    with gr.Tab("1️⃣ Tambah Data Sampling"):
        img1 = gr.Image(type="pil", label="Gambar Sampling")
        lbl = gr.Textbox(label="Label")
        out1 = gr.Textbox(label="Status")
        btn1 = gr.Button("Tambah")

        btn1.click(tambah_sampling, [img1, lbl], out1)

    with gr.Tab("2️⃣ Tebak Gambar"):
        img2 = gr.Image(type="pil", label="Gambar Uji")
        out2 = gr.Textbox(label="Hasil")
        btn2 = gr.Button("Tebak")

        btn2.click(tebak_gambar, img2, out2)

app.launch()