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

ALPHA = 10
BETTA = 1
GAMMA = 1.0
M_1 = 0.4
M_2 = 0.8


class CNNEncoder(nn.Module):
    def __init__(self, embed_dim: int, vocab_size: int):
        super(CNNEncoder, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.conv1 = nn.Conv2d(1, 100, (2, embed_dim))

    def forward(self, x):
        x = self.embed(x)
        x = x.unsqueeze(0)
        x = x.unsqueeze(0)
        x = self.conv1(x)
        x = x.squeeze(3)
        x = torch.tanh(x)
        x = F.avg_pool1d(x, x.size()[2]).squeeze(2)
        return x


def cosine_similarity(x, y):
    return F.cosine_similarity(x, y)


class ProtoNet(nn.Module):
    def __init__(self):
        super(ProtoNet, self).__init__()

    def forward(self, x_in, x_ood, support_set, labels):
        prototypes = support_set.mean(1)
        return prototypes


class OProtoNet(nn.Module):
    def __init__(self, cnn_encoder: CNNEncoder, distance_function=cosine_similarity):
        super(OProtoNet, self).__init__()
        self.cnn_encoder = cnn_encoder
        self.proto_network = ProtoNet()
        self.distance_function = distance_function

    def forward(self, x_in, x_ood, support_set, labels):
        x_in = self.cnn_encoder.forward(x_in)
        x_ood = self.cnn_encoder.forward(x_ood)

        encoded_support_set = []

        for class_row in support_set.split(1):
            encoded_sents = []
            for sentence in class_row.squeeze_().split(1):
                encoded_sents.append(self.cnn_encoder.forward(sentence.squeeze_()))
            encoded_support_set.append(encoded_sents)

        class_tensors = []
        for i in range(support_set.shape[0]):
            sentence_tensors = []
            for j in range(support_set.shape[1]):
                sentence_tensors.append(torch.tensor(encoded_support_set[i][j].clone().detach().requires_grad_(True)))
            class_tensors.append(torch.cat(sentence_tensors))

        support_set = torch.stack(class_tensors)

        prototypes = self.proto_network.forward(x_in, x_ood, support_set, labels)

        dists = self.distance_function(x_in, prototypes)
        log_p_y = F.log_softmax(-dists * ALPHA, dim=0)

        loss_in = -log_p_y[0]

        max_dist = torch.max(self.distance_function(x_ood, prototypes) - M_1)

        zero_tensor = torch.FloatTensor([0.0])

        loss_ood = torch.max(max_dist, zero_tensor)

        dist = M_2 - self.distance_function(x_in, prototypes[0].unsqueeze_(0))

        loss_gt = torch.max(zero_tensor.expand_as(dist), dist)

        return loss_in + BETTA * loss_ood.item() + GAMMA * loss_gt.item()