import torch
import numpy as np
import copy
import torch.nn.functional as F
from torch.nn import Parameter

class ScoreFunction(torch.nn.Module):
    def __init__(self, N_words=10) -> None:
        super().__init__()
        self.d_model = 256
        self.n_head = 4
        #self.dropout = torch.nn.Dropout(0.1)
        self.norm = torch.nn.LayerNorm(self.d_model)
        self.Attention = torch.nn.MultiheadAttention(self.d_model, self.n_head, dropout=0)
        self.Linear1 = torch.nn.Linear(N_words,256)
        self.Linear2 = torch.nn.Linear(256,128)
        self.Linear3 = torch.nn.Linear(128,32)
        self.Linear4 = torch.nn.Linear(32,1)
        self.Activation1 = torch.nn.ReLU()
        self.Activation2 = torch.nn.Sigmoid()
        self.register_parameter('bias',Parameter(torch.zeros(1)))

    def forward(self, input):
        #NOTE: input has shape NxN_word
        output = self.Linear1(input)
        output0 = self.Attention(output, output, value=output, attn_mask=None, key_padding_mask=None)[0]
        output = output0*0.2 + output*0.8
        output = self.norm(output)
        output = self.Activation1(output)
        output = self.Linear2(output)
        output = self.Activation1(output)
        output = self.Linear3(output)
        output = self.Activation1(output)
        output = self.Linear4(output) + self.bias
        output = self.Activation2(output)
        return output