import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import copy
import numpy as np

def clones(module, N):
    '''Produce N identical layers.'''
    '''当添加 nn.ModuleList 作为 nn.Module 对象的一个成员时（即当我们添加模块到我们的网络时），所有 nn.ModuleList 内部的 nn.Module 的 parameter 也被添加作为 我们的网络的 parameter'''
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])  

class AttentionLayer(nn.Module):
    def __init__(self, embed_size, h, is_share=False, drop=0.0):
        super(AttentionLayer, self).__init__()
        self.is_share = is_share
        self.h = h
        self.embed_size = embed_size # 300
        self.d_k = embed_size // h
        self.drop_p = drop
        # domain的数目
        self.num_domain = 6
        if is_share:
            self.linear = nn.Linear(embed_size, embed_size)
            self.linears = [self.linear, self.linear, self.linear] 
        else:
            # 3层全连接---将原始fea映射成Q,K,V
            self.linears = clones(nn.Linear(embed_size, embed_size), 3)
        if self.drop_p > 0:
            self.dropout = nn.Dropout(drop)
        # self.linear_soft = nn.Linear(embed_size, self.num_domain)
        self.linear_soft = nn.Linear(embed_size, 1)
        

    def forward(self, inp, mask=None):
        nbatches = inp.size(0)
        inp_transfer = inp.view(-1, self.embed_size)
        scores1 = self.linear_soft(inp_transfer).view(nbatches, self.num_domain, 1)
        # scores2:  torch.Size([64, 6, 1])
        scores2 = F.softmax(scores1, dim=1)
        scores3 = scores2.repeat(1, 1, 300)
        weighted_repre = torch.mul(inp, scores3)
        # print('scores2: ', scores2)
        # print('scores2: ', np.shape(scores2))
        # exit()

        # query, key, value = \
        #     [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
        #      for l, x in zip(self.linears, (inp, inp, inp))]
        
        # scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k)     
        # if mask is not None:
        #     scores = scores.masked_fill(mask == 0, -1e9)
        # p_attn = F.softmax(scores, dim=-1)
        # if self.drop_p > 0:
        #     p_attn = self.dropout(p_attn)
        # x = torch.matmul(p_attn, value) 
        # x = x.transpose(1, 2).contiguous() \
        #      .view(nbatches, -1, self.h * self.d_k)
        return weighted_repre

class FeedForward(nn.Module):
    def __init__(self, embed_size, hidden, drop=0.0):
        super(FeedForward, self).__init__()
        self.fc1 = nn.Linear(embed_size, hidden)
        self.fc2 = nn.Linear(hidden, embed_size)
        self.dropout = nn.Dropout(drop)

    def forward(self, x):
        return self.fc2(self.dropout(F.relu(self.fc1(x))))

class SelfAttention(nn.Module):
    def __init__(self, embed_size, hid_size, h, drop=0.0):
        super(SelfAttention, self).__init__()
        self.h = h # num of heads --16
        self.att_layer = AttentionLayer(embed_size, h, drop=drop)
        self.feed_forward_layer = FeedForward(embed_size, hid_size, drop=drop)
        self.dropout = nn.Dropout(drop)

    def forward(self, local_emb, mask=None):
        weighted_repre = self.att_layer(local_emb, mask=mask)  
        # exit()
        # self_att_emb = local_emb + self.dropout(x)
        # x = self.feed_forward_layer(self_att_emb)      
        # fw_emb = self_att_emb + self.dropout(x)

        # return fw_emb
        return weighted_repre

