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


class EmbeddingNet(nn.Module):
    def __init__(self):
        super(EmbeddingNet, self).__init__()
        self.convnet = nn.Sequential(nn.Conv1d(7, 64, 3, padding=1), nn.ELU(),
                                     nn.BatchNorm1d(64),
                                     nn.MaxPool1d(2, stride=2),

                                     nn.Conv1d(64, 128, 3, padding=1), nn.ELU(),
                                     nn.BatchNorm1d(128),
                                     nn.MaxPool1d(2, stride=2),

                                     nn.Conv1d(128, 256, 3, padding=1), nn.ELU(),
                                     nn.BatchNorm1d(256),
                                     nn.MaxPool1d(2, stride=2),

                                     nn.Conv1d(256, 512, 3, padding=1), nn.ELU(),
                                     nn.BatchNorm1d(512),
                                     nn.MaxPool1d(2, stride=2)
                                     )

        self.fc = nn.Sequential(nn.Linear(9216, 64),
                                nn.BatchNorm1d(64),
                                nn.ELU(),
                                nn.Dropout(0.4),

                                nn.Linear(64, 32),
                                nn.BatchNorm1d(32),
                                nn.ELU(),

                                # nn.Linear(32, 2),


                                )

    def forward(self, x):
        output = self.convnet(x)
        output = output.view(output.size()[0], -1)
        output = self.fc(output)
        return output

    def get_embedding(self, x):
        return self.forward(x)


class EmbeddingNetL2(EmbeddingNet):
    def __init__(self):
        super(EmbeddingNetL2, self).__init__()

    def forward(self, x):
        output = super(EmbeddingNetL2, self).forward(x)
        output /= output.pow(2).sum(1, keepdim=True).sqrt()
        return output

    def get_embedding(self, x):
        return self.forward(x)


class ClassificationNet(nn.Module):
    def __init__(self, embedding_net, n_classes):
        super(ClassificationNet, self).__init__()
        self.embedding_net = embedding_net
        self.n_classes = n_classes
        self.nonlinear = nn.PReLU()
        self.fc1 = nn.Linear(32, n_classes)

    def forward(self, x):
        output = self.embedding_net(x)
        output = self.nonlinear(output)
        scores = F.log_softmax(self.fc1(output), dim=-1)
        return scores

    def get_embedding(self, x):
        return self.nonlinear(self.embedding_net(x))


class SiameseNet(nn.Module):
    def __init__(self, embedding_net):
        super(SiameseNet, self).__init__()
        self.embedding_net = embedding_net

    def forward(self, x1, x2):
        output1 = self.embedding_net(x1)
        output2 = self.embedding_net(x2)
        return output1, output2

    def get_embedding(self, x):
        return self.embedding_net(x)


class TripletNet(nn.Module):
    def __init__(self, embedding_net):
        super(TripletNet, self).__init__()
        self.embedding_net = embedding_net

    def forward(self, x1, x2, x3):
        output1 = self.embedding_net(x1)
        output2 = self.embedding_net(x2)
        output3 = self.embedding_net(x3)
        return output1, output2, output3

    def get_embedding(self, x):
        return self.embedding_net(x)
