"""GNN 模型：RGNN，移植自 TorchEEG，主要用于 EEG 情绪识别中的正则化图神经网络建模。"""

from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from zyplib.nn import BaseModule


class LazyLoader:
    def __init__(self, lib_name):
        self._lib_name = lib_name
        self._module = None

    def __getattr__(self, name):
        if self._module is None:
            try:
                self._module = __import__(self._lib_name)
            except ImportError:
                raise ImportError(
                    f"To use this functionality, you need to install `{self._lib_name}`. "
                    f"Please refer to https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html"
                )
        return getattr(self._module, name)


pyg = LazyLoader('torch_geometric')


def _scatter_add_1d(src: torch.Tensor, index: torch.Tensor, dim_size: int) -> torch.Tensor:
    out = torch.zeros(dim_size, dtype=src.dtype, device=src.device)
    out.index_add_(0, index, src)
    return out


def _batch_edge_index(edge_index: torch.Tensor, num_nodes: int, batch_size: int) -> torch.Tensor:
    offsets = torch.arange(batch_size, device=edge_index.device).repeat_interleave(edge_index.shape[1]) * num_nodes
    return edge_index.repeat(1, batch_size) + offsets.unsqueeze(0)


def _unpack_graph_input(
    data_or_x,
    edge_index: Optional[torch.Tensor] = None,
    batch: Optional[torch.Tensor] = None,
):
    if edge_index is None and hasattr(data_or_x, 'x') and hasattr(data_or_x, 'edge_index'):
        data = data_or_x
        x = data.x
        edge_index = data.edge_index
        batch = getattr(data, 'batch', None)
        num_graphs = getattr(data, 'num_graphs', None)
        return x, edge_index, batch, num_graphs

    if edge_index is None:
        raise ValueError('edge_index is required when the input is not a PyG Batch/Data object.')

    x = data_or_x
    if x.ndim == 3:
        batch_size, num_nodes, num_features = x.shape
        if edge_index.numel() > 0 and int(edge_index.max().item()) < num_nodes:
            edge_index = _batch_edge_index(edge_index, num_nodes, batch_size)
        x = x.reshape(batch_size * num_nodes, num_features)
        if batch is None:
            batch = torch.arange(batch_size, device=x.device).repeat_interleave(num_nodes)
        return x, edge_index, batch, batch_size

    if x.ndim == 2:
        if batch is None:
            raise ValueError('batch is required when x is already flattened to shape [num_nodes_total, num_features].')
        num_graphs = int(batch.max().item()) + 1 if batch.numel() > 0 else 0
        return x, edge_index, batch, num_graphs

    raise ValueError('Expected x to have shape [batch, num_nodes, num_features] or [num_nodes_total, num_features].')


def maybe_num_electrodes(index: torch.Tensor, num_electrodes: Union[int, None] = None) -> int:
    return int(index.max().item()) + 1 if num_electrodes is None else num_electrodes


def add_remaining_self_loops(edge_index: torch.Tensor,
                             edge_weight: Union[torch.Tensor, None] = None,
                             fill_value: float = 1.0,
                             num_electrodes: Union[int, None] = None) -> Tuple[torch.Tensor, torch.Tensor]:
    num_electrodes = maybe_num_electrodes(edge_index, num_electrodes)
    row, col = edge_index

    mask = row != col

    inv_mask = ~mask
    # print("inv_mask", inv_mask)

    loop_weight = torch.full((num_electrodes, ),
                             fill_value,
                             dtype=None if edge_weight is None else edge_weight.dtype,
                             device=edge_index.device)

    if edge_weight is not None:
        assert edge_weight.numel() == edge_index.size(1)
        remaining_edge_weight = edge_weight[inv_mask]

        if remaining_edge_weight.numel() > 0:
            loop_weight[row[inv_mask]] = remaining_edge_weight

        edge_weight = torch.cat([edge_weight[mask], loop_weight], dim=0)
    loop_index = torch.arange(
        0, num_electrodes, dtype=row.dtype, device=row.device)
    loop_index = loop_index.unsqueeze(0).repeat(2, 1)
    edge_index = torch.cat([edge_index[:, mask], loop_index], dim=1)

    return edge_index, edge_weight


class RGNN(BaseModule):
    r'''
    Regularized Graph Neural Networks (RGNN). For more details, please refer to the following information.

    - Paper: Zhong P, Wang D, Miao C. EEG-based emotion recognition using regularized graph neural networks[J]. IEEE Transactions on Affective Computing, 2020.
    - URL: https://ieeexplore.ieee.org/abstract/document/9091308
    - Related Project: https://github.com/zhongpeixiang/RGNN

    Below is a recommended suite for use in emotion recognition tasks:

    .. code-block:: python

        from torcheeg.datasets import SEEDDataset
        from torcheeg.models import RGNN
        from torcheeg.transforms.pyg import ToG
        from torcheeg.datasets.constants import SEED_STANDARD_ADJACENCY_MATRIX
        from torch_geometric.data import DataLoader


        dataset = SEEDDataset(root_path='./Preprocessed_EEG',
                              offline_transform=transforms.BandDifferentialEntropy(),
                              online_transform=ToG(SEED_STANDARD_ADJACENCY_MATRIX),
                              label_transform=transforms.Compose([
                                  transforms.Select('emotion'),
                                  transforms.Lambda(lambda x: int(x) + 1),
                              ]),
                              num_worker=8)

        model = RGNN(adj=torch.Tensor(SEED_STANDARD_ADJACENCY_MATRIX),
                     in_channels=5,
                     num_electrodes=62,
                     hid_channels=32,
                     num_layers=2,
                     num_classes=3,
                     dropout=0.7,
                     learn_edge_weights=True)

        x, y = next(iter(DataLoader(dataset, batch_size=64)))
        model(x)

    Args:
        adj (torch.Tensor): The adjacency matrix corresponding to the EEG representation, where 1.0 means the node is adjacent and 0.0 means the node is not adjacent. The matrix shape should be [num_electrodes, num_electrodes].
        num_electrodes (int): The number of electrodes. (default: :obj:`62`)
        in_channels (int): The feature dimension of each electrode. (default: :obj:`5`)
        num_layers (int): The number of graph convolutional layers. (default: :obj:`2`)
        hid_channels (int): The number of hidden nodes in the first fully connected layer. (default: :obj:`32`)
        num_classes (int): The number of classes to predict. (default: :obj:`3`)
        dropout (float): Probability of an element to be zeroed in the dropout layers at the output fully-connected layer. (default: :obj:`0.7`)
        learn_edge_weights (bool): Whether to learn a set of parameters to adjust the adjacency matrix. (default: :obj:`True`)
    '''

    def __init__(self,
                 adj: Union[torch.Tensor, list],
                 num_electrodes: int = 62,
                 in_channels: int = 5,
                 num_layers: int = 2,
                 hid_channels: int = 32,
                 num_classes: int = 3,
                 dropout: float = 0.7,
                 learn_edge_weights: bool = True):

        super(RGNN, self).__init__()
        self.num_electrodes = num_electrodes
        self.in_channels = in_channels
        self.num_electrodes = num_electrodes
        self.hid_channels = hid_channels
        self.num_layers = num_layers
        self.num_classes = num_classes
        self.dropout = dropout
        self.learn_edge_weights = learn_edge_weights

        self.xs, self.ys = torch.tril_indices(
            self.num_electrodes, self.num_electrodes, offset=0)
        if isinstance(adj, list):
            adj = torch.tensor(adj)
        adj = adj.reshape(self.num_electrodes, self.num_electrodes)[
            self.xs, self.ys]  # strict lower triangular values
        self.adj = nn.Parameter(adj, requires_grad=learn_edge_weights)

        class NewSGConv(pyg.nn.SGConv):
            def __init__(self,
                         in_channels: int,
                         out_channels: int,
                         num_layers: int = 1,
                         cached: bool = False,
                         bias: bool = True):
                super(NewSGConv, self).__init__(in_channels, out_channels,
                                                K=num_layers, cached=cached, bias=bias)

            # allow negative edge weights
            @staticmethod
            def norm(edge_index: torch.Tensor,
                     num_electrodes: int,
                     adj: torch.Tensor,
                     improved: bool = False,
                     dtype: Union[torch.dtype, None] = None):
                if adj is None:
                    adj = torch.ones((edge_index.size(1), ),
                                     dtype=dtype, device=edge_index.device)

                fill_value = 1 if not improved else 2
                edge_index, adj = add_remaining_self_loops(
                    edge_index, adj, fill_value, num_electrodes)
                row, col = edge_index
                deg = _scatter_add_1d(torch.abs(adj), row, dim_size=num_electrodes)
                deg_inv_sqrt = deg.pow(-0.5)
                deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0

                return edge_index, deg_inv_sqrt[row] * adj * deg_inv_sqrt[col]

            def forward(self, x: torch.Tensor, edge_index: torch.Tensor, adj: Union[None, torch.Tensor] = None) -> torch.Tensor:
                """"""
                if not self.cached or self.cached_result is None:
                    edge_index, edge_weight = NewSGConv.norm(
                        edge_index, x.size(0), adj, dtype=x.dtype)

                    for k in range(self.K):
                        x = self.propagate(
                            edge_index, x=x, edge_weight=edge_weight)
                    self.cached_result = x

                return self.lin(self.cached_result)

            def message(self, x_j: torch.Tensor, edge_weight: torch.Tensor) -> torch.Tensor:
                # x_j: (batch_size*num_electrodes*num_electrodes, in_channels)
                # edge_weight: (batch_size*num_electrodes*num_electrodes, )
                return edge_weight.view(-1, 1) * x_j

        self.conv1 = NewSGConv(in_channels=in_channels,
                               out_channels=hid_channels, num_layers=num_layers)
        self.fc = nn.Linear(hid_channels, num_classes)

    def forward(
        self,
        data_or_x,
        edge_index: Optional[torch.Tensor] = None,
        batch: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        r'''
        Args:
            data_or_x: PyG 的 Batch/Data，或张量形式的节点特征。
            edge_index (torch.Tensor | None): 当 data_or_x 为张量时，图边索引，shape 为 :obj:`[2, num_edges]`。
            batch (torch.Tensor | None): 当 data_or_x 为扁平节点张量时，对应每个节点所属图的 batch 索引。

        Returns:
            torch.Tensor[number of sample, number of classes]: the predicted probability that the samples belong to the classes.
        '''
        x, edge_index, batch, num_graphs = _unpack_graph_input(data_or_x, edge_index, batch)
        if batch is None:
            raise ValueError('batch information is required for graph-level pooling.')

        adj = torch.zeros(
            (self.num_electrodes, self.num_electrodes), device=edge_index.device)
        adj[self.xs.to(adj.device), self.ys.to(adj.device)] = self.adj
        adj = adj + adj.transpose(1, 0) - torch.diag(adj.diagonal())
        adj = adj.reshape(-1).repeat(num_graphs)
        x = F.relu(self.conv1(x, edge_index, adj))

        x = pyg.nn.global_add_pool(x, batch, size=num_graphs)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.fc(x)

        return x
