import torch
from ogb.utils.features import get_bond_feature_dims

from torch_geometric.graphgym.register import (
    register_edge_encoder,
    register_node_encoder,
)


@register_node_encoder('example')
class ExampleNodeEncoder(torch.nn.Module):
    """Provides an encoder for integer node features.

    Args:
        num_classes (int): The number of classes for the embedding mapping to
            learn.
    """
    def __init__(self, emb_dim, num_classes=None):
        super().__init__()

        self.encoder = torch.nn.Embedding(num_classes, emb_dim)
        torch.nn.init.xavier_uniform_(self.encoder.weight.data)

    def forward(self, batch):
        # Encode just the first dimension if more exist
        batch.x = self.encoder(batch.x[:, 0])

        return batch


@register_edge_encoder('example')
class ExampleEdgeEncoder(torch.nn.Module):
    def __init__(self, emb_dim):
        super().__init__()

        self.bond_embedding_list = torch.nn.ModuleList()
        full_bond_feature_dims = get_bond_feature_dims()

        for dim in full_bond_feature_dims:
            emb = torch.nn.Embedding(dim, emb_dim)
            torch.nn.init.xavier_uniform_(emb.weight.data)
            self.bond_embedding_list.append(emb)

    def forward(self, batch):
        bond_embedding = 0
        for i in range(batch.edge_feature.shape[1]):
            bond_embedding += \
                self.bond_embedding_list[i](batch.edge_attr[:, i])

        batch.edge_attr = bond_embedding
        return batch
