# coding:utf-8
#
# unike/module/model/CompGCN.py
#
# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 19, 2023
# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on May 26, 2024
#
# 该脚本定义了 CompGCN 类.

"""
CompGCN - 这是一种在图卷积网络中整合多关系信息的新框架，它利用知识图谱嵌入技术中的各种组合操作，将实体和关系共同嵌入到图中。
"""

import dgl
import torch
import typing
from torch import nn
import dgl.function as fn
from .Model import Model
import torch.nn.functional as F
from typing_extensions import override

class CompGCN(Model):

    """
    ``CompGCN`` :cite:`CompGCN` 发表于 ``2020`` 年，这是一种在图卷积网络中整合多关系信息的新框架，它利用知识图谱嵌入技术中的各种组合操作，将实体和关系共同嵌入到图中。

    正三元组的评分函数的值越大越好，如果想获得更详细的信息请访问 :ref:`CompGCN <compgcn>`。

    例子::

        from unike.module.model import CompGCN
        from unike.module.loss import CompGCNLoss
        from unike.module.strategy import CompGCNSampling
        from unike.config import Trainer, GraphTester
        
        # define the model
        compgcn = CompGCN(
        	ent_tol = dataloader.get_ent_tol(),
        	rel_tol = dataloader.get_rel_tol(),
        	dim = 100
        )
        
        # define the loss function
        model = CompGCNSampling(
        	model = compgcn,
        	loss = CompGCNLoss(model = compgcn),
        	ent_tol = dataloader.get_ent_tol()
        )
        
        # test the model
        tester = GraphTester(model = compgcn, data_loader = dataloader, use_gpu = True, device = 'cuda:0', prediction = "tail")
        
        # train the model
        trainer = Trainer(model = model, data_loader = dataloader.train_dataloader(),
        	epochs = 2000, lr = 0.0001, use_gpu = True, device = 'cuda:0',
        	tester = tester, test = True, valid_interval = 50, log_interval = 50,
        	save_interval = 50, save_path = '../../checkpoint/compgcn.pth'
        )
        trainer.run()
    """

    def __init__(
        self,
        ent_tol: int,
        rel_tol: int,
        dim: int,
        opn: str = 'mult',
        fet_drop: float = 0.2,
        hid_drop: float = 0.3,
        margin: float = 40.0,
        decoder_model: str = 'ConvE'):

        """创建 RGCN 对象。
        
        :param ent_tol: 实体的个数
        :type ent_tol: int
        :param rel_tol: 关系的个数
        :type rel_tol: int
        :param dim: 实体和关系嵌入向量的维度
        :type dim: int
        :param opn: 组成运算符：'mult'、'sub'、'corr'
        :type opn: str
        :param fet_drop: 用于 'ConvE' 解码器，用于卷积特征的 dropout
        :type fet_drop: float
        :param hid_drop: 用于 'ConvE' 解码器，用于隐藏层的 dropout
        :type hid_drop: float
        :param margin: 用于 'TransE' 解码器，gamma。
        :type margin: float
        :param decoder_model: 用什么得分函数作为解码器: 'ConvE'、'DistMult'、'TransE'
        :type decoder_model: str
		"""

        super(CompGCN, self).__init__(ent_tol, rel_tol)

        #: 实体和关系嵌入向量的维度
        self.dim: int = dim
        #: 组成运算符：'mult'、'sub'、'corr'
        self.opn: str = opn
        #: 用什么得分函数作为解码器: 'ConvE'、'DistMult'
        self.decoder_model: str = decoder_model

        #------------------------------CompGCN--------------------------------------------------------------------
        #: 根据实体个数，创建的实体嵌入
        self.ent_emb: torch.nn.parameter.Parameter = nn.Parameter(torch.Tensor(self.ent_tol, self.dim))
        #: 根据关系个数，创建的关系嵌入
        self.rel_emb: torch.nn.parameter.Parameter = nn.Parameter(torch.Tensor(self.rel_tol, self.dim))

        nn.init.xavier_normal_(self.ent_emb, gain=nn.init.calculate_gain('relu'))
        nn.init.xavier_normal_(self.rel_emb, gain=nn.init.calculate_gain('relu'))

        #: CompGCNCov
        self.GraphCov: CompGCNCov = CompGCNCov(self.dim, self.dim * 2, torch.tanh, bias = 'False', drop_rate = 0.1, opn = self.opn)
        #: 用于 :py:attr:`GraphCov` 输出结果
        self.drop: torch.nn.Dropout = nn.Dropout(0.3)
        #: 最后计算得分时的偏置
        self.bias: torch.nn.parameter.Parameter = nn.Parameter(torch.zeros(self.ent_tol))
        #-----------------------------ConvE-----------------------------------------------------------------------
        #: 用于 'ConvE' 解码器，头实体嵌入向量和关系嵌入向量的 BatchNorm
        self.bn0: torch.nn.BatchNorm2d = torch.nn.BatchNorm2d(1)
        #: 用于 'ConvE' 解码器，卷积层
        self.conv1: torch.nn.Conv2d = torch.nn.Conv2d(1, 200, (7, 7), 1, 0, bias=False)
        #: 用于 'ConvE' 解码器，卷积特征的 BatchNorm
        self.bn1: torch.nn.Conv2d = torch.nn.BatchNorm2d(200)
        #: 用于 'ConvE' 解码器，卷积特征的 Dropout
        self.fet_drop: torch.nn.Dropout = torch.nn.Dropout2d(fet_drop)
        flat_sz_h = 4 * self.dim // 20 - 7 + 1
        flat_sz_w = 20 - 7 + 1
        flat_sz = flat_sz_h * flat_sz_w * 200
        #: 用于 'ConvE' 解码器，隐藏层层
        self.fc: torch.nn.Linear = torch.nn.Linear(flat_sz, self.dim*2)
        #: 用于 'ConvE' 解码器，隐藏层的 Dropout
        self.hid_drop: torch.nn.Dropout = torch.nn.Dropout(hid_drop)
        #: 用于 'ConvE' 解码器，隐藏层的 BatchNorm
        self.bn2: torch.nn.BatchNorm1d = torch.nn.BatchNorm1d(self.dim*2)
        #-----------------------------TransE-----------------------------------------------------------------------
        #: 用于 TransE 得分函数
        self.margin: torch.nn.parameter.Parameter = nn.Parameter(torch.Tensor([margin]))
        self.margin.requires_grad = False

    @override
    def forward(
        self,
        graph: dgl.DGLGraph,
        relation: torch.Tensor,
        norm: torch.Tensor,
        triples: torch.Tensor) -> torch.Tensor:

        """
		定义每次调用时执行的计算。
		:py:class:`torch.nn.Module` 子类必须重写 :py:meth:`torch.nn.Module.forward`。
        
        :param graph: 子图
        :type graph: dgl.DGLGraph
        :param relation: 子图的关系
        :type relation: torch.Tensor
        :param norm: 关系的归一化系数
        :type norm: torch.Tensor
        :param triples: 三元组
        :type triples: torch.Tensor
        :returns: 三元组的得分
        :rtype: torch.Tensor
		"""

        head, rela = triples[:,0], triples[:, 1]
        x, r = self.ent_emb, self.rel_emb
        x, r = self.GraphCov(graph, x, r, relation, norm)
        x = self.drop(x)
        head_emb = torch.index_select(x, 0, head)
        rela_emb = torch.index_select(r, 0, rela)

        if self.decoder_model.lower() == 'conve':
           score = self.conve(head_emb, rela_emb, x)
        elif self.decoder_model.lower() == 'distmult':
            score = self.distmult(head_emb, rela_emb, x)
        elif self.decoder_model.lower() == 'transe':
            score = self.transe(head_emb, rela_emb, x)
        else:
            raise ValueError("please choose decoder (TransE/DistMult/ConvE)")

        return score

    def conve(
        self,
        sub_emb: torch.Tensor,
        rel_emb: torch.Tensor,
        all_ent: torch.Tensor) -> torch.Tensor:

        """计算 ConvE 作为解码器时三元组的得分。
        
        :param sub_emb: 头实体的嵌入向量
        :type sub_emb: torch.Tensor
        :param rel_emb: 关系的嵌入向量
        :type rel_emb: torch.Tensor
        :param all_ent: 全部实体的嵌入向量
        :type all_ent: torch.Tensor
        :returns: 三元组的得分
        :rtype: torch.Tensor"""

        stack_input = self.concat(sub_emb, rel_emb)
        x = self.bn0(stack_input)
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.fet_drop(x)
        x = x.view(x.shape[0], -1)
        x = self.fc(x)
        x = self.hid_drop(x)
        x = self.bn2(x)
        x = F.relu(x)
        x = torch.mm(x, all_ent.transpose(1, 0))
        x += self.bias.expand_as(x)
        score = torch.sigmoid(x)
        return score

    def concat(
        self,
        ent_embed: torch.Tensor,
        rel_embed: torch.Tensor) -> torch.Tensor:

        """ConvE 作为解码器时，用于拼接头实体嵌入向量和关系嵌入向量。
        
        :param ent_embed: 头实体的嵌入向量
        :type ent_embed: torch.Tensor
        :param rel_embed: 关系的嵌入向量
        :type rel_embed: torch.Tensor
        :returns: ConvE 解码器的输入特征
        :rtype: torch.Tensor"""

        ent_embed = ent_embed.view(-1, 1, self.dim*2)
        rel_embed = rel_embed.view(-1, 1, self.dim*2)
        stack_input = torch.cat([ent_embed, rel_embed], 1)
        stack_input = stack_input.reshape(-1, 1, 4*self.dim//20, 20)
        return stack_input

    def distmult(
        self,
        head_emb: torch.Tensor,
        rela_emb: torch.Tensor,
        all_ent: torch.Tensor) -> torch.Tensor:

        """计算 DistMult 作为解码器时三元组的得分。
        
        :param sub_emb: 头实体的嵌入向量
        :type sub_emb: torch.Tensor
        :param rel_emb: 关系的嵌入向量
        :type rel_emb: torch.Tensor
        :param all_ent: 全部实体的嵌入向量
        :type all_ent: torch.Tensor
        :returns: 三元组的得分
        :rtype: torch.Tensor"""

        obj_emb = head_emb * rela_emb
        x = torch.mm(obj_emb, all_ent.transpose(1, 0))
        x += self.bias.expand_as(x)
        score = torch.sigmoid(x)
        return score

    def transe(
        self,
        head_emb: torch.Tensor,
        rela_emb: torch.Tensor,
        all_ent: torch.Tensor) -> torch.Tensor:

        """计算 TransE 作为解码器时三元组的得分。
        
        :param sub_emb: 头实体的嵌入向量
        :type sub_emb: torch.Tensor
        :param rel_emb: 关系的嵌入向量
        :type rel_emb: torch.Tensor
        :param all_ent: 全部实体的嵌入向量
        :type all_ent: torch.Tensor
        :returns: 三元组的得分
        :rtype: torch.Tensor"""

        obj_emb = head_emb + rela_emb
        x = self.margin - torch.norm(obj_emb.unsqueeze(1) - all_ent, p=1, dim=2)
        score = torch.sigmoid(x)
        return score

    @override
    def predict(
        self,
        data: dict[str, typing.Union[dgl.DGLGraph, torch.Tensor]],
        mode: str) -> torch.Tensor:

        """CompGCN 的推理方法。
        
        :param data: 数据。
        :type data: dict[str, typing.Union[dgl.DGLGraph, torch.Tensor]]
        :param mode: 在 CompGCN 时，无用，只为了保证推理函数形式一致
        :type mode: str
        :returns: 三元组的得分
        :rtype: torch.Tensor
		"""
        
        triples    = data['positive_sample']
        graph      = data['graph']
        relation   = data['rela']
        norm       = data['norm'] 

        head, rela = triples[:,0], triples[:, 1]
        x, r = self.ent_emb, self.rel_emb
        x, r = self.GraphCov(graph, x, r, relation, norm)
        x = self.drop(x)
        head_emb = torch.index_select(x, 0, head)
        rela_emb = torch.index_select(r, 0, rela)

        if self.decoder_model.lower() == 'conve':
           score = self.conve(head_emb, rela_emb, x)
        elif self.decoder_model.lower() == 'distmult':
            score = self.distmult(head_emb, rela_emb, x)
        elif self.decoder_model.lower() == 'transe':
            score = self.transe(head_emb, rela_emb, x)
        else:
            raise ValueError("please choose decoder (TransE/DistMult/ConvE)")

        return score

class CompGCNCov(nn.Module):

    """``CompGCN`` :cite:`CompGCN` 图神经网络模块。"""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        act: typing.Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
        bias: bool = True,
        drop_rate: float = 0.,
        opn: str = 'corr'):

        """创建 CompGCN 对象。
        
        :param in_channels: 输入的特征维度
        :type in_channels: int
        :param out_channels: 输出的特征维度
        :type out_channels: int
        :param act: 激活函数
        :type act: typing.Callable[[torch.Tensor], torch.Tensor]
        :param bias: 是否有偏置
        :type bias: bool
        :param drop_rate: Dropout rate
        :type drop_rate: float
        :param opn: 组成运算符：'mult'、'sub'、'corr'
        :type opn: str
		"""

        super(CompGCNCov, self).__init__()

        #: 输入的特征维度
        self.in_channels: int = in_channels
        #: 输出的特征维度
        self.out_channels: int = out_channels
        self.rel_wt = None
        #: 关系嵌入向量
        self.rel: torch.nn.parameter.Parameter = None
        #: 组成运算符：'mult'、'sub'、'corr'
        self.opn: str = opn
        #：图神经网络的权重矩阵，用于原始关系
        self.in_w: torch.nn.parameter.Parameter = self.get_param([in_channels, out_channels])
        #：图神经网络的权重矩阵，用于相反关系
        self.out_w: torch.nn.parameter.Parameter = self.get_param([in_channels, out_channels])
        #: 用于原始关系和相反关系转换后输出结果的 Dropout
        self.drop: torch.nn.Dropout = nn.Dropout(drop_rate)
        #: 自循环关系嵌入向量的转换矩阵
        self.loop_rel: torch.nn.parameter.Parameter = self.get_param([1, in_channels])
        #：图神经网络的权重矩阵，用于自循环关系
        self.loop_w: torch.nn.parameter.Parameter = self.get_param([in_channels, out_channels])
        #: 偏置
        self.bias: torch.nn.Parameter = nn.Parameter(torch.zeros(out_channels)) if bias else None
        #: BatchNorm
        self.bn: torch.nn.BatchNorm1d = torch.nn.BatchNorm1d(out_channels)
        #: 激活函数
        self.act: typing.Callable[[torch.Tensor], torch.Tensor] = act
        #: 关系嵌入向量的转换矩阵
        self.w_rel: torch.nn.parameter.Parameter = self.get_param([in_channels, out_channels])

    def get_param(
        self,
        shape: list[int]) -> torch.nn.parameter.Parameter:

        """获得权重矩阵。
        
        :param shape: 权重矩阵的 shape
        :type shape: list[int]
        :returns: 权重矩阵
        :rtype: torch.nn.parameter.Parameter
        """

        param = nn.Parameter(torch.Tensor(*shape))
        nn.init.xavier_normal_(param, gain=nn.init.calculate_gain('relu'))
        return param

    def forward(
        self,
        graph: dgl.DGLGraph,
        ent_emb: torch.nn.parameter.Parameter,
        rel_emb: torch.nn.parameter.Parameter,
        edge_type: torch.Tensor,
        edge_norm: torch.Tensor) -> tuple[torch.nn.parameter.Parameter, torch.nn.parameter.Parameter]:
        
        """
        定义每次调用时执行的计算。
        :py:class:`torch.nn.Module` 子类必须重写 :py:meth:`torch.nn.Module.forward`。
        
        :param graph: 子图
        :type graph: dgl.DGLGraph
        :param ent_emb: 实体嵌入向量
        :type ent_emb: torch.nn.parameter.Parameter
        :param rel_emb: 关系嵌入向量
        :type rel_emb: torch.nn.parameter.Parameter
        :param edge_type: 关系 ID
        :type edge_type: torch.Tensor
        :param norm: 关系的归一化系数
        :type norm: torch.Tensor
        :returns: 更新后的实体嵌入和关系嵌入
        :rtype: tuple[torch.nn.parameter.Parameter, torch.nn.parameter.Parameter]
        """

        graph = graph.local_var()
        graph.ndata['h'] = ent_emb
        graph.edata['type'] = edge_type
        graph.edata['norm'] = edge_norm
        if self.rel_wt is None:
            self.rel = rel_emb
        else:
            self.rel = torch.mm(self.rel_wt, rel_emb)
        graph.update_all(self.message_func, fn.sum(msg='msg', out='h'), self.reduce_func)
        ent_emb = graph.ndata.pop('h') + torch.mm(self.comp(ent_emb, self.loop_rel), self.loop_w) / 3
        if self.bias is not None:
            ent_emb = ent_emb + self.bias
        ent_emb = self.bn(ent_emb)

        return self.act(ent_emb), torch.matmul(self.rel, self.w_rel)

    def message_func(self, edges: dgl.udf.EdgeBatch):

        """
        消息函数。
        """

        edge_type = edges.data['type']
        edge_num = edge_type.shape[0]
        edge_data = self.comp(edges.src['h'], self.rel[edge_type])
        msg = torch.cat([torch.matmul(edge_data[:edge_num // 2, :], self.in_w),
                         torch.matmul(edge_data[edge_num // 2:, :], self.out_w)])
        msg = msg * edges.data['norm'].reshape(-1, 1)
        return {'msg': msg}

    def comp(
        self,
        h: torch.Tensor,
        r: torch.Tensor) -> torch.Tensor:

        """组成运算：'mult'、'sub'、'corr'
        
        :param h: 头实体嵌入向量
        :type h: torch.Tensor
        :param r: 关系嵌入向量
        :type r: torch.Tensor
        :returns: 组合后的边数据
        :rtype: torch.Tensor
        """

        def com_mult(a, b):

            """复数乘法"""

            r1, i1 = a.real, a.imag
            r2, i2 = b.real, b.imag
            real = r1 * r2 - i1 * i2
            imag = r1 * i2 + i1 * r2
            return torch.complex(real, imag)

        def conj(a):

            """共轭复数"""

            a.imag = -a.imag
            return a

        def ccorr(a, b):

            """corr 运算"""

            return torch.fft.irfft(com_mult(conj(torch.fft.rfft(a)), torch.fft.rfft(b)), a.shape[-1])

        if self.opn == 'mult':
            return h * r
        elif self.opn == 'sub':
            return h - r
        elif self.opn == 'corr':
            return ccorr(h, r.expand_as(h))
        else:
            raise KeyError(f'composition operator {self.opn} not recognized.')

    def reduce_func(self, nodes: dgl.udf.NodeBatch):

        """聚合函数"""
        
        return {'h': self.drop(nodes.data['h']) / 3}

def get_compgcn_hpo_config() -> dict[str, dict[str, typing.Any]]:
    
    """返回 :py:class:`CompGCN` 的默认超参数优化配置。
    
    默认配置为::
	
        parameters_dict = {
            'model': {
                'value': 'CompGCN'
            },
            'dim': {
                'values': [100, 150, 200]
            },
            'opn': {
                'value': 'mult'
            },
            'fet_drop': {
                'value': 0.2
            },
            'hid_drop': {
                'value': 0.3
            },
            'margin': {
                'value': 40.0
            },
            'decoder_model': {
                'value': 'ConvE'
            }
        }
        
    :returns: :py:class:`CompGCN` 的默认超参数优化配置
    :rtype: dict[str, dict[str, typing.Any]]
    """
    
    parameters_dict = {
        'model': {
            'value': 'CompGCN'
        },
        'dim': {
            'values': [100, 150, 200]
        },
        'opn': {
            'value': 'mult'
        },
        'fet_drop': {
            'value': 0.2
        },
        'hid_drop': {
            'value': 0.3
        },
        'margin': {
            'value': 40.0
        },
        'decoder_model': {
            'value': 'ConvE'
        }
    }
    
    return parameters_dict