import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as T
import gradio as gr
from PIL import Image, ImageOps

import os

# On ZeroGPU Spaces the `spaces` package exists; elsewhere fall back to a no-op decorator
try:
    import spaces
except ImportError:  # not on a ZeroGPU Space: provide a do-nothing stand-in
    class spaces:
        @staticmethod
        def GPU(fn):
            return fn

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

# Must match the values used in training
EMBED_DIM, ATTENTION_DIM, DECODER_DIM = 256, 256, 512
CKPT_PATH = "best_model.pt"


class EncoderCNN(nn.Module):
    def __init__(self):
        super().__init__()
        # weights=None: the trained ResNet weights come from the checkpoint
        resnet = models.resnet50(weights=None)
        self.resnet = nn.Sequential(*list(resnet.children())[:-2])

    def forward(self, images):
        features = self.resnet(images)
        b, c, h, w = features.shape
        return features.view(b, c, h * w).permute(0, 2, 1)


class BahdanauAttention(nn.Module):
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super().__init__()
        self.W_encoder = nn.Linear(encoder_dim, attention_dim)
        self.W_decoder = nn.Linear(decoder_dim, attention_dim)
        self.V = nn.Linear(attention_dim, 1)

    def forward(self, encoder_out, decoder_hidden):
        wh = self.W_encoder(encoder_out)
        ws = self.W_decoder(decoder_hidden).unsqueeze(1)
        scores = self.V(torch.tanh(wh + ws)).squeeze(2)
        weights = F.softmax(scores, dim=1)
        context = (encoder_out * weights.unsqueeze(2)).sum(dim=1)
        return context, weights


class DecoderWithAttention(nn.Module):
    def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size, encoder_dim=2048):
        super().__init__()
        self.vocab_size = vocab_size
        self.attention = BahdanauAttention(encoder_dim, decoder_dim, attention_dim)
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.dropout = nn.Dropout(0.5)
        self.lstm_cell = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim)
        self.fc = nn.Linear(decoder_dim, vocab_size)


class ImageCaptioner(nn.Module):
    def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size):
        super().__init__()
        self.encoder = EncoderCNN()
        self.decoder = DecoderWithAttention(attention_dim, embed_dim, decoder_dim, vocab_size)
        self.init_h = nn.Linear(2048, decoder_dim)
        self.init_c = nn.Linear(2048, decoder_dim)

    def init_hidden_state(self, encoder_out):
        mean = encoder_out.mean(dim=1)
        return torch.tanh(self.init_h(mean)), torch.tanh(self.init_c(mean))


# ---- load checkpoint ----
if not os.path.exists(CKPT_PATH):  # e.g. running from the GitHub repo without the weights
    from huggingface_hub import hf_hub_download
    CKPT_PATH = hf_hub_download(repo_id="birukale/Image-Caption", filename="best_model.pt", repo_type="space")

ckpt = torch.load(CKPT_PATH, map_location="cpu")
itos = ckpt["itos"]
stoi = {w: i for i, w in itos.items()}
model = ImageCaptioner(ATTENTION_DIM, EMBED_DIM, DECODER_DIM, len(itos)).to(device)
model.load_state_dict(ckpt["model"])
model.eval()

transform = T.Compose([
    T.Resize((224, 224)), T.ToTensor(),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])


@spaces.GPU
@torch.no_grad()
def caption_image(image, max_len=30):
    if image is None:
        return "Please upload an image."
    image = ImageOps.exif_transpose(image).convert("RGB")
    img = transform(image).unsqueeze(0).to(device)

    enc = model.encoder(img)
    h, c = model.init_hidden_state(enc)
    word = torch.tensor([stoi["<start>"]], device=device)

    out = []
    for _ in range(max_len):
        emb = model.decoder.embedding(word)
        context, _ = model.decoder.attention(enc, h)
        h, c = model.decoder.lstm_cell(torch.cat([emb, context], dim=1), (h, c))
        word = model.decoder.fc(h).argmax(dim=1)
        if word.item() == stoi["<end>"]:
            break
        out.append(itos[word.item()])
    return " ".join(out)


demo = gr.Interface(
    fn=caption_image,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Textbox(label="Generated caption"),
    title="Image Captioning: ResNet-50 + Bahdanau Attention + LSTM",
    description="Trained from scratch on Flickr8k. Works best on everyday scenes: people, dogs, sports, streets, outdoors.",
)

if __name__ == "__main__":
    demo.launch()