from triacontagon.data import Data
from triacontagon.sampling import fixed_unigram_candidate_sampler, \
    get_true_classes, \
    negative_sample_adj_mat, \
    negative_sample_data, \
    get_edges_and_degrees
import triacontagon.sampling
from triacontagon.decode import dedicom_decoder
import torch
import time
import pytest


def test_fixed_unigram_candidate_sampler_01():
    true_classes = torch.tensor([[-1],
        [-1],
        [ 3],
        [ 2],
        [-1]])
    num_repeats = torch.tensor([0, 0, 1, 1, 0])
    unigrams = torch.tensor([0., 0., 1., 1., 0.], dtype=torch.float64)
    distortion = 0.75
    res = fixed_unigram_candidate_sampler(true_classes, num_repeats,
        unigrams, distortion)
    print('res:', res)


def test_fixed_unigram_candidate_sampler_02():
    foo_bar = torch.tensor([
        [0, 1, 0, 1],
        [0, 0, 0, 1],
        [0, 1, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 1, 1]
    ], dtype=torch.float32)

    # bar_foo = foo_bar.transpose(0, 1).to_sparse().coalesce()
    bar_foo = foo_bar.to_sparse().coalesce()

    true_classes, row_count = get_true_classes(bar_foo)
    print('true_classes:', true_classes)
    print('row_count:', row_count)

    edges_pos, degrees = get_edges_and_degrees(bar_foo)
    print('degrees:', degrees)

    res = fixed_unigram_candidate_sampler(true_classes, row_count,
        degrees, 0.75)
    print('res:', res)


def test_get_true_classes_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]
    ], dtype=torch.float).to_sparse()

    true_classes, row_count = get_true_classes(adj_mat)
    print('true_classes:', true_classes)

    true_classes = torch.repeat_interleave(true_classes, row_count, dim=0)

    assert torch.all(true_classes == torch.tensor([
        [1, 3],
        [1, 3],
        [4, -1],
        [0, 1],
        [0, 1],
        [2, 4],
        [2, 4],
        [1, -1]
    ]))


def test_get_true_classes_02():
    adj_mat = torch.rand(2000, 2000).round().to_sparse()

    t = time.time()
    true_classes, row_count = get_true_classes(adj_mat)
    print('Elapsed:', time.time() - t)

    print('true_classes.shape:', true_classes.shape)


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.to_sparse())

    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)


def test_fixed_unigram_candidate_sampler_new_01():
    if 'fixed_unigram_candidate_sampler_new' not in dir(triacontagon.sampling):
        pytest.skip('fixed_unigram_candidate_sampler_new not found')
    x = (torch.rand((10, 10)) < .05).to(torch.float32).to_sparse()
    true_classes, row_count = get_true_classes(x)
    edges, degrees = get_edges_and_degrees(x)
    # import pdb
    # pdb.set_trace()
    _ = triacontagon.sampling.fixed_unigram_candidate_sampler_new(true_classes,
        row_count, degrees, 0.75)