from icosagon.data import Data
from icosagon.bulkdec import BulkDecodeLayer
from icosagon.input import OneHotInputLayer
from icosagon.convlayer import DecagonLayer
import torch


def test_bulk_decode_layer_01():
    data = Data()
    data.add_node_type('Dummy', 100)
    fam = data.add_relation_family('Dummy-Dummy', 0, 0, False)
    fam.add_relation_type('Dummy Relation 1',
        torch.rand((100, 100), dtype=torch.float32).round().to_sparse())

    in_layer = OneHotInputLayer(data)
    d_layer = DecagonLayer(in_layer.output_dim, 32, data)
    dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
        keep_prob=1., activation=lambda x: x)
    seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)

    pred = seq(None)

    assert isinstance(pred, list)
    assert len(pred) == len(data.relation_families)
    assert isinstance(pred[0], torch.Tensor)
    assert len(pred[0].shape) == 3
    assert len(pred[0]) == len(data.relation_families[0].relation_types)
    assert pred[0].shape[1] == data.node_types[0].count
    assert pred[0].shape[2] == data.node_types[0].count


def test_bulk_decode_layer_02():
    data = Data()
    data.add_node_type('Foo', 100)
    data.add_node_type('Bar', 50)
    fam = data.add_relation_family('Foo-Bar', 0, 1, False)
    fam.add_relation_type('Foobar Relation 1',
        torch.rand((100, 50), dtype=torch.float32).round().to_sparse(),
        torch.rand((50, 100), dtype=torch.float32).round().to_sparse())

    in_layer = OneHotInputLayer(data)
    d_layer = DecagonLayer(in_layer.output_dim, 32, data)
    dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
        keep_prob=1., activation=lambda x: x)
    seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)

    pred = seq(None)

    assert isinstance(pred, list)
    assert len(pred) == len(data.relation_families)
    assert isinstance(pred[0], torch.Tensor)
    assert len(pred[0].shape) == 3
    assert len(pred[0]) == len(data.relation_families[0].relation_types)
    assert pred[0].shape[1] == data.node_types[0].count
    assert pred[0].shape[2] == data.node_types[1].count


def test_bulk_decode_layer_03():
    data = Data()
    data.add_node_type('Foo', 100)
    data.add_node_type('Bar', 50)
    fam = data.add_relation_family('Foo-Bar', 0, 1, False)
    fam.add_relation_type('Foobar Relation 1',
        torch.rand((100, 50), dtype=torch.float32).round().to_sparse(),
        torch.rand((50, 100), dtype=torch.float32).round().to_sparse())
    fam.add_relation_type('Foobar Relation 2',
        torch.rand((100, 50), dtype=torch.float32).round().to_sparse(),
        torch.rand((50, 100), dtype=torch.float32).round().to_sparse())

    in_layer = OneHotInputLayer(data)
    d_layer = DecagonLayer(in_layer.output_dim, 32, data)
    dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
        keep_prob=1., activation=lambda x: x)
    seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)

    pred = seq(None)

    assert isinstance(pred, list)
    assert len(pred) == len(data.relation_families)
    assert isinstance(pred[0], torch.Tensor)
    assert len(pred[0].shape) == 3
    assert len(pred[0]) == len(data.relation_families[0].relation_types)
    assert pred[0].shape[1] == data.node_types[0].count
    assert pred[0].shape[2] == data.node_types[1].count


def test_bulk_decode_layer_03_big():
    data = Data()
    data.add_node_type('Foo', 2000)
    data.add_node_type('Bar', 2100)
    fam = data.add_relation_family('Foo-Bar', 0, 1, False)
    fam.add_relation_type('Foobar Relation 1',
        torch.rand((2000, 2100), dtype=torch.float32).round().to_sparse(),
        torch.rand((2100, 2000), dtype=torch.float32).round().to_sparse())
    fam.add_relation_type('Foobar Relation 2',
        torch.rand((2000, 2100), dtype=torch.float32).round().to_sparse(),
        torch.rand((2100, 2000), dtype=torch.float32).round().to_sparse())

    in_layer = OneHotInputLayer(data)
    d_layer = DecagonLayer(in_layer.output_dim, 32, data)
    dec_layer = BulkDecodeLayer(input_dim=d_layer.output_dim, data=data,
        keep_prob=1., activation=lambda x: x)
    seq = torch.nn.Sequential(in_layer, d_layer, dec_layer)

    pred = seq(None)

    assert isinstance(pred, list)
    assert len(pred) == len(data.relation_families)
    assert isinstance(pred[0], torch.Tensor)
    assert len(pred[0].shape) == 3
    assert len(pred[0]) == len(data.relation_families[0].relation_types)
    assert pred[0].shape[1] == data.node_types[0].count
    assert pred[0].shape[2] == data.node_types[1].count