from triacontagon.data import Data
from triacontagon.sampling import negative_sample_adj_mat, \
    negative_sample_data
from triacontagon.decode import dedicom_decoder
import torch


def test_negative_sample_adj_mat_01():
    adj_mat = torch.tensor([
        [0, 1, 0, 1, 0],
        [0, 0, 0, 0, 1],
        [1, 1, 0, 0, 0],
        [0, 0, 1, 0, 1],
        [0, 1, 0, 0, 0]
    ])

    print('adj_mat:', adj_mat)

    adj_mat_neg = negative_sample_adj_mat(adj_mat)

    print('adj_mat_neg:', adj_mat_neg.to_dense())


def test_negative_sample_data_01():
    d = Data()
    d.add_vertex_type('Gene', 5)

    d.add_edge_type('Gene-Gene', 0, 0, [
        torch.tensor([
            [0, 1, 0, 1, 0],
            [0, 0, 0, 0, 1],
            [1, 1, 0, 0, 0],
            [0, 0, 1, 0, 1],
            [0, 1, 0, 0, 0]
        ], dtype=torch.float).to_sparse()
    ], dedicom_decoder)

    d_neg = negative_sample_data(d)