import torch

from torch_geometric.datasets import KarateClub
from torch_geometric.nn import GCNConv
from torch_geometric.visualization import influence


class Net(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, out_channels)
        self.conv2 = GCNConv(out_channels, out_channels)

    def forward(self, x, edge_index):
        x = torch.nn.functional.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x


def test_influence():
    data = KarateClub()[0]
    x = torch.randn(data.num_nodes, 8)

    out = influence(Net(x.size(1), 16), x, data.edge_index)
    assert out.size() == (data.num_nodes, data.num_nodes)
    assert torch.allclose(out.sum(dim=-1), torch.ones(data.num_nodes),
                          atol=1e-04)