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

class StrongTextCNN(nn.Module):

    def __init__(
        self,
        vocab_size,
        embed_dim=300,
        num_filters=128,
        dropout=0.4
    ):
        super().__init__()

        self.embedding = nn.Embedding(
            vocab_size,
            embed_dim,
            padding_idx=0
        )

        self.convs = nn.ModuleList([
            nn.Conv1d(embed_dim, num_filters, k)
            for k in [2,3,4,5]
        ])

        self.bn = nn.BatchNorm1d(num_filters * 4)

        self.fc1 = nn.Linear(num_filters * 4, 256)
        self.fc2 = nn.Linear(256, 64)
        self.out = nn.Linear(64, 1)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):

        x = self.embedding(x)
        x = x.permute(0,2,1)

        conv_outputs = []

        for conv in self.convs:

            c = F.relu(conv(x))

            p = F.max_pool1d(
                c,
                kernel_size=c.shape[2]
            ).squeeze(2)

            conv_outputs.append(p)

        x = torch.cat(conv_outputs, dim=1)

        x = self.bn(x)

        x = self.dropout(x)

        x = F.relu(self.fc1(x))
        x = self.dropout(x)

        x = F.relu(self.fc2(x))
        x = self.dropout(x)

        x = self.out(x)

        return x.squeeze(1)