import torch
import torch.nn as nn
from torch.nn.functional import max_pool1d

from utility.model_parameter import Configuration, ModelParameter


class CNNLayer(nn.Module):
    def __init__(self, config: Configuration, vocab_size=30000, use_embeddings=True, embed_dim=-1, **kwargs):
        super(CNNLayer, self).__init__()

        # set parameters
        self.max_seq_length = config.get_int(ModelParameter.MAX_LENGTH)
        self.use_gpu = torch.cuda.is_available()
        if embed_dim == -1:
            self.embedding_dim = config.get_int(ModelParameter.EMBEDDING_SIZE)
        else:
            self.embedding_dim = embed_dim
        self.max_length = config.get_int(ModelParameter.MAX_LENGTH)
        self.use_embeddings = use_embeddings
        self.conv_out_channels = config.get_int(ModelParameter.CHANNELS)
        self.filter_sizes = [2]

        # create and initialize layers
        self.embedding = nn.Embedding(vocab_size, self.embedding_dim)
        self.relu = nn.ReLU()
        self.convolutions = nn.ModuleList(
            [nn.Conv2d(1, self.conv_out_channels, (K, self.embedding_dim)) for K in self.filter_sizes])
        self.dropout = nn.Dropout(0.3)

    def get_output_length(self):
        return len(self.filter_sizes) * self.conv_out_channels

    def forward(self, samples, **kwargs):
        encoded_samples = self.encode(samples)
        return encoded_samples

    def encode(self, samples):
        x = self.embedding(samples)
        x = x.unsqueeze(1)
        x = [self.relu(conv(x)).squeeze(3) for conv in self.convolutions]
        x = [max_pool1d(i, i.size(2)).squeeze(2) for i in x]
        x = self.dropout(torch.cat(x, 1))
        return x

[((14, 23, 14, 48), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((25, 25, 25, 69), 'torch.nn.Embedding', 'nn.Embedding', ({(25, 38, 25, 48): 'vocab_size', (25, 50, 25, 68): 'self.embedding_dim'}, {}), '(vocab_size, self.embedding_dim)', True, 'import torch.nn as nn\n'), ((26, 20, 26, 29), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((29, 23, 29, 38), 'torch.nn.Dropout', 'nn.Dropout', ({(29, 34, 29, 37): '0.3'}, {}), '(0.3)', True, 'import torch.nn as nn\n'), ((43, 25, 43, 40), 'torch.cat', 'torch.cat', ({(43, 35, 43, 36): 'x', (43, 38, 43, 39): '1'}, {}), '(x, 1)', False, 'import torch\n'), ((28, 13, 28, 74), 'torch.nn.Conv2d', 'nn.Conv2d', ({(28, 23, 28, 24): '1', (28, 26, 28, 48): 'self.conv_out_channels', (28, 50, 28, 73): '(K, self.embedding_dim)'}, {}), '(1, self.conv_out_channels, (K, self.embedding_dim))', True, 'import torch.nn as nn\n')]