from torch import nn
import torch


class Doubletower(nn.Module):
    def __init__(self, length):
        self.length = length
        super(Doubletower, self).__init__()
        self.conv1l = nn.Conv1d(1, 32, kernel_size=3, stride=1, padding=1)
        self.relu1l = nn.ReLU()
        self.pool1l = nn.MaxPool1d(2)
        self.conv2l = nn.Conv1d(32, 64, kernel_size=3, stride=1, padding=1)
        self.relu2l = nn.ReLU()
        self.pool2l = nn.MaxPool1d(2)
        self.conv3l = nn.Conv1d(64, 128, kernel_size=3, stride=1, padding=1)
        self.relu3l = nn.ReLU()
        self.pool3l = nn.MaxPool1d(2)
        self.fcl = nn.Linear(self.length * 16, 256)
        self.conv1r = nn.Conv1d(1, 32, kernel_size=3, stride=1, padding=1)
        self.relu1r = nn.ReLU()
        self.pool1r = nn.MaxPool1d(2)
        self.conv2r = nn.Conv1d(32, 64, kernel_size=3, stride=1, padding=1)
        self.relu2r = nn.ReLU()
        self.pool2r = nn.MaxPool1d(2)
        self.conv3r = nn.Conv1d(64, 128, kernel_size=3, stride=1, padding=1)
        self.relu3r = nn.ReLU()
        self.pool3r = nn.MaxPool1d(2)
        self.fcr = nn.Linear(self.length * 16, 256)
        self.sigmoid = nn.Sigmoid()

    def forward(self, X):
        x = self.conv1l(X[0])
        x = self.relu1l(x)
        x = self.pool1l(x)
        x = self.conv2l(x)
        x = self.relu2l(x)
        x = self.pool2l(x)
        x = self.conv3l(x)
        x = self.relu3l(x)
        x = self.pool3l(x)
        x = x.view(-1, self.length * 16)
        x = self.fcl(x)
        x = self.sigmoid(x).reshape(256)
        y = self.conv1r(X[1])
        y = self.relu1r(y)
        y = self.pool1r(y)
        y = self.conv2r(y)
        y = self.relu2r(y)
        y = self.pool2r(y)
        y = self.conv3r(y)
        y = self.relu3r(y)
        y = self.pool3r(y)
        y = y.view(-1, self.length * 16)
        y = self.fcr(y)
        y = self.sigmoid(y).reshape(256)
        X = torch.dot(x, y)
        return X

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