import math
import tensorlayerx as tlx
from gammagl.layers.conv import MessagePassing
from tensorlayerx.nn import ModuleDict, Linear, Parameter, ParameterDict
from gammagl.utils import segment_softmax
from gammagl.layers.conv.hgt_conv import HGTConv
from scipy import sparse
import numpy as np
from gammagl.data import HeteroGraph

def test_hgt_conv_same_dimensions():
    x_dict={
        'author':tlx.random_normal(shape=(4,16)),
        'paper':tlx.random_normal(shape=(6,16))
    }
    
    index1=tlx.convert_to_tensor([0, 0, 1, 0, 3, 1, 0, 2, 0, 2, 0, 0, 2, 0, 3, 3, 1, 0, 2, 3],dtype=tlx.int32)
    index2=tlx.convert_to_tensor([5, 3, 2, 3, 3, 1, 1, 5, 1, 5, 5, 0, 3, 2, 1, 5, 1, 5, 5, 1],dtype=tlx.int32)

    edge_index_dict={
        ("author","writes","paper"):tlx.stack([index1,index2]),
        ("paper","written_by","author"):tlx.stack([index2,index1]),
    }

    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
    conv=HGTConv(16,16,metadata,heads=2)
    out_dict1=conv(x_dict,edge_index_dict)
    assert len(out_dict1) == 2 
    assert tlx.get_tensor_shape(out_dict1['author']) == [4, 16] 
    assert tlx.get_tensor_shape(out_dict1['paper']) == [6, 16]

"""
def test_hgt_conv_different_dimensions():
    x_dict = {
        'author': tlx.random_normal(shape=(4, 16)),
        'paper': tlx.random_normal(shape=(6, 32)),
    }

    index1 = tlx.convert_to_tensor([3, 1, 0, 3, 1, 3, 0, 0, 2, 1, 0, 3, 1, 1, 2, 1, 2, 3, 2, 2],dtype=tlx.int32)
    index2 = tlx.convert_to_tensor([2, 3, 5, 2, 5, 1, 1, 2, 5, 1, 3, 1, 1, 2, 2, 1, 2, 2, 2, 1],dtype=tlx.int32)

    edge_index_dict = {
        ('author', 'writes', 'paper'): tlx.stack([index1, index2]),
        ('paper', 'written_by', 'author'): tlx.stack([index2, index1]),
    }

    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))

    conv = HGTConv(in_channels={
        'author': 16,
        'paper': 32
    }, out_channels=32, metadata=metadata, heads=2)
    out_dict1 = conv(x_dict, edge_index_dict) #有问题,hgt_conv的121计算out时维度不匹配
    #assert out_dict1['author'].size() == (4, 32)
    #assert out_dict1['paper'].size() == (6, 32)
"""

def test_hgt_conv_lazy():
    x_dict={
        'author':tlx.random_normal(shape=(4,16)),
        'paper':tlx.random_normal(shape=(6,16))
    }

    index1=tlx.convert_to_tensor([0, 0, 1, 0, 3, 1, 0, 2, 0, 2, 0, 0, 2, 0, 3, 3, 1, 0, 2, 3],dtype=tlx.int32)
    index2=tlx.convert_to_tensor([5, 3, 2, 3, 3, 1, 1, 5, 1, 5, 5, 0, 3, 2, 1, 5, 1, 5, 5, 1],dtype=tlx.int32)



    edge_index_dict={
        ("author","writes","paper"):tlx.stack([index1,index2]),
        ("paper","written_by","author"):tlx.stack([index2,index1]),
    }

    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))

    conv=HGTConv(None,16,metadata,heads=2) #此处的None在py-g中为-1
    out_dict1=conv(x_dict,edge_index_dict)
    assert len(out_dict1) == 2
    assert tlx.get_tensor_shape(out_dict1['author']) == [4, 16]
    assert tlx.get_tensor_shape(out_dict1['paper']) == [6, 16]

"""
def test_hgt_conv_out_of_place():
    data = HeteroGraph()
    data["author"].x=tlx.random_normal(shape=(4,16))
    data["paper"].x=tlx.random_normal(shape=(6,32))

    index1=tlx.convert_to_tensor([0, 0, 1, 0, 3, 1, 0, 2, 0, 2, 0, 0, 2, 0, 3, 3, 1, 0, 2, 3],dtype=tlx.int32)
    index2=tlx.convert_to_tensor([5, 3, 2, 3, 3, 1, 1, 5, 1, 5, 5, 0, 3, 2, 1, 5, 1, 5, 5, 1],dtype=tlx.int32)

    data['author', 'paper'].edge_index = tlx.stack([index1, index2], axis=0)
    data['paper', 'author'].edge_index = tlx.stack([index2, index1], axis=0)

    conv = HGTConv(None, 64, data.metadata(), heads=1)
    x_dict, edge_index_dict = data.x_dict, data.edge_index_dict
    assert x_dict['author'].shape == (4, 16)
    assert x_dict['paper'].shape== (6, 32)
    _ = conv(x_dict, edge_index_dict)#有问题,hgt_conv的121计算out时维度不匹配
"""
