import torch
import torch._utils
import torch.nn as nn
import torchvision.models as models
from typing import Tuple
from source.config import Config


class Encoder(nn.Module):
    def __init__(self, image_emb_dim: int, device: torch.device):
        """ Image encoder to obtain features from images. Contains pretrained Resnet50 with last layer removed
            and a linear layer with the output dimension of (BATCH, image_emb_dim)
        """

        super(Encoder, self).__init__()
        self.image_emb_dim = image_emb_dim
        self.device = device

        # pretrained Resnet50 model with freezed parameters
        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        for param in resnet.parameters():
            param.requires_grad_(False)

        # remove last layer
        modules = list(resnet.children())[:-1]
        self.resnet = nn.Sequential(*modules)

        # define a final classifier
        self.fc = nn.Linear(in_features=resnet.fc.in_features, out_features=self.image_emb_dim)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        """ Forward operation of encoder, passing images through resnet and then linear layer.

        Args:
            > images (torch.Tensor): (BATCH, 3, 224, 224)

        Returns:
            > features (torch.Tensor): (BATCH, IMAGE_EMB_DIM)
        """

        features = self.resnet(images)
        # features: (BATCH, 2048, 1, 1)

        features = features.reshape(features.size(0), -1).to(self.device)
        # features: (BATCH, 2048)

        features = self.fc(features).to(self.device)
        # features: (BATCH, IMAGE_EMB_DIM)

        return features


class Decoder(nn.Module):
    def __init__(self,
                 image_emb_dim: int,
                 word_emb_dim: int,
                 hidden_dim: int,
                 num_layers: int,
                 vocab_size: int,
                 device: torch.device):
        """
        Decoder taking as input for the LSTM layer the concatenation of features obtained from the encoder
        and embedded captions obtained from the embedding layer. Hidden and cell states are randomly initialized.
        Final classifier is a linear layer with output dimension of the size of a vocabulary.
        """

        super(Decoder, self).__init__()

        self.config = Config()

        self.image_emd_dim = image_emb_dim
        self.word_emb_dim = word_emb_dim
        self.hidden_dim = hidden_dim
        self.num_layer = num_layers
        self.vocab_size = vocab_size
        self.device = device

        self.hidden_state_0 = nn.Parameter(torch.zeros((self.num_layer, 1, self.hidden_dim)))
        self.cell_state_0 = nn.Parameter(torch.zeros((self.num_layer, 1, self.hidden_dim)))

        self.lstm = nn.LSTM(input_size=self.image_emd_dim + self.word_emb_dim,
                            hidden_size=self.hidden_dim,
                            num_layers=self.num_layer,
                            bidirectional=False)

        self.fc = nn.Sequential(
            nn.Linear(in_features=self.hidden_dim, out_features=self.vocab_size),
            nn.LogSoftmax(dim=2)
        )

    def forward(self,
                embedded_captions: torch.Tensor,
                features: torch.Tensor,
                hidden: torch.Tensor,
                cell: torch.Tensor) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        Forward operation of (word-by-word) decoder.
        The LSTM input (concatenation of embedded_captions and features) is passed through LSTM and then linear layer.

        Args:

            > embedded_captions(torch.Tensor): (SEQ_LENGTH = 1, BATCH, WORD_EMB_DIM)
            > features (torch.Tensor): (1, BATCH, IMAGE_EMB_DIM)
            > hidden (torch.Tensor): (NUM_LAYER, BATCH, HIDDEN_DIM)
            > cell (torch.Tensor): (NUM_LAYER, BATCH, HIDDEN_DIM)

        Returns:

            > output (torch.Tensor): (1, BATCH, VOCAB_SIZE)
            > (hidden, cell) (torch.Tensor, torch.Tensor): (NUM_LAYER, BATCH, HIDDEN_DIM)
        """

        lstm_input = torch.cat((embedded_captions, features), dim=2)

        output, (hidden, cell) = self.lstm(lstm_input, (hidden, cell))
        # output : (SEQ_LENGTH, BATCH, HIDDEN_DIM)
        # hidden : (NUM_LAYER, BATCH, HIDDEN_DIM)

        output = output.to(self.device)

        output = self.fc(output)
        # output : (SEQ_LENGTH, BATCH, VOCAB_SIZE)

        return output, (hidden, cell)