import gradio as gr
from transformers import AutoImageProcessor, AutoModel
import torch
from PIL import Image
import os
import json
import uuid
import numpy as np

# Загрузка модели DINOv2
processor = AutoImageProcessor.from_pretrained("facebook/dino-vits16")
model = AutoModel.from_pretrained("facebook/dino-vits16")

# Функция для извлечения эмбеддинга
def extract_features(img):
    inputs = processor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    return outputs.last_hidden_state[:, 0].squeeze().numpy()

# Косинусное расстояние между двумя эмбеддингами
def cosine_similarity(vec1, vec2):
    a = np.array(vec1)
    b = np.array(vec2)
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# Путь к базе
db_path = "embeddings2.json"

# Сохранение фото в базу
def save_embedding(image, building_name):
    if not building_name:
        return "Ошибка: нужно указать название здания."
    
    embedding = extract_features(image).tolist()
    entry = {
        "id": str(uuid.uuid4()),
        "building_name": building_name,
        "embedding": embedding
    }

    if os.path.exists(db_path):
        with open(db_path, "r") as f:
            data = json.load(f)
    else:
        data = []

    data.append(entry)

    with open(db_path, "w") as f:
        json.dump(data, f, indent=2)

    return f"Фото сохранено для здания: {building_name}"

# Поиск похожего здания
def identify_building(image):
    if not os.path.exists(db_path):
        return "База данных пуста. Добавьте здания."

    with open(db_path, "r") as f:
        data = json.load(f)

    if not data:
        return "База данных пуста."

    embedding = extract_features(image).tolist()

    similarities = []
    for item in data:
        score = cosine_similarity(item["embedding"], embedding)
        similarities.append({
            "building_name": item["building_name"],
            "score": score
        })

    similarities.sort(key=lambda x: x["score"], reverse=True)
    best = similarities[0]

    return f"Похоже на: {best['building_name']}\nСовпадение: {best['score']:.4f}"

# Интерфейс
with gr.Blocks() as demo:
    gr.Markdown("## 🏛 Распознавание зданий и пополнение базы")

    with gr.Tab("🔍 Найти здание"):
        with gr.Row():
            img_input = gr.Image(type="pil")
            recognize_button = gr.Button("Распознать здание")
        result_text = gr.Textbox()
        recognize_button.click(fn=identify_building, inputs=img_input, outputs=result_text)

    with gr.Tab("➕ Добавить новое здание"):
        with gr.Row():
            img_save = gr.Image(type="pil")
            building_name = gr.Textbox(label="Название здания")
            save_button = gr.Button("Сохранить в базу")
        save_result = gr.Textbox()
        save_button.click(fn=save_embedding, inputs=[img_save, building_name], outputs=save_result)

demo.launch()