import gradio as gr
from PIL import Image
import torch
import numpy as np
import faiss

from transformers import (
    GitProcessor,
    GitForCausalLM,
    AutoTokenizer,
    AutoModelForCausalLM,
    CLIPProcessor,
    CLIPModel
)
from sentence_transformers import SentenceTransformer
from datasets import load_dataset

device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")

tokenizer_llama = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model_llama = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
).eval()

text_encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device).eval()
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Загрузка только первых 10000 изображений через streaming
MAX_IMAGES = 10_000
dataset_stream = load_dataset("huggan/wikiart", split="train", streaming=True)
first_10000 = [x for i, x in enumerate(dataset_stream) if i < MAX_IMAGES]

image_index = faiss.read_index("image_index_llama.faiss")
text_index = faiss.read_index("text_index_llama.faiss")

def clean_caption(text):
    return text.replace("[ unused0 ]", "").strip()

def generate_captions(image: Image.Image):
    inputs = git_processor(images=image, return_tensors="pt")["pixel_values"].to(device)

    captions = []
    with torch.no_grad():
        deterministic_ids = git_model.generate(
            pixel_values=inputs,
            max_new_tokens=30,
            do_sample=False
        )
        captions.append(clean_caption(git_processor.tokenizer.decode(deterministic_ids[0], skip_special_tokens=True)))

        sampled_ids = git_model.generate(
            pixel_values=inputs,
            max_new_tokens=30,
            do_sample=True,
            top_k=100,
            temperature=0.8,
            num_return_sequences=2
        )
        sampled = git_processor.tokenizer.batch_decode(sampled_ids, skip_special_tokens=True)
        captions.extend([clean_caption(c) for c in sampled])

    return captions

def refine_caption(base, desc1, desc2):
    prompt = f"""
Given the base caption that is true and factual:
\"{base}\"

And two descriptive captions:
1) {desc1}
2) {desc2}

Write a short, coherent description that is faithful to the base caption but incorporates descriptive elements from captions 1 and 2 without contradicting the original meaning.
"""
    inputs = tokenizer_llama(prompt, return_tensors="pt").to(model_llama.device)
    with torch.no_grad():
        output = model_llama.generate(**inputs, max_new_tokens=100, do_sample=False)
        text = tokenizer_llama.decode(output[0], skip_special_tokens=True)
        answer = text[len(prompt):].strip()
        for prefix in ["Example:", "example:"]:
            if answer.startswith(prefix):
                answer = answer[len(prefix):].strip()
        return answer

def get_text_embedding(text):
    emb = text_encoder.encode([text], normalize_embeddings=False).astype("float32")
    faiss.normalize_L2(emb)
    return emb

def get_image_embedding(image):
    inputs = clip_processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        image_features = clip_model.get_image_features(**inputs)
    emb = image_features.cpu().numpy().astype("float32")
    faiss.normalize_L2(emb)
    return emb

def get_results_with_images(embedding, index, top_k=2):
    D, I = index.search(embedding, top_k)
    results = []
    for idx in I[0]:
        if idx >= MAX_IMAGES:
            continue
        try:
            item = first_10000[idx]
            img = item["image"]
            caption = item["caption"]
            caption_text = f"ID: {idx}\n{caption}"
            results.append((img, caption_text))
        except IndexError:
            continue
    return results

def search_similar_images(image: Image.Image):
    captions = generate_captions(image)
    refined = refine_caption(captions[0], captions[1], captions[2])

    text_emb = get_text_embedding(refined)
    image_emb = get_image_embedding(image)

    text_results = get_results_with_images(text_emb, text_index)
    image_results = get_results_with_images(image_emb, image_index)

    return refined, text_results, image_results

demo = gr.Interface(
    fn=search_similar_images,
    inputs=gr.Image(label="Загрузите изображение", type="pil"),
    outputs=[
        gr.Textbox(label="📜 Сгенерированное описание"),
        gr.Gallery(label="🔍 Похожие по описанию (caption)", height="auto", columns=2),
        gr.Gallery(label="🎨 Похожие по изображению (CLIP)", height="auto", columns=2)
    ],
    title="🎨 Semantic WikiArt Search",
    description="Загрузите изображение. Модель сгенерирует описание, получит эмбеддинги и найдёт похожие картины по описанию и изображению."
)

demo.launch()