import torch
import torch.nn as nn
from config.global_config import GlobalConfig
from config.model_config import TextEncoderConfig
from torch.nn import Module, Parameter
import numpy as np
# from layers import HGNN_conv

def trans_to_cuda(variable):
    if torch.cuda.is_available():
        return variable.cuda()
    else:
        return variable

def _convert_sp_mat_to_sp_tensor(X):
    # coo = X.tocoo().astype(np.float32)
    coo = X.tocoo()
    indices = np.vstack((coo.row, coo.col))
    # print('indices: ', indices)
    # indices = np.mat([coo.row, coo.col]).transpose()
    i = torch.LongTensor(indices)
    v = torch.FloatTensor(coo.data) 
    # shape = torch.Size(coo)
    shape = coo.shape
    # print('shape: ', shape)
    sp_tensor = torch.sparse.FloatTensor(i, v, torch.Size(shape))
    # 第三个参数不能是shape--应该是原始矩阵coo
    # sp_tensor = torch.sparse.FloatTensor(i, v, shape)
    # print('----sp_tensor---------')
    # return tf.SparseTensor(indices, coo.data, coo.shape)
    # return torch.sparse.FloatTensor(indices, coo.data, coo.shape)
    return sp_tensor


class HyperConv(Module):
    def __init__(self, config: TextEncoderConfig, subgraphs_G, reversed_subgraphs_mapping_w, emb_size=300):
        super(HyperConv, self).__init__()
        self.emb_size = emb_size
        self.layers = config.num_layers
        self.embedding = nn.Embedding(
            config.vocab_size,
            config.embed_size,
            padding_idx=config.pad_index)
        # 注意每一条都要to到对应的device
        self.embedding = self.embedding.to(GlobalConfig.device)
        # Initial embedding.
        # if config.embed_init is not None:
        #     # 加载预训练好的模型，freeze：若为True，表示训练过程不更新，默认为True，False表示表征随着训练过程一起更新
        #     self.embedding = self.embedding.from_pretrained(
        #         config.embed_init, freeze=False)
        # 数据集中的全部context都被转换了---接下来使用index取出每个context对应的邻接矩阵
        self.context_len = config.train_len

        self.subgraphs_G = {}
        for i in range(self.context_len):
            self.subgraphs_G[i] = {}
            self.subgraphs_G[i]['G'] = _convert_sp_mat_to_sp_tensor(subgraphs_G[i]['G'])
            self.subgraphs_G[i]['E'] = _convert_sp_mat_to_sp_tensor(subgraphs_G[i]['E'])
            # print('Success! ' +str(i))
        # self.subgraphs_G = self.subgraphs_G.to(GlobalConfig.device)
        # self.subgraphs_G = self.subgraphs_G.to(GlobalConfig.device)
        self.reversed_subgraphs_mapping_w = reversed_subgraphs_mapping_w


    def forward(self, index_list):
        # print('index_list:', index_list)
        # batch_context_hyrepre = []
        # btach_index---对于每个context取出包含单词的embedding以及邻接矩阵
        index_hy = 0
        batch_len_utt = []
        batch_len_word = []

        for context_index in index_list:
            # print('context_index:', context_index)
            adjacency_G = self.subgraphs_G[context_index]['G']
            adjacency_E = self.subgraphs_G[context_index]['E']
            word_num = np.shape(adjacency_E)[1]
            utt_num = np.shape(adjacency_E)[0]
            # print('word_num: ', self.reversed_subgraphs_mapping_w[context_index])
            # print('adjacency_G: ', np.shape(adjacency_G))
            # print('adjacency_E: ', np.shape(adjacency_E))
            # exit()
            word_index_list = torch.tensor(self.reversed_subgraphs_mapping_w[context_index])
            word_index_list = word_index_list.to(GlobalConfig.device)
            # print('-----adjacency_word_success!-----')
            # 获取单词初始embedding
            # word_embeddings = torch.take(self.embedding, torch.tensor(word_index_list))
            word_embeddings =self.embedding(word_index_list)
            # print(np.shape(word_embeddings))
            # exit()
            for i in range(self.layers):
                word_embeddings = torch.sparse.mm(trans_to_cuda(adjacency_G), word_embeddings)
                edge_embeddings = torch.sparse.mm(trans_to_cuda(adjacency_E), word_embeddings)
            # edge_embedding肯定是10个---每个context最多有10个--如果不够10个，则补全，不然的话，则截断
            # 每个utterance中最多有30个单词
            # print('word_embeddings:', np.shape(word_embeddings))
            # print('edge_embeddings:', np.shape(edge_embeddings))

            if index_hy == 0:
                # batch_edge_hyrepre = edge_embeddings
                # batch_node_hyrepre = word_embeddings
                batch_edge_hyrepre = torch.unsqueeze(edge_embeddings, 0)
                # batch_node_hyrepre = torch.unsqueeze(word_embeddings, 0) 
                # print('batch_node_hyrepre:', np.shape(batch_node_hyrepre))
                # print('batch_edge_hyrepre:', np.shape(batch_edge_hyrepre))
            else:
                # batch_edge_hyrepre = torch.cat((batch_edge_hyrepre, edge_embeddings), 0)
                # batch_node_hyrepre = torch.cat((batch_node_hyrepre, word_embeddings), 0)
                batch_edge_hyrepre = torch.cat((batch_edge_hyrepre, torch.unsqueeze(edge_embeddings, 0)), 0)
                # batch_node_hyrepre = torch.cat((batch_node_hyrepre, torch.unsqueeze(word_embeddings, 0)), 0)
                # print('batch_node_hyrepre:', np.shape(batch_node_hyrepre))
                # print('batch_edge_hyrepre:', np.shape(batch_edge_hyrepre))
            index_hy = index_hy + 1
            # exit()
            batch_len_utt.append(utt_num)
            batch_len_word.append(word_num)
        # return batch_edge_hyrepre, batch_node_hyrepre, batch_len_utt, batch_len_word
        return batch_edge_hyrepre

class domainHyper(Module):
    def __init__(self, config:TextEncoderConfig, emb_size_global=768, emb_size_utter=300):
        super(domainHyper, self).__init__()
        self.emb_size_global = emb_size_global
        self.emb_size_utter = emb_size_utter
        self.domain_layers = config.num_layers
        # self.subgraphs_domain_node = _convert_sp_mat_to_sp_tensor(subgraphs_domain_node)
        # self.subgraphs_domain_edge = subgraphs_domain_edge
        # self.train_global_repre = torch.stack([torch.LongTensor(each_repre) for each_repre in train_global_repre])
        # self.train_global_repre = self.train_global_repre.float()
        # self.train_global_repre = self.train_global_repre.to(GlobalConfig.device)
        # self.subgraphs_domain_node = self.subgraphs_domain_node.to(GlobalConfig.device)
        self.glu_domain = nn.Linear(self.emb_size_global, self.emb_size_utter)
    
    def forward(self, global_repre, fusion_repre, index_list, subgraphs_domain_node):
        # 768--300
        # 23369, 300
        self.subgraphs_domain_node = _convert_sp_mat_to_sp_tensor(subgraphs_domain_node)
        global_repre2 = self.glu_domain(global_repre)
        # print(np.shape(global_repre2))
        # self.global_repre2 = self.glu_domain(self.train_global_repre)
        # print('global_repre2:', np.shape(global_repre2))

        # 替换值
        index = torch.LongTensor(index_list)
        index = index.unsqueeze(1)
        index = index.to(GlobalConfig.device)

        # 使用fusion的feature替换原始数据中的表示
        for i in range(len(index_list)):
            context_index = index_list[i]
            fusion_repre_specific = fusion_repre[i]
            global_repre2[context_index] = fusion_repre_specific

        # print('global_repre2:', np.shape(global_repre2))

        # print('index:', index)
        # content_select1 = torch.index_select(global_repre2, 0, index)
        # print('content_select1_before_replace:', np.shape(content_select1))
        # 23369, 300
        # global_repre2.scatter(0, index, fusion_repre)
        # global_repre2 = global_repre2.to(GlobalConfig.device)
        # content_select2 = torch.index_select(global_repre2, 0, index)
        # print('content_select1_after_replace:', np.shape(content_select2))
        # print(np.shape(global_repre2))
        # exit()
        # 超图卷积
        for i in range(self.domain_layers):
            # self.global_repre2 = torch.sparse.mm(trans_to_cuda(self.subgraphs_domain_node), self.global_repre2)
            global_repre2 = torch.sparse.mm(trans_to_cuda(self.subgraphs_domain_node), global_repre2)
        # 更新之后的dialogue的表征
        # 64, 300
        content_update = torch.index_select(global_repre2, 0, torch.tensor(index_list))
        # print('content_update:', np.shape(content_update))
        return content_update




class UtterFusion(Module):
    def __init__(self, config: TextEncoderConfig, batch_size =64, emb_size_local=768, emb_size_global=300):
        super(UtterFusion, self).__init__()
        self.emb_size_local = emb_size_local
        self.emb_size_global = emb_size_global
        self.batch_size = batch_size
        self.pos_embedding = nn.Embedding(200, self.emb_size_global)
        # self.w_2 = nn.Parameter(torch.Tensor(self.emb_size_global, 1))
        self.w_1 = nn.Parameter(torch.Tensor(2 * self.emb_size_global, self.emb_size_global))
        self.w_2 = nn.Parameter(torch.Tensor(self.emb_size_global, 2))
        # self.w_2 = self.w_2.to(GlobalConfig.device)
        # self.w_2 = trans_to_cuda(self.w_2)
        self.glu1 = nn.Linear(10*self.emb_size_local, 10*self.emb_size_global)
        # self.glu1 = self.glu1.to(GlobalConfig.device)
        # self.glu1 = trans_to_cuda(self.glu1)
        self.glu2 = nn.Linear(10*self.emb_size_global, 10*self.emb_size_global) 
        # self.glu2 = self.glu2.to(GlobalConfig.device)
        # self.glu2 = trans_to_cuda(self.glu2)
    
    def forward(self, local_texts_vec, edge_embeddings):
        # print('-------local_texts_vec1 begin----------------')
        # local_texts_vec = local_texts_vec.view(self.batch_size, -1)
        local_texts_vec = local_texts_vec.view(-1, 10*self.emb_size_local)
        local_texts_vec1 = self.glu1(local_texts_vec)
        # local_texts_vec1 = self.glu1((local_texts_vec))
        # print('-------edge_embeddings1 begin----------------')
        # edge_embeddings = edge_embeddings.view(self.batch_size, -1)
        edge_embeddings = edge_embeddings.view(-1, 10*self.emb_size_global)
        edge_embeddings1 = self.glu2(edge_embeddings)
        nh = torch.sigmoid(local_texts_vec1 + edge_embeddings1) #激活函数与论文中的tanh不是一个
        # 64,10,2 --- 64,10,300/64,10,300
        # nh = nh.view(self.batch_size, 10 , -1)
        nh = nh.view(-1, 10, self.emb_size_global)
        # 64, 10, 2
        # print('torch.matmul(nh, self.w_2):', np.shape(torch.matmul(nh, self.w_2)))
        beta = torch.softmax(torch.matmul(nh, self.w_2), dim=-1)
        # print('beta:', np.shape(beta))
        local_beta =beta[:,:,0]
        edge_beta = beta[:,:,1]
        # 64,10,30
        # local_texts_vec2 = local_texts_vec1.view(self.batch_size, 10, -1)
        local_texts_vec2 = local_texts_vec1.view(-1, 10, 300)
        edge_embeddings2 = edge_embeddings1.view(-1, 10, 300)
        # edge_embeddings2 = edge_embeddings1.view(self.batch_size, 10, -1)
        # 64, 10, 1
        local_beta2 = local_beta[:, :, None]
        edge_beta2 = edge_beta[:, :, None]

        local_beta3 = local_beta2.repeat(1, 1, 300)
        # local_beta3 = local_beta2.repeat(1, 1, 30)
        # edge_beta3 = edge_beta2.repeat(1, 1, 30)
        edge_beta3 = edge_beta2.repeat(1, 1, 300)

        # print('local_beta3:', np.shape(local_beta3))
        # print('edge_beta3:', np.shape(edge_beta3))

        # utter_fusion1 = local_texts_vec1*local_beta[:, :, None] + edge_embeddings1*edge_beta[:, :, None]  
        # utter_fusion1 = local_texts_vec2*local_beta[:, :, None] + edge_embeddings2*edge_beta[:, :, None]  
        # print('local_texts_vec2:', np.shape(local_texts_vec2))
        # print('local_beta3:', np.shape(local_beta3))
        # print('edge_embeddings2:', np.shape(edge_embeddings2))
        # print('edge_beta3:', np.shape(edge_beta3))
        # print('torch.mul(local_texts_vec2, local_beta3):', np.shape(torch.mul(local_texts_vec2, local_beta3)))
        # print('torch.mul(edge_embeddings2, edge_beta3):', np.shape(torch.mul(edge_embeddings2, edge_beta3)))

        utter_fusion1 = torch.mul(local_texts_vec2, local_beta3) +  torch.mul(edge_embeddings2, edge_beta3)
        len_utter = utter_fusion1.shape[1]
        pos_emb = self.pos_embedding.weight[:len_utter]   
        batch_len = utter_fusion1.shape[0]  
        pos_emb = pos_emb.unsqueeze(0).repeat(batch_len, 1, 1)
        # position和embedding拼接之后，是600--再映射成300--64，10，300
        utter_fusion2 = torch.matmul(torch.cat([pos_emb, utter_fusion1], -1), self.w_1)
        utter_fusion3 = utter_fusion2.sum(axis=1)
        # beta = torch.matmul(nh, self.w_2)
        return utter_fusion3

class TextEncoder(nn.Module):
    """ Text encoder."""
    # 这里的cofig指的是TextEncoderConfig
    def __init__(self, config: TextEncoderConfig):
        super(TextEncoder, self).__init__()

        # torch.nn.Embedding(numembeddings,embeddingdim)是创建一个词嵌入模型，numembeddings代表一共有多少个词, embedding_dim代表你想要为每个词创建一个多少维的向量来表示它
        # 所有vocab中为pad_index的词对应的词向量都是0
        self.embedding = nn.Embedding(
            config.vocab_size,
            config.embed_size,
            padding_idx=config.pad_index)
        # 注意每一条都要to到对应的device
        self.embedding = self.embedding.to(GlobalConfig.device)

        # Initial embedding.
        if config.embed_init is not None:
            # 加载预训练好的模型，freeze：若为True，表示训练过程不更新，默认为True，False表示表征随着训练过程一起更新
            self.embedding = self.embedding.from_pretrained(
                config.embed_init, freeze=False)

        # LSTM.
        self.rnn = nn.LSTM(
            input_size=config.embed_size,
            hidden_size=config.hidden_size,
            num_layers=config.num_layers,
            bidirectional=config.num_directions == 2,
            batch_first=False,
            dropout=config.dropout)
        self.rnn = self.rnn.to(GlobalConfig.device)
        # self.HyperGraph = HyperConv(config.n_hyper)

        # 根据index--抽取每个context中每条utterance中对应单词的表征
        # self.xx = torch.take(self.embedding, config.reversed_subgraphs_mapping_w) 

        # HGNN
        # 输入：300 / 输出：300

        # self.nodes, self.edges = HGNN_conv(input_dim = 300,
        #                         output_dim = 300,
        #                         adj = config.subgraphs_G, n_hyper = config.n_hyper,  dropout = config.dropout, act = nn.ReLU)(xx)   
        # self.nodes = self.nodes.to(GlobalConfig.device)
        # self.edges = self.edges.to(GlobalConfig.device)
                 
        # nodes, edges = HGNN_conv(input_dim = self.args.emsize,
        #                         output_dim = self.args.emsize,
        #                         adj = self.subgraphs_G[i],
        #                         act = tf.nn.relu,
        #                         dropout = self.dropout_graph,
        #                         n_hyper = self.n_hyper)(xx)        


    #怎么对它调用呢，是直接使用TextEncoder类，就可以调用了吗
    # BiLSTM编码每一条utterance
    def forward(self, input_seq, input_lengths):
        """Forward.

        Args:
            input_seq: (batch_size, seq_len)
            input_lengths: (batch_size, )

        Returns:
            output: (batch_size, num_layers * num_directions * hidden_size)

        """
        batch_size = input_lengths.size(0)

        input_seq = input_seq.transpose(0, 1)
        # (seq_len, batch_size)

        # 取出embedding
        # 取出邻接矩阵
        embedded = self.embedding(input_seq)
        # (seq_len, batch_size, embed_size)

        hiddens, (h_n, _) = self.rnn(embedded)
        # hiddens: (seq_len, batch, num_directions * hidden_size)
        # h_n: (num_layers * num_directions, batch_size, hidden_size)
        h_n = h_n.transpose(0, 1)
        # h_n: (batch_size, num_layers * num_directions, hidden_size)
        # view函数相当于reshape
        output = h_n.contiguous().view(batch_size, -1)
        # (batch_size, num_layers * num_directions * hidden_size)

        return output, hiddens
