import pprint

import torch
import yaml
from torch_geometric.data.data import Data

from hmpn.common.hmpn_util import make_batch
from hmpn.get_hmpn import get_hmpn_from_graph


def main():
    # load the config file
    with open('config.yaml', 'r') as file:
        config = yaml.safe_load(file)
    pprint.pprint(config)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}\n\n")

    # build a graph with 2 nodes with 3 features, fully connected edges with 2 features and 4 global features
    graph = Data(x=torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32),  # 2 nodes with 3 features each
                 edge_index=torch.tensor([[0, 1, 1, 0], [1, 0, 0, 1]]),  # fully connected
                 edge_attr=torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]], dtype=torch.float32),  # 4 edges, 2 features
                 u=torch.tensor([[1, 2, 3, 4]], dtype=torch.float32),  # shape (1, num_features)
                 )
    graph.to(device)
    print(f"Graph: {graph}\n\n")

    # build the message passing network from this graph. In this case, the network is a homogeneous message passing
    # network since the graph is homogeneous
    hmpn = get_hmpn_from_graph(example_graph=graph,
                               latent_dimension=config.get("latent_dimension"),
                               base_config=config.get("base"),
                               unpack_output=True,
                               device=device)
    print(f"Network: {hmpn}\n\n")

    # run the network on the graph, getting a latent representation of all node, edge and global features as an output
    batch = make_batch(graph)
    print(f"Batch: {batch}")
    out = hmpn(batch)
    print(f"Output: {out}\n\n")


if __name__ == '__main__':
    main()
