import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import Batch
from torch_geometric.nn import SGConv, global_add_pool
from torch_scatter import scatter_add

from typing import Union, Tuple


def maybe_num_electrodes(index: torch.Tensor, num_electrodes: Union[int, None] = None) -> int:
    return 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 NewSGConv(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(torch.abs(adj), row, dim=0, 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, norm = NewSGConv.norm(edge_index, x.size(0), adj, dtype=x.dtype)

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

        return self.lin(self.cached_result)

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


class RGNN(torch.nn.Module):
    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

        dataset = SEEDDataset(io_path=f'./seed',
                              root_path='./Preprocessed_EEG',
                              online_transform=transforms.ToG(SEED_ADJACENCY_MATRIX),,
                              label_transform=transforms.Compose([
                                  transforms.Select(['emotion']),
                                  transforms.Lambda(x: x + 1)
                              ]))
        model = RGNN(adj=torch.Tensor(SEED_ADJACENCY_MATRIX),
                     in_channels=4,
                     num_electrodes=62,
                     hid_channels=32,
                     num_layers=2,
                     num_classes=2,
                     dropout=0.7,
                     domain_adaptation=False,
                     beta=0.0,
                     learn_edge_weights=True)

    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. (defualt: :obj:`62`)
        in_channels (int): The feature dimension of each electrode. (defualt: :obj:`4`)
        num_layers (int): The number of graph convolutional layers. (defualt: :obj:`2`)
        hid_channels (int): The number of hidden nodes in the first fully connected layer. (defualt: :obj:`32`)
        num_classes (int): The number of classes to predict. (defualt: :obj:`2`)
        dropout (float): Probability of an element to be zeroed in the dropout layers at the output fully-connected layer. (defualt: :obj:`0.7`)
        learn_edge_weights (bool): Whether to learn a set of parameters to adjust the adjacency matrix. (defualt: :obj:`True`)
    '''
    def __init__(self,
                 adj: torch.Tensor,
                 num_electrodes: int = 62,
                 in_channels: int = 4,
                 num_layers: int = 2,
                 hid_channels: int = 32,
                 num_classes: int = 2,
                 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)
        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)

        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: Batch) -> torch.Tensor:
        r'''
        Args:
            x (torch_geometric.data.Batch): EEG signal representation, the ideal input shape is :obj:`[n, 62, 4]`. Here, :obj:`n` corresponds to the batch size, :obj:`62` corresponds to :obj:`num_electrodes`, and :obj:`4` corresponds to :obj:`in_channels`.

        Returns:
            torch.Tensor[number of sample, number of classes]: the predicted probability that the samples belong to the classes.
        '''
        batch_size = data.num_graphs
        x, edge_index = data.x, data.edge_index
        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())  # copy values from lower tri to upper tri
        adj = adj.reshape(-1).repeat(batch_size)
        x = F.relu(self.conv1(x, edge_index, adj))

        x = global_add_pool(x, data.batch, size=batch_size)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.fc(x)

        return x