#
# Copyright (C) Stanislaw Adaszewski, 2020
# License: GPLv3
#


from icosagon.data import Data, \
    _equal, \
    RelationFamily
from icosagon.decode import DEDICOMDecoder
import torch
import pytest


def test_equal_01():
    x = torch.rand((10, 10))
    y = torch.rand((10, 10)).round().to_sparse()

    print('x == x ?')
    assert torch.all(_equal(x, x))
    print('y == y ?')
    assert torch.all(_equal(y, y))
    print('x == y ?')
    with pytest.raises(ValueError):
        _equal(x, y)

    print('y == z ?')
    z = torch.rand((10, 10)).round().to_sparse()
    assert not torch.all(_equal(y, z))


def test_relation_family_01():
    d = Data()
    d.add_node_type('Whatever', 10)

    fam = RelationFamily(d, 'Dummy-Dummy', 0, 0, True, DEDICOMDecoder)

    with pytest.raises(ValueError):
        fam.add_relation_type('Dummy-Dummy', None, None)

    with pytest.raises(ValueError):
        fam.add_relation_type('Dummy-Dummy', 'bad-value', None)

    with pytest.raises(ValueError):
        fam.add_relation_type('Dummy-Dummy', None, 'bad-value')

    with pytest.raises(ValueError):
        fam.add_relation_type('Dummy-Dummy', torch.rand((5, 5)), None)

    with pytest.raises(ValueError):
        fam.add_relation_type('Dummy-Dummy', None, torch.rand((5, 5)))

    with pytest.raises(ValueError):
        fam.add_relation_type('Dummy-Dummy', torch.rand((10, 10)), torch.rand((10, 10)))

    with pytest.raises(ValueError):
        fam.add_relation_type('Dummy-Dummy', torch.rand((10, 10)), None)


def test_relation_family_02():
    d = Data()
    d.add_node_type('A', 10)
    d.add_node_type('B', 5)

    fam = RelationFamily(d, 'A-B', 0, 1, True, DEDICOMDecoder)

    with pytest.raises(ValueError):
        fam.add_relation_type('A-B', torch.rand((10, 5)).round(),
            torch.rand((5, 10)).round())


def test_relation_family_03():
    d = Data()
    d.add_node_type('A', 10)
    d.add_node_type('B', 5)

    fam = RelationFamily(d, 'A-B', 0, 1, True, DEDICOMDecoder)

    fam.add_relation_type('A-B', torch.rand((10, 5)).round())

    assert torch.all(fam.relation_types[0].adjacency_matrix.transpose(0, 1) == \
        fam.relation_types[0].adjacency_matrix_backward)


def test_data_01():
    d = Data()
    d.add_node_type('Gene', 1000)
    d.add_node_type('Drug', 100)
    dummy_0 = torch.zeros((100, 1000))
    dummy_1 = torch.zeros((1000, 100))
    dummy_2 = torch.zeros((100, 100))
    dummy_3 = torch.zeros((1000, 1000))

    fam = d.add_relation_family('Drug-Gene', 1, 0, True)
    fam.add_relation_type('Target', dummy_0)

    fam = d.add_relation_family('Gene-Gene', 0, 0, True)
    fam.add_relation_type('Interaction', dummy_3)

    fam = d.add_relation_family('Drug-Drug', 1, 1, True)
    fam.add_relation_type('Side Effect: Nausea', dummy_2)
    fam.add_relation_type('Side Effect: Infertility', dummy_2)
    fam.add_relation_type('Side Effect: Death', dummy_2)

    print(d)


def test_data_02():
    d = Data()
    d.add_node_type('Gene', 1000)
    d.add_node_type('Drug', 100)

    dummy_0 = torch.zeros((100, 1000))
    dummy_1 = torch.zeros((1000, 100))
    dummy_2 = torch.zeros((100, 100))
    dummy_3 = torch.zeros((1000, 1000))

    fam = d.add_relation_family('Drug-Gene', 1, 0, True)
    with pytest.raises(ValueError):
        fam.add_relation_type('Target', dummy_1)

    fam = d.add_relation_family('Gene-Gene', 0, 0, True)
    with pytest.raises(ValueError):
        fam.add_relation_type('Interaction', dummy_2)

    fam = d.add_relation_family('Drug-Drug', 1, 1, True)
    with pytest.raises(ValueError):
        fam.add_relation_type('Side Effect: Nausea', dummy_3)
    with pytest.raises(ValueError):
        fam.add_relation_type('Side Effect: Infertility', dummy_3)
    with pytest.raises(ValueError):
        fam.add_relation_type('Side Effect: Death', dummy_3)
    print(d)


def test_data_03():
    d = Data()
    d.add_node_type('Gene', 1000)
    d.add_node_type('Drug', 100)
    fam = d.add_relation_family('Drug-Gene', 1, 0, True)
    with pytest.raises(ValueError):
        fam.add_relation_type('Target', None)
    print(d)