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__()
        # Загружаем полную модель ResNet50 с предобученными весами
        self.resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        # Замораживаем все параметры
        for param in self.resnet.parameters():
            param.requires_grad = False
            
    def forward(self, imgs):
        # Используем все слои до avgpool
        x = self.resnet.conv1(imgs)
        x = self.resnet.bn1(x)
        x = self.resnet.relu(x)
        x = self.resnet.maxpool(x)
        
        x = self.resnet.layer1(x)
        x = self.resnet.layer2(x)
        x = self.resnet.layer3(x)
        x = self.resnet.layer4(x)
        
        x = self.resnet.avgpool(x)
        return torch.flatten(x, 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)
        self.num_layers = num_layers

    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.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.text_fe.num_layers, 1, 1)
        cell = torch.zeros_like(hidden)
        outputs, _ = self.text_fe(captions, (hidden, cell))
        outputs = self.fc_out(outputs)
        return outputs