from icosagon.convolve import GraphConv, \
    DropoutGraphConvActivation, \
    MultiDGCA
import torch


def _test_graph_conv_01(use_sparse: bool):
    adj_mat = torch.rand((10, 20))
    adj_mat[adj_mat < .5] = 0
    adj_mat = torch.ceil(adj_mat)

    node_reprs = torch.eye(20)

    graph_conv = GraphConv(20, 20, adj_mat.to_sparse() \
        if use_sparse else adj_mat)
    graph_conv.weight = torch.eye(20)

    res = graph_conv(node_reprs)
    assert torch.all(res == adj_mat)


def _test_graph_conv_02(use_sparse: bool):
    adj_mat = torch.rand((10, 20))
    adj_mat[adj_mat < .5] = 0
    adj_mat = torch.ceil(adj_mat)

    node_reprs = torch.eye(20)

    graph_conv = GraphConv(20, 20, adj_mat.to_sparse() \
        if use_sparse else adj_mat)
    graph_conv.weight = torch.eye(20) * 2

    res = graph_conv(node_reprs)
    assert torch.all(res == adj_mat * 2)


def _test_graph_conv_03(use_sparse: bool):
    adj_mat = torch.tensor([
        [1, 0, 1, 0, 1, 0], # [1, 0, 0]
        [1, 0, 1, 0, 0, 1], # [1, 0, 0]
        [1, 1, 0, 1, 0, 0], # [0, 1, 0]
        [0, 0, 0, 1, 0, 1], # [0, 1, 0]
        [1, 1, 1, 1, 1, 1], # [0, 0, 1]
        [0, 0, 0, 1, 1, 1]  # [0, 0, 1]
    ], dtype=torch.float32)

    expect = torch.tensor([
        [1, 1, 1],
        [1, 1, 1],
        [2, 1, 0],
        [0, 1, 1],
        [2, 2, 2],
        [0, 1, 2]
    ], dtype=torch.float32)

    node_reprs = torch.eye(6)

    graph_conv = GraphConv(6, 3, adj_mat.to_sparse() \
        if use_sparse else adj_mat)
    graph_conv.weight = torch.tensor([
        [1, 0, 0],
        [1, 0, 0],
        [0, 1, 0],
        [0, 1, 0],
        [0, 0, 1],
        [0, 0, 1]
    ], dtype=torch.float32)

    res = graph_conv(node_reprs)
    assert torch.all(res == expect)


def test_graph_conv_dense_01():
    _test_graph_conv_01(use_sparse=False)


def test_graph_conv_dense_02():
    _test_graph_conv_02(use_sparse=False)


def test_graph_conv_dense_03():
    _test_graph_conv_03(use_sparse=False)


def test_graph_conv_sparse_01():
    _test_graph_conv_01(use_sparse=True)


def test_graph_conv_sparse_02():
    _test_graph_conv_02(use_sparse=True)


def test_graph_conv_sparse_03():
    _test_graph_conv_03(use_sparse=True)