import torch
from torch import nn


class Convkb(nn.Module):
    def __init__(self, length):
        self.length = length
        super(Convkb, self).__init__()
        self.conv1 = nn.Conv1d(2, 1, kernel_size=1, stride=1, padding=0)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv1d(2, 1, kernel_size=1, stride=1, padding=0)
        self.relu2 = nn.ReLU()
        self.conv3 = nn.Conv1d(2, 1, kernel_size=1, stride=1, padding=0)
        self.relu3 = nn.ReLU()
        self.fc = nn.Linear(self.length * 3, 1)

    def forward(self, x):
        X1 = self.conv1(x)
        X1 = self.relu1(X1)
        X2 = self.conv2(x)
        X2 = self.relu1(X2)
        X3 = self.conv3(x)
        X3 = self.relu1(X3)
        X = torch.cat([X1, X2, X3], dim=2)
        X = X.view(-1, self.length * 3)
        X = self.fc(X)
        return X

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