import torch
from torch import nn
from torchvision.models import resnet18

from config import GlobalConfig
from config.model_config import ImageEncoderConfig


class ImageEncoder(nn.Module):
    """Image Encoder."""

    def __init__(self, config: ImageEncoderConfig):
        super(ImageEncoder, self).__init__()
        # 使用预训练好的ResNet作为CNN
        # --->pretrained(bool)：If True, returns a model pre-trained on ImageNet;
        self.cnn = resnet18(pretrained=True)
        self.cnn = nn.Sequential(*list(self.cnn.children())[:-1])
        self.cnn = self.cnn.to(GlobalConfig.device)
        self.image_size = config.image_size
        self.num_channels = config.num_channels
        image_feat_size = (self.image_size ** 2) * self.num_channels
        self.linear = nn.Linear(config.text_feat_size, image_feat_size)
        self.linear = self.linear.to(GlobalConfig.device)
        # dim = 2--->表征对于第2维开始操作(dim从0开始计数)
        self.softmax = nn.Softmax(dim=2).to(GlobalConfig.device)

    def forward(self, image, encoded_text):
        """

        Args:
            image: (batch_size, 3, image_size, image_size)
            encoded_text: (batch_size,
                           num_layers * num_directions * hidden_size)

        Returns:
            encoded_image: (batch_size, image_feat_size)

        """
        batch_size = image.size(0)
        # CNN编码图片
        image = self.cnn(image)
        # (batch_size, self.num_channels, self.image_size, self.image_size)
        # 输入文本，输出图片维度
        score = self.linear(encoded_text)
        # (batch_size, self.num_channels * self.image_size * self.image_size)
        # view函数的作用为重构张量的维度，相当于numpy中resize---相当于reshape
        score = score.view(batch_size, self.num_channels,
                           self.image_size * self.image_size)
        # (batch_size, self.num_channels, self.image_size * self.image_size)
        # 相当于使用文本表示，对于图片分配一个权重，这种做法好奇怪，为什么不进行文本和图片的交互呢
        score = self.softmax(score)
        # SOFTMAX IS PERFORMED TO DIM 2

        image = image.view(
            (batch_size, self.num_channels, self.image_size * self.image_size))
        # (batch_size, self.num_channels, self.image_size * self.image_size)
        # 加权之后的图片表征
        image = torch.mul(image, score)
        # (batch_size, self.num_channels, self.image_size * self.image_size)
        # 对第2维进行操作
        image_feat = image.sum(2)
        # (batch_size, num_channels)
        return image_feat


if __name__ == '__main__':
    cnn = resnet18(pretrained=True)
