import decagon_pytorch.convolve
import decagon.deep.layers
import torch
import tensorflow as tf
import numpy as np


def prepare_data():
    np.random.seed(0)
    latent = np.random.random((5, 10)).astype(np.float32)
    latent[latent < .5] = 0
    latent = np.ceil(latent)
    adjacency_matrices = []
    for _ in range(5):
        adj_mat = np.random.random((len(latent),) * 2).astype(np.float32)
        adj_mat[adj_mat < .5] = 0
        adj_mat = np.ceil(adj_mat)
        adjacency_matrices.append(adj_mat)
    return latent, adjacency_matrices


def sparse_graph_conv_torch():
    torch.random.manual_seed(0)
    latent, adjacency_matrices = prepare_data()
    print('latent.dtype:', latent.dtype)
    latent = torch.tensor(latent).to_sparse()
    adj_mat = adjacency_matrices[0]
    adj_mat = torch.tensor(adj_mat).to_sparse()
    print('adj_mat.dtype:', adj_mat.dtype,
        'latent.dtype:', latent.dtype)
    conv = decagon_pytorch.convolve.SparseGraphConv(10, 10,
        adj_mat)
    latent = conv(latent)
    return latent


def dense_to_sparse_tf(x):
    a, b = np.where(x)
    indices = np.array([a, b]).T
    values = x[a, b]
    return tf.sparse.SparseTensor(indices, values, x.shape)



def sparse_graph_conv_tf():
    torch.random.manual_seed(0)
    latent, adjacency_matrices = prepare_data()
    conv_torch = decagon_pytorch.convolve.SparseGraphConv(10, 10,
        torch.tensor(adjacency_matrices[0]).to_sparse())
    weight = tf.constant(conv_torch.weight.detach().numpy())
    latent = dense_to_sparse_tf(latent)
    adj_mat = dense_to_sparse_tf(adjacency_matrices[0])
    latent = tf.sparse_tensor_dense_matmul(latent, weight)
    latent = tf.sparse_tensor_dense_matmul(adj_mat, latent)
    return latent


def test_sparse_graph_conv():
    latent_torch = sparse_graph_conv_torch()
    latent_tf = sparse_graph_conv_tf()
    assert np.all(latent_torch.detach().numpy() == latent_tf.eval(session = tf.Session()))


def test_sparse_dropout_grap_conv_activation():
    pass


def test_sparse_multi_dgca():
    pass