from typing import Callable, Tuple, Optional, List

import torch
from torch_geometric.data import Batch


class ModelInterface(torch.nn.Module):
    r"""
    Provides the signature for any main model to be trained under PyDGN

    Args:
        dim_node_features (int): dimension of node features (according to the :class:`~pydgn.data.dataset.DatasetInterface` property)
        dim_edge_features (int): dimension of edge features (according to the :class:`~pydgn.data.dataset.DatasetInterface` property)
        dim_target (int): dimension of the target (according to the :class:`~pydgn.data.dataset.DatasetInterface` property)
        readout_class (Callable[...,:class:`torch.nn.Module`]): class of the module implementing the readout. This is optional, but useful to put different readouts to try in the config file
        config (dict): config dictionary containing all the necessary hyper-parameters plus additional information (if needed)
    """
    def __init__(self,
                 dim_node_features: int,
                 dim_edge_features: int,
                 dim_target: int,
                 readout_class: Callable[...,torch.nn.Module],
                 config: dict):
        super().__init__()
        self.dim_node_features = dim_node_features
        self.dim_edge_features = dim_edge_features
        self.dim_target = dim_target
        self.readout_class = readout_class
        self.config = config

    def forward(self, data: Batch) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[object]]]:
        r"""
        Performs a forward pass over a batch of graphs

        Args:
            data (:class:`torch_geometric.data.Batch`): a batch of graphs

        Returns:
            a tuple (model's output, [optional] node embeddings, [optional] additional outputs
        """
        raise NotImplementedError('You need to implement this method!')


class ReadoutInterface(torch.nn.Module):
    r"""
    Provides the signature for any readout to be trained under PyDGN

    Args:
        dim_node_features (int): dimension of node features (according to the :class:`~pydgn.data.dataset.DatasetInterface` property)
        dim_edge_features (int): dimension of edge features (according to the :class:`~pydgn.data.dataset.DatasetInterface` property)
        dim_target (int): dimension of the target (according to the :class:`~pydgn.data.dataset.DatasetInterface` property)
        config (dict): config dictionary containing all the necessary hyper-parameters plus additional information (if needed)
    """
    def __init__(self,
                 dim_node_features: int,
                 dim_edge_features: int,
                 dim_target: int,
                 config: dict):
        super().__init__()
        self.dim_node_features = dim_node_features
        self.dim_edge_features = dim_edge_features
        self.dim_target = dim_target
        self.config = config

    def forward(self,
                node_embeddings: torch.tensor,
                batch: torch.Tensor,
                **kwargs) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[List[object]]]:
        r"""
        Performs a forward pass over a batch of graphs

        Args:
            node_embeddings (:class:`torch_geometric.data.Batch`): the node embeddings
            batch (:class:`torch.Tensor`): the usual ``batch`` object of PyG
            kwargs (dict): additional and optional arguments

        Returns:
            a tuple (model's output, [optional] node embeddings, [optional] additional outputs
        """
        raise NotImplementedError('You need to implement this method!')