#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/8/2 21:33
# @Author  : Raymound luo
# @Mail    : luolinhao1998@gmail.com
# @File    : CGNN.py
# @Software: PyCharm
# @Describe:

import torch
import torch.nn as nn
import dgl.function as fn
import json
import torch.nn.functional as F


class ContextGNN(nn.Module):
    def __init__(self, g, model_config):
        super(ContextGNN, self).__init__()
        self.g = g
        self.primary_type = model_config['primary_type']
        self.auxiliary_embedding = model_config['auxiliary_embedding']
        self.K_length = model_config['K_length']
        self.primary_emb = nn.Embedding(g.number_of_nodes(
            self.primary_type), model_config['embedding_dim'])
        self.enable_gru = model_config['gru']
        self.enable_add_init = model_config['add_init']
        self.ntypes = g.ntypes
        self.etypes = g.etypes
        if self.auxiliary_embedding == "emb":
            self.auxiliary_emb = nn.ModuleDict(
                {ntype: nn.Embedding(g.number_of_nodes(ntype), model_config['embedding_dim']) for ntype in g.ntypes if
                 ntype != self.primary_type})
        else:
            self.auxiliary_tans_fuc = nn.ModuleDict({
                etype: EmbTransformer(model_config['in_dim'], model_config['out_dim'], self.auxiliary_embedding) for
                etype in
                g.etypes
            })

        self.multihead_cgnn = nn.ModuleList(
            [MultiHeadCGNN(model_config['in_dim'], model_config['out_dim'],
                           model_config['num_heads'],
                           merge=model_config['merge'], g_agg_type=model_config['g_agg_type'],
                           drop_out=model_config['drop_out'],
                           cgnn_non_linear=model_config['cgnn_non_linear'],
                           multi_attn_linear=model_config['multi_attn_linear'],
                           ntypes=g.ntypes,
                           etypes=g.etypes,
                           graph_attention=model_config['graph_attention'],
                           kq_linear_out_dim=model_config['kq_linear_out_dim'],
                           path_attention=model_config['path_attention'],
                           c_linear_out_dim=model_config['c_linear_out_dim'],
                           enable_bilinear=model_config['enable_bilinear']) for hop in
             range(self.K_length + 1)  # self.K_length + 1 (self.K_length + 1) // 2 + 1
             ])
        self.loss_weight = nn.Parameter(torch.zeros(self.K_length + 1))

        if self.enable_gru:
            self.gru_gate = nn.GRUCell(
                model_config['out_dim'], model_config['out_dim'])
        self.reset_parameter()

    def reset_parameter(self):
        if self.auxiliary_embedding == "emb":
            [nn.init.normal_(emb.weight)
             for emb in self.auxiliary_emb.values()]
        else:
            [nn.init.normal_(func.linear.weight)
             for func in self.auxiliary_tans_fuc.values()]

    def embedding_transformer(self, primary_feature):
        '''
        Get HIN feature dict for each node type
        :param primary_feature: Primary Graph feature
        :return: h_dict, {ntype, feature}
        '''
        if self.auxiliary_embedding == "emb":
            h_dict = {ntype: emb.weight for ntype, emb in self.auxiliary_emb.items()}
            h_dict[self.primary_type] = primary_feature
        else:
            h_dict = {}
            non_init_graph = self.g.ntypes[:]  # Copy list
            h_dict[self.primary_type] = primary_feature
            non_init_graph.remove(self.primary_type)
            while non_init_graph:
                for srctype, etype, dsttype in self.g.canonical_etypes:
                    if srctype not in non_init_graph and dsttype in non_init_graph:
                        h_dict[dsttype] = self.auxiliary_tans_fuc[etype](
                            self.g, h_dict[srctype], srctype, etype, dsttype
                        )
                        non_init_graph.remove(dsttype)
        return h_dict

    def dump_cgnn_attention_matrix(self, path):
        '''
        Dump multi head attention matrix to json
        :param path: path to json file
        :return:
        '''
        with open(path, 'w') as f:
            multi_head_attention_matrix = {}
            for k, layer in enumerate(self.multihead_cgnn):
                hop_multi_head_matrix = layer.dump_multi_head_attention_matrix()
                multi_head_attention_matrix["length_{}".format(
                    k)] = hop_multi_head_matrix

            json.dump(multi_head_attention_matrix, f)

    def forward(self, k):
        '''

        :param k: k-hop
        :return: K_length feature emb h_feature of primary type
        '''
        h_dict = self.embedding_transformer(self.primary_emb.weight)  # Generate emb for auxiliary graph
        # KCI
        # k = k // 2 + 1
        for hop in range(k):
            new_h_dict = self.multihead_cgnn[hop](self.g, h_dict)  # multi head CGNN

            if self.enable_gru:  # Enable Gru
                new_primary_feature = self.gru_gate(
                    new_h_dict[self.primary_type], h_dict[self.primary_type]
                )  # GRU gate for primary graph
                new_h_dict[self.primary_type] = new_primary_feature
            if self.enable_add_init:
                new_h_dict[self.primary_type] += self.primary_emb.weight
            h_dict = new_h_dict
        return h_dict[self.primary_type]

    def _context_score(self, src, dst, p_emb, context_emb):
        logits = torch.sum((p_emb[src] * context_emb[src])
                           * (p_emb[dst] * context_emb[dst]), dim=1)
        return logits

    def get_loss(self, k_hop, pos_src, pos_dst, neg_src, neg_dst, p_emb, p_context_emb):
        # k_length = k_hop // 2
        k_length = k_hop - 1
        pos_score = self._context_score(pos_src, pos_dst, p_emb, p_context_emb)
        neg_score = self._context_score(neg_src, neg_dst, p_emb, p_context_emb)
        # pos_score = torch.clamp(pos_score, max=10, min=-10)
        pos_loss = -F.logsigmoid(pos_score).view(-1).mean()
        # neg_score = torch.clamp(neg_score, max=10, min=-10)
        neg_loss = -F.logsigmoid(-neg_score).view(-1).mean()
        # Multi loss
        weight = torch.exp(-self.loss_weight[k_length])
        return (pos_loss + neg_loss) * weight + self.loss_weight[
            k_length]  # torch.mean(torch.cat((pos_score, neg_score)))


class EmbTransformer(nn.Module):
    def __init__(self, in_feat, out_feat, apply_linear='linear'):
        super(EmbTransformer, self).__init__()
        self.linear = nn.Linear(in_feat, out_feat)
        self.apply_linear = apply_linear

    def forward(self, g, src_h, srctype, etype, dsttype):
        g.nodes[srctype].data['h'] = src_h
        g[etype].update_all(fn.copy_src('h', 'm'), fn.sum(msg='m', out='h'))
        if self.apply_linear == 'non_linear':
            return torch.relu(self.linear(g.nodes[dsttype].data['h']))
        else:
            return self.linear(g.nodes[dsttype].data['h'])


class MultiHeadCGNN(nn.Module):
    def __init__(self, in_dim, out_dim, num_heads, merge='linear', g_agg_type="mean", drop_out=0.,
                 cgnn_non_linear=False,
                 multi_attn_linear=False, ntypes=None, etypes=None, graph_attention=True, kq_linear_out_dim=None,
                 path_attention=True,
                 c_linear_out_dim=None, enable_bilinear=False):
        super(MultiHeadCGNN, self).__init__()
        self.heads = nn.ModuleList()
        cgnn_out_dim = out_dim // num_heads if merge == 'stack' or merge == 'linear' else out_dim
        for i in range(num_heads):
            self.heads.append(
                CGNNLayer(in_dim, cgnn_out_dim, g_agg_type, drop_out, cgnn_non_linear,
                          multi_attn_linear, ntypes, etypes, graph_attention, kq_linear_out_dim, path_attention,
                          c_linear_out_dim, enable_bilinear))
        self.merge = merge
        # W_2
        self.cgnn_non_linear = cgnn_non_linear
        if merge == 'linear':
            self.fc = nn.Linear(out_dim, out_dim)
            self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_normal_(self.fc.weight)

    def dump_multi_head_attention_matrix(self):
        '''
        Reutrn attention matrix of each head
        :return: multi_head_matrix_dict {h_head, attention matrix of head h}
        '''
        multi_head_matrix_dict = {}
        for h, head in enumerate(self.heads):
            head_matrix = head.dump_attention_matrix()
            multi_head_matrix_dict['head_{}'.format(h)] = head_matrix
        return multi_head_matrix_dict

    def forward(self, g, h):
        head_outs = [attn_head(g, h) for attn_head in self.heads]
        h_dict = {}
        if self.merge == 'linear':
            # concat on the output feature dimension (dim=1)
            for ntype in g.ntypes:
                n_feature = torch.cat([feature[ntype]
                                       for feature in head_outs], dim=1)
                h_dict[ntype] = self.fc(n_feature)
        elif self.merge == 'mean':
            # merge using average
            for ntype in g.ntypes:
                n_feature = torch.stack([feature[ntype]
                                         for feature in head_outs], dim=1)
                h_dict[ntype] = torch.mean(n_feature, dim=1)
        elif self.merge == 'stack':
            for ntype in g.ntypes:
                n_feature = torch.cat([feature[ntype]
                                       for feature in head_outs], dim=1)
                h_dict[ntype] = n_feature
        return h_dict


class CGNNLayer(nn.Module):
    def __init__(self, in_dim, out_dim, g_agg_type="mean", drop_out=0., cgnn_non_linear=False,
                 multi_attn_linear=False,
                 ntypes=None, etypes=None, graph_attention=True, kq_linear_out_dim=None, path_attention=True,
                 c_linear_out_dim=None, enable_bilinear=False):
        super(CGNNLayer, self).__init__()
        self.graph_attention_matrix = {}
        self.g_agg_type = g_agg_type  # Graph aggregation type
        self.multi_attn_linear = multi_attn_linear
        self.cgnn_non_linear = cgnn_non_linear
        self.graph_attention = graph_attention
        self.path_attention = path_attention
        self.enable_bilinear = enable_bilinear
        if self.multi_attn_linear:
            self.k_linear = nn.ModuleDict()
            self.q_linear = nn.ModuleDict()
            self.c_linear = nn.ModuleDict()
            for ntype in ntypes:
                self.k_linear[ntype] = nn.Linear(
                    in_dim, kq_linear_out_dim, bias=False)
                self.q_linear[ntype] = nn.Linear(
                    in_dim, kq_linear_out_dim, bias=False)
                self.c_linear[ntype] = nn.Linear(
                    in_dim, c_linear_out_dim, bias=False)
        else:
            # K-linear
            self.k_linear = nn.Linear(in_dim, kq_linear_out_dim, bias=False)
            # Q-linear
            self.q_linear = nn.Linear(in_dim, kq_linear_out_dim, bias=False)
            self.c_linear = nn.Linear(in_dim, c_linear_out_dim, bias=False)
        # W_1
        if self.cgnn_non_linear:
            self.fc = nn.Sequential(
                nn.Linear(in_dim, out_dim),
                nn.ReLU()
            )
        else:
            self.fc = nn.Linear(in_dim, out_dim)
        self.node_dropout = nn.Dropout2d(drop_out)
        self.attn_fc = nn.Linear(2 * c_linear_out_dim, 1, bias=False)
        if enable_bilinear:
            self.bilinear = nn.ModuleDict()
            for etype in etypes:
                self.bilinear[etype] = nn.Bilinear(kq_linear_out_dim, kq_linear_out_dim, 1)
        self.reset_parameters()

    def reset_parameters(self):
        """Reinitialize learnable parameters."""
        if self.multi_attn_linear:
            for linear in self.k_linear.values():
                nn.init.xavier_normal_(linear.weight)
            for linear in self.q_linear.values():
                nn.init.xavier_normal_(linear.weight)
        else:
            nn.init.xavier_normal_(self.k_linear.weight)
            nn.init.xavier_normal_(self.q_linear.weight)
        if self.cgnn_non_linear:
            nn.init.xavier_normal_(self.fc[0].weight)
        else:
            nn.init.xavier_normal_(self.fc.weight)

    def graph_aggregation(self, h):
        '''
        How to aggregate graph information
        :param h: Node embedding (N, D)
        :return: Graph representation (D)
        '''
        if self.g_agg_type == "mean":
            return torch.mean(h, dim=0)
        elif self.g_agg_type == "sum":
            return torch.sum(h, dim=0)
        return

    def dump_attention_matrix(self):
        '''
        Return the attention matrix of each hop
        :return: graph_attention_matrix attention_matrix {etype, attention score}
        '''
        return self.graph_attention_matrix

    def context_edge_attention(self, edges):
        # s_ij = W(h_i||h_j)
        z = torch.cat((edges.src['C_h'], edges.dst['C_h']), dim=1)
        c_A = self.attn_fc(z)
        # c_A = torch.sum(edges.src['Q_h'] * edges.dst['K_h'], dim=1) / (edges.dst['K_h'].shape[-1] ** (1 / 2))
        return {'c_A': c_A.squeeze(-1)}

    def context_attention(self, g):
        '''
        Context edges attention score
        :param g:
        :return: g with k,q feature for each node
        '''
        for ntype in g.ntypes:
            if self.multi_attn_linear:
                g.nodes[ntype].data['C_h'] = self.c_linear[ntype](
                    g.nodes[ntype].data['h'])
            else:
                g.nodes[ntype].data['C_h'] = self.c_linear(
                    g.nodes[ntype].data['h'])
        for etype in g.etypes:
            g.apply_edges(self.context_edge_attention, etype=etype)
        return g

    def msg_func(self, edges):
        return {'h': edges.src['h'], 'C_h': edges.src['C_h'], 'g_A': edges.data['g_A'], 'c_A': edges.data['c_A']}

    def reduce_func(self, nodes):
        # S_j
        # sim_matrix = torch.tanh(torch.bmm(nodes.mailbox['C_h'], nodes.mailbox['C_h'].transpose(1, 2)) / (
        #         nodes.mailbox['C_h'].shape[-1] ** (1 / 2)))
        # sim_J = torch.mean(sim_matrix, dim=2)
        # \sigma(s_ij-S_j)
        e = F.leaky_relu(nodes.mailbox['c_A'])  # nodes.mailbox['c_A'] - sim_J
        alpha = F.softmax(e, dim=1).unsqueeze(-1)
        h = nodes.mailbox['g_A'][:,
            0].unsqueeze(-1) * torch.sum(alpha * nodes.mailbox['h'], dim=1)
        return {'h': h}

    def forward(self, g, h):
        '''
        CGNN
        :param g: graph
        :param h: feature dict {ntype, feature}
        :return: h: feature dict {ntype, feature}
        '''
        # The input is a dictionary of node features for each type
        funcs = {}
        graph_feature = {}
        for ntype in g.ntypes:
            h_feat = h[ntype].unsqueeze(0)  # (1, N , D)
            feat = self.node_dropout(h_feat)
            feat = feat.squeeze(0)  # (N, D)
            graph_feature[ntype] = self.graph_aggregation(
                feat)  # Graph Aggregation
            g.nodes[ntype].data['h'] = feat

        # Context Edges attention:
        if self.path_attention:
            g = self.context_attention(g)

        # Graph Attention
        if self.graph_attention:
            l_hop_graph_attention = {}
            for q_dsttype in g.ntypes:
                context_feature = []
                if self.multi_attn_linear:
                    q_h = self.q_linear[q_dsttype](graph_feature[q_dsttype])
                else:
                    q_h = self.q_linear(graph_feature[q_dsttype])  # k linear

                # Compute attention score for each other graph
                for k_srctype, k_etype, k_dsttype in g.canonical_etypes:  # Score only exists when two graphs are connected
                    if q_dsttype != k_dsttype:
                        continue
                    if self.multi_attn_linear:
                        k_h = self.k_linear[k_srctype](
                            graph_feature[k_srctype])
                    else:
                        k_h = self.k_linear(
                            graph_feature[k_srctype])  # q linear
                    # A = g.number_of_edges(q_etype)  # V
                    if self.enable_bilinear:
                        z = self.bilinear[k_etype](k_h, q_h)
                    else:
                        z = torch.sum(k_h * q_h) / \
                            (k_h.shape[-1] ** (1 / 2))  # ATThead
                    context_feature.append((k_etype, z))
                c_feature = torch.stack(
                    [score for _, score in context_feature])
                c_attention_matrix = torch.softmax(c_feature, dim=0)
                for i, (etype, score) in enumerate(context_feature):
                    att_score = c_attention_matrix[i]
                    g.edges[etype].data['g_A'] = att_score.expand(
                        g.number_of_edges(etype))  # Bind attention score to edge feature
                    l_hop_graph_attention[etype] = float(
                        c_attention_matrix[i].detach().cpu().numpy())
            self.graph_attention_matrix = l_hop_graph_attention

        # h: source node feature, g_A: Graph attention score between two graphs, c_A: Attentnion score between two nodes, m: result
        for etype in g.etypes:
            if self.path_attention:
                funcs[etype] = (self.msg_func, self.reduce_func)
            elif self.graph_attention:
                funcs[etype] = (fn.u_mul_e('h', 'g_A', 'm'), fn.sum('m', 'h'))
            else:
                funcs[etype] = (fn.copy_src('h', 'm'), fn.sum('m', 'h'))
        g.multi_update_all(funcs, 'sum')
        return {ntype: self.fc(g.nodes[ntype].data['h']) for ntype in g.ntypes}
