from torch import nn
import torch


class Cnn(nn.Module):
    def __init__(self, length):
        self.length = length
        super(Cnn, self).__init__()
        self.conv1 = nn.Conv1d(1, 32, kernel_size=3, stride=1, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool1d(2)
        self.conv2 = nn.Conv1d(32, 64, kernel_size=3, stride=1, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool1d(2)
        self.conv3 = nn.Conv1d(64, 128, kernel_size=3, stride=1, padding=1)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool1d(2)
        self.fc = nn.Linear(self.length * 32, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = self.relu3(x)
        x = self.pool3(x)
        x = x.view(-1, self.length * 32)
        x = self.fc(x)
        x = self.sigmoid(x)
        return x

    def prosess(self, va, vb):
        tensor = torch.Tensor([va, vb])
        tensor = tensor.reshape([1, 1, len(va) + len(vb)])
        return tensor
