import torch

from torch import nn

import os

class LinearModel(nn.Module):

    def __init__(self, label_count, embedding_type, embedding, freeze_embedding=True):

        super(LinearModel, self).__init__()

        self.embedding_type = embedding_type

        if self.embedding_type in ["word", "cui"]:

            self.embedding = nn.Embedding.from_pretrained(embedding)

            self.input_dim = self.embedding.weight.shape[1]

            if freeze_embedding:

                self.embedding.weight.required_grad = False

        if self.embedding_type == "bert":

            self.embedding = embedding

            self.input_dim = 768

            if freeze_embedding:

                for name, param in self.embedding.named_parameters():

                    param.requires_grad = False

        self.linear = nn.Linear(self.input_dim * 2, label_count)

        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, x0, x1, length_0=None, length_1=None, label=None):

        count = x0.shape[0]

        x = torch.cat((x0, x1), dim=0)

        emb = self.embedding(x)

        #print(x.shape, emb.shape, length_0.shape)

        if self.embedding_type == "word":

            emb = torch.sum(emb, dim=1)

            length = torch.cat((length_0, length_1)).reshape(-1, 1).expand_as(emb)

            emb = emb / length

        if self.embedding_type == "cui":

            pass

        if self.embedding_type == "bert":

            emb = emb[1]

        emb_0 = emb[0:count]

        emb_1 = emb[count:]

        feature = torch.cat((emb_0, emb_1), dim=1)

        pred = self.linear(feature)

        if label is not None:

            loss = self.loss_fn(pred, label)

            return pred, loss

        return pred, 0.