import torch


def gather_scatter_python(input, edges, directed: bool = False):
    """
    Python implementation of gather_scatter for aggregating features of
    neighbor nodes in a graph.

    Given a directed graph: v0 -> v1 -> v2 the updated feature for v1 depends
    on v2 in order to be consistent with Morris et al. AAAI 2019
    (https://arxiv.org/abs/1810.02244). This only affects
    directed graphs; for undirected graphs v1 will depend on both v0 and v2,
    no matter which way the edges are physically stored.

    Args:
        input: Tensor of shape (num_vertices, input_dim).
        edges: Tensor of edge indices of shape (num_edges, 2).
        directed: bool indicating if edges are directed.

    Returns:
        output: Tensor of same shape as input.
    """
    if not (input.dim() == 2):
        raise ValueError("input can only have 2 dimensions.")
    if not (edges.dim() == 2):
        raise ValueError("edges can only have 2 dimensions.")
    if not (edges.shape[1] == 2):
        raise ValueError("edges must be of shape (num_edges, 2).")

    num_vertices, input_feature_dim = input.shape
    num_edges = edges.shape[0]
    output = torch.zeros_like(input)
    idx0 = edges[:, 0].view(num_edges, 1).expand(num_edges, input_feature_dim)
    idx1 = edges[:, 1].view(num_edges, 1).expand(num_edges, input_feature_dim)

    output = output.scatter_add(0, idx0, input.gather(0, idx1))
    if not directed:
        output = output.scatter_add(0, idx1, input.gather(0, idx0))
    return output