from points.datasets import get_dataset

from torch_geometric.transforms import RadiusGraph


def print_dataset(train_dataset, test_dataset):
    num_nodes = num_edges = 0
    for data in train_dataset:
        data = RadiusGraph(0.2)(data)
        num_nodes += data.num_nodes
        num_edges += data.num_edges
    for data in test_dataset:
        data = RadiusGraph(0.2)(data)
        num_nodes += data.num_nodes
        num_edges += data.num_edges

    num_graphs = len(train_dataset) + len(test_dataset)
    print('Graphs', num_graphs)
    print('Nodes', num_nodes / num_graphs)
    print('Edges', (num_edges // 2) / num_graphs)
    print('Label rate', len(train_dataset) / num_graphs)
    print()


print_dataset(*get_dataset(num_points=1024))