from icosagon.fastmodel import FastModel
from icosagon.data import Data
from icosagon.trainprep import prepare_training, \
    TrainValTest
import torch
import time


def _make_symmetric(x: torch.Tensor):
    x = (x + x.transpose(0, 1)) / 2
    return x


def _symmetric_random(n_rows, n_columns):
    return _make_symmetric(torch.rand((n_rows, n_columns),
        dtype=torch.float32).round().to_sparse())


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

    fam = d.add_relation_family('Drug-Gene', 1, 0, True)
    fam.add_relation_type('Target',
        torch.rand((100, 1000), dtype=torch.float32).round().to_sparse())

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

    fam = d.add_relation_family('Drug-Drug', 1, 1, True)
    for i in range(500):
        fam.add_relation_type('Side Effect: Nausea %d' % i,
            _symmetric_random(100, 100))
        fam.add_relation_type('Side Effect: Infertility %d' % i,
            _symmetric_random(100, 100))
        fam.add_relation_type('Side Effect: Death %d' % i,
            _symmetric_random(100, 100))
    return d


def test_fast_model_01():
    d = _some_data_with_interactions()
    prep_d = prepare_training(d, TrainValTest(.8, .1, .1))
    model = FastModel(prep_d)
    for i in range(10):
        t = time.time()
        _ = model(None)
        print('Model forward took:', time.time() - t)