from torch.functional import F
import torch.nn as nn


class CNN(nn.Module):
    def __init__(self, dense, channels, kernels, paddings, drop_out,
                 vocab_size, embedding_length, weights):
        super(CNN, self).__init__()

        """
        Arguments
        ---------
        output_size : 2 = (pos, neg)
        in_channels : Number of input channels. Here it is 1 as the input data has dimension = (batch_size, num_seq, embedding_length)
        kernel_heights : A list consisting of 3 different kernel_heights. Convolution will be performed 3 times and finally results from each kernel_height will be concatenated.
        drop_out : Probability of retaining an activation node during dropout operation
        vocab_size : Size of the vocabulary containing unique words
        embedding_length : Embedding dimension of GloVe word embeddings
        weights : Pre-trained GloVe word_embeddings which we will use to create our word_embedding look-up table
        --------

        """
        self.dense = dense
        self.channels = channels
        self.kernels = kernels
        self.paddings = paddings
        self.vocab_size = vocab_size
        self.embedding_length = embedding_length

        self.word_embeddings = self.__create_embedding(weights)
        self.conv1 = self.__create_convolution(0)
        self.conv2 = self.__create_convolution(1)
        self.conv3 = self.__create_convolution(2)
        self.lstm = nn.LSTM(input_size=self.dense[0], hidden_size=self.dense[1], dropout=drop_out)
        self.denses = self.__create_dense()

    def __create_dense(self):
        return nn.Sequential(
            nn.Linear(self.dense[0], self.dense[1]),
            nn.ReLU(),

            nn.Linear(self.dense[1], self.dense[2]),
            nn.ReLU(),

            nn.Linear(self.dense[2], self.dense[3]),
            nn.ReLU(),

            nn.Linear(self.dense[3], self.dense[4]),
            nn.ReLU(),

            nn.Softmax()
        )

    def __create_embedding(self, weights):
        word_embeddings = nn.Embedding(self.vocab_size, self.embedding_length)
        word_embeddings.weight = nn.Parameter(weights, requires_grad=False)
        return word_embeddings

    def __create_convolution(self, step):
        return nn.Sequential(
            nn.Conv2d(
                in_channels=self.channels[step],
                out_channels=self.channels[step+1],
                kernel_size=(self.kernels[step], self.embedding_length if step == 0 else 1),
                padding=(self.paddings[step], 0),
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(self.kernels[step], 1), padding=(1, 0), stride=(1, 1)),
        )

    def forward(self, x, hidden=None):
        """
        The idea of the Convolutional Neural Netwok for Text Classification is very simple. We perform convolution operation on the embedding matrix
        whose shape for each batch is (num_seq, embedding_length) with kernel of varying height but constant width which is same as the embedding_length.
        We will be using ReLU activation after the convolution operation and then for each kernel height, we will use max_pool operation on each tensor
        and will filter all the maximum activation for every channel and then we will concatenate the resulting tensors. This output is then fully connected
        to the output layers consisting two units which basically gives us the logits for both positive and negative classes.

        Parameters
        ----------
        x: input_sentences of shape = (batch_size, num_sequences)
        hidden: embedding from prev step

        Returns
        -------
        Output of the linear layer containing logits for pos & neg class.
        logits.size() = (batch_size, output_size)

        """

        x = self.word_embeddings(x)
        max_out1 = self.conv1(x)
        max_out2 = self.conv2(max_out1)
        max_out3 = self.conv3(max_out2)

        # all_out = torch.cat((max_out1, max_out2, max_out3), 0)
        flatten = max_out3.view(max_out3.shape[0], 1, max_out3.shape[1] * max_out3.shape[2] * max_out3.shape[3])
        lstm, hidden = self.lstm(flatten, hidden)
        output = self.denses(lstm)
        return output, hidden, max_out3