import torch
from torch import nn

from torch_geometric.utils import to_dense_batch

from models.base.graph_encoder import GraphEncoder
from models.base.config import GraphPromptEncoderConfig


class BaseGraphPromptEncoder(nn.Module):
    """
    Base class to define the graph prompt encoder for encoding graph-based structural information.
    
    Args:
        config ([`GraphPromptEncoderConfig`]): The configuration of the graph prompt encoder.
    """
    def __init__(self, config: GraphPromptEncoderConfig):
        super().__init__()
        # Graph encoder settings
        self.graph_encoder = GraphEncoder(config)
    
    def reshape_node_embeds(
        self,
        x: torch.Tensor,
        batch_index: torch.Tensor,
    )-> tuple[torch.Tensor, torch.Tensor]:
        """
        Reshape the node embeddings with batch axis to match the input word embeddings.
        
        Args:
            x (torch.Tensor): Node embeddings of shape (num_nodes, output_size).
            batch_index (torch.Tensor): Batch indices for each node (shape: [num_nodes]).
        
        Returns:
            tuple[torch.Tensor, torch.Tensor]: A tuple containing the reshaped node embeddings and the corresponding node mask.
        """
        # node_embeds shape: [batch_size, max_num_nodes, output_size]
        # node_mask shape: [batch_size, max_num_nodes] (dtype: torch.bool)
        node_embeds, node_mask = to_dense_batch(x, batch_index, fill_value=0.0)
        
        return node_embeds, node_mask.long()