import torch
from triacontagon.model import Model, \
    TrainingBatch, \
    _per_layer_required_vertices
from triacontagon.data import Data
from triacontagon.decode import dedicom_decoder
from triacontagon.util import common_one_hot_encoding


def test_per_layer_required_vertices_01():
    d = Data()
    d.add_vertex_type('Gene', 4)
    d.add_vertex_type('Drug', 5)

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

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

    d.add_edge_type('Drug-Drug', 1, 1, [ torch.tensor([
        [0, 0, 1, 0, 1],
        [0, 0, 0, 1, 1],
        [1, 0, 0, 0, 0],
        [0, 1, 0, 0, 1],
        [1, 1, 0, 1, 0]
    ]).to_sparse() ], dedicom_decoder)

    batch = TrainingBatch(0, 1, 0, torch.tensor([
        [0, 1]
    ]))

    res = _per_layer_required_vertices(d, batch, 5)
    print('res:', res)


def test_model_convolve_01():
    d = Data()
    d.add_vertex_type('Gene', 4)
    d.add_vertex_type('Drug', 5)

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

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

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

    model = Model(d, [9, 32, 64], keep_prob=1.0,
        conv_activation = torch.sigmoid,
        dec_activation = torch.sigmoid)

    repr_1 = torch.eye(9)
    repr_1[4:, 4:] = 0
    repr_2 = torch.eye(9)
    repr_2[:4, :4] = 0

    in_layer_repr = [
        repr_1[:4, :].to_sparse(),
        repr_2[4:, :].to_sparse()
    ]

    _ = model.convolve(in_layer_repr)


def test_model_decode_01():
    d = Data()
    d.add_vertex_type('Gene', 100)

    gene_gene = torch.rand(100, 100).round()
    gene_gene = gene_gene - torch.diag(torch.diag(gene_gene))
    d.add_edge_type('Gene-Gene', 0, 0, [
        gene_gene.to_sparse()
    ], dedicom_decoder)

    b = TrainingBatch(0, 0, 0, torch.tensor([
        [0, 1],
        [10, 51],
        [50, 60],
        [70, 90],
        [98, 99]
    ]), torch.ones(5))

    in_repr = common_one_hot_encoding([100])

    in_repr = [ in_repr[0].to_dense() ]

    m = Model(d, [100], 1.0, torch.sigmoid, torch.sigmoid)

    _ = m.decode(in_repr, b)