import torch
import torch.nn as nn
import torchvision.models as models

class img_fe_class(nn.Module):
    def __init__(self):
        super(img_fe_class, self).__init__()
        weights = models.ResNet50_Weights.DEFAULT
        self.resnet = models.resnet50(weights=weights)
        self.resnet = nn.Sequential(*list(self.resnet.children())[:-1])
        for param in self.resnet.parameters():
            param.requires_grad = False

    def forward(self, imgs):
        features = self.resnet(imgs)
        return features.reshape(features.size(0), -1)

class text_fe_class(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_size, num_layers):
        super(text_fe_class, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.rnn = nn.LSTM(embedding_dim, hidden_size, num_layers, batch_first=True)

    def forward(self, captions, hidden):
        embeddings = self.embedding(captions)
        outputs, hidden = self.rnn(embeddings, hidden)
        return outputs, hidden

class image_captioning_model(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_size, num_layers):
        super(image_captioning_model, self).__init__()
        self.img_fe = img_fe_class()
        self.text_fe = text_fe_class(vocab_size, embedding_dim, hidden_size, num_layers)
        self.fc_out = nn.Linear(hidden_size, vocab_size)
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.img_to_hidden = nn.Linear(2048, hidden_size)  # 2048 - размер выхода ResNet50

    def forward(self, imgs, captions):
        img_features = self.img_fe(imgs)
        hidden = self.img_to_hidden(img_features).unsqueeze(0).repeat(self.num_layers, 1, 1)
        cell = torch.zeros_like(hidden)  # Для LSTM нужно также cell state
        outputs, _ = self.text_fe(captions, (hidden, cell))
        outputs = self.fc_out(outputs)
        return outputs

    def generate_caption(self, img, tokenizer, max_len=30, device="cuda"):
        self.eval()
        with torch.no_grad():
            img = img.unsqueeze(0).to(device)
            img_features = self.img_fe(img)
            hidden = self.img_to_hidden(img_features).unsqueeze(0).repeat(self.num_layers, 1, 1)
            cell = torch.zeros_like(hidden)
            caption = [tokenizer.tok_to_ind["<BOS>"]]
            for _ in range(max_len):
                inputs = torch.LongTensor(caption[-1]).unsqueeze(0).unsqueeze(0).to(device)
                outputs, (hidden, cell) = self.text_fe(inputs, (hidden, cell))
                outputs = self.fc_out(outputs[:, -1, :])
                next_token = outputs.argmax(dim=-1).item()
                caption.append(next_token)
                if next_token == tokenizer.tok_to_ind["<EOS>"]:
                    break
            return tokenizer.decode(caption)