import torch
import torch.nn.functional as F
import gradio as gr
from PIL import Image
import pickle
import json
from huggingface_hub import hf_hub_download

from model import image_captioning_model
from tokenizer import CaptionTokenizer  # Важно!

# Загрузка файлов
repo_id = "Pumponchik/ImageCaption"
repo_type = "space"
model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin", repo_type=repo_type)
config_path = hf_hub_download(repo_id=repo_id, filename="config.json", repo_type=repo_type)
tokenizer_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.pkl", repo_type=repo_type)
transform_path = hf_hub_download(repo_id=repo_id, filename="image_transforms.pkl", repo_type=repo_type)


# Загрузка конфигурации
with open(config_path, "r") as f:
    full_config = json.load(f)

# Фильтрация параметров для модели
model_params = {
    "vocab_size": full_config["vocab_size"],
    "embedding_dim": full_config["embedding_dim"],
    "hidden_size": full_config["hidden_size"],
    "num_layers": full_config["num_layers"]
}

# Инициализация модели
model = image_captioning_model(**model_params)

# Загрузка весов модели
model.load_state_dict(torch.load(model_path, map_location="cpu"), strict=False)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

tokenizer = CaptionTokenizer.load(tokenizer_path)


# Загрузка трансформаций
with open(transform_path, "rb") as f:
    transform = pickle.load(f)

def wrap_text(text, max_width):
    words = text.split(" ")
    result = [[words[0]]]
    for word in words[1:]:
        if len(" ".join(result[-1])) + len(word) + 1 > max_width:
            result[-1] = " ".join(result[-1])
            result.append([])
        result[-1].append(word)
    result[-1] = " ".join(result[-1])
    return "\n".join(result)

def generate(model, image, tokenizer, max_seq_len=20, top_p=None, top_k=1):
    assert top_p is None or top_k is None, "Не используйте одновременно top_p и top_k"

    model.eval()
    device = next(model.parameters()).device
    image = transform(image).to(device).unsqueeze(0)
    batch_size = image.size(0)

    with torch.no_grad():
        img_emb = model.img_fe(image)
        h_0 = model.img_to_hidden(img_emb).unsqueeze(0).repeat(model.text_fe.num_layers, 1, 1)
        c_0 = h_0.clone()

        input_token = torch.full((batch_size,), tokenizer.tok_to_ind['<BOS>'], dtype=torch.long, device=device)
        result_tokens = [[] for _ in range(batch_size)]
        mask = torch.ones(batch_size, dtype=torch.bool, device=device)

        for _ in range(max_seq_len):
            input_emb = model.text_fe.embedding(input_token).unsqueeze(1)
            output, (h_0, c_0) = model.text_fe.rnn(input_emb, (h_0, c_0))
            logits = model.fc_out(output.squeeze(1))
            probs = F.softmax(logits, dim=-1)

            next_tokens = torch.zeros(batch_size, dtype=torch.long, device=device)
            for i in range(batch_size):
                if not mask[i]:
                    continue
                prob = probs[i]
                if top_k is not None:
                    sorted_probs, sorted_indices = torch.sort(prob, descending=True)
                    sorted_probs = sorted_probs[:top_k]
                    sorted_indices = sorted_indices[:top_k]
                    sorted_probs = sorted_probs / sorted_probs.sum()
                    next_token_idx = torch.multinomial(sorted_probs, 1)
                    next_tokens[i] = sorted_indices[next_token_idx]
                elif top_p is not None:
                    sorted_probs, sorted_indices = torch.sort(prob, descending=True)
                    cumulative_probs = torch.cumsum(sorted_probs, dim=0)
                    mask_p = cumulative_probs <= top_p
                    sorted_probs = sorted_probs[mask_p]
                    sorted_indices = sorted_indices[mask_p]
                    sorted_probs = sorted_probs / sorted_probs.sum()
                    next_token_idx = torch.multinomial(sorted_probs, 1)
                    next_tokens[i] = sorted_indices[next_token_idx]
                else:
                    next_tokens[i] = torch.multinomial(prob, 1)

                token_id = next_tokens[i].item()
                result_tokens[i].append(token_id)
                if token_id == tokenizer.tok_to_ind['<EOS>']:
                    mask[i] = False

            input_token = next_tokens
            if not mask.any():
                break

        result_text = []
        for tokens in result_tokens:
            caption = " ".join([tokenizer.ind_to_tok.get(t, "<UNK>") for t in tokens 
                               if t not in [tokenizer.tok_to_ind.get("<BOS>", -1), 
                                           tokenizer.tok_to_ind.get("<EOS>", -1), 
                                           tokenizer.tok_to_ind.get("<PAD>", -1)]])
            result_text.append(caption)

    return result_text[0] if batch_size == 1 else result_text

def generate_caption(image):
    image = image.convert("RGB")
    max_seq_len = full_config.get("max_len", 20)
    caption = generate(model, image, tokenizer, max_seq_len=max_seq_len, top_k=5)
    caption = wrap_text(caption, 43)
    return caption

iface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="Загрузите изображение"),
    outputs=gr.Textbox(label="Сгенерированное описание"),
    title="Генерация описаний изображений",
    description="Загрузите изображение, и модель создаст текстовое описание."
)
iface.launch(
    server_name="0.0.0.0",
    server_port=7860,
)