"""
This type stub file was generated by pyright.
"""

from .rpc import Request, Response

"""A set of graph services of getting subgraphs from DistGraph"""
__all__ = ['sample_neighbors', 'in_subgraph', 'find_edges']
SAMPLING_SERVICE_ID = ...
INSUBGRAPH_SERVICE_ID = ...
EDGES_SERVICE_ID = ...
OUTDEGREE_SERVICE_ID = ...
INDEGREE_SERVICE_ID = ...
class SubgraphResponse(Response):
    """The response for sampling and in_subgraph"""
    def __init__(self, global_src, global_dst, global_eids) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown, Unknown]:
        ...
    


class FindEdgeResponse(Response):
    """The response for sampling and in_subgraph"""
    def __init__(self, global_src, global_dst, order_id) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown, Unknown]:
        ...
    


class SamplingRequest(Request):
    """Sampling Request"""
    def __init__(self, nodes, fan_out, edge_dir=..., prob=..., replace=...) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown, Unknown, Unknown, Unknown]:
        ...
    
    def process_request(self, server_state): # -> SubgraphResponse:
        ...
    


class EdgesRequest(Request):
    """Edges Request"""
    def __init__(self, edge_ids, order_id) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown]:
        ...
    
    def process_request(self, server_state): # -> FindEdgeResponse:
        ...
    


class InDegreeRequest(Request):
    """In-degree Request"""
    def __init__(self, n, order_id) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown]:
        ...
    
    def process_request(self, server_state): # -> InDegreeResponse:
        ...
    


class InDegreeResponse(Response):
    """The response for in-degree"""
    def __init__(self, deg, order_id) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown]:
        ...
    


class OutDegreeRequest(Request):
    """Out-degree Request"""
    def __init__(self, n, order_id) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown]:
        ...
    
    def process_request(self, server_state): # -> OutDegreeResponse:
        ...
    


class OutDegreeResponse(Response):
    """The response for out-degree"""
    def __init__(self, deg, order_id) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self): # -> tuple[Unknown, Unknown]:
        ...
    


class InSubgraphRequest(Request):
    """InSubgraph Request"""
    def __init__(self, nodes) -> None:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def __getstate__(self):
        ...
    
    def process_request(self, server_state): # -> SubgraphResponse:
        ...
    


def merge_graphs(res_list, num_nodes): # -> DGLHeteroGraph:
    """Merge request from multiple servers"""
    ...

LocalSampledGraph = ...
def sample_neighbors(g, nodes, fanout, edge_dir=..., prob=..., replace=...): # -> DGLHeteroGraph:
    """Sample from the neighbors of the given nodes from a distributed graph.

    For each node, a number of inbound (or outbound when ``edge_dir == 'out'``) edges
    will be randomly chosen.  The returned graph will contain all the nodes in the
    original graph, but only the sampled edges.

    Node/edge features are not preserved. The original IDs of
    the sampled edges are stored as the `dgl.EID` feature in the returned graph.

    This version provides an experimental support for heterogeneous graphs.
    When the input graph is heterogeneous, the sampled subgraph is still stored in
    the homogeneous graph format. That is, all nodes and edges are assigned with
    unique IDs (in contrast, we typically use a type name and a node/edge ID to
    identify a node or an edge in ``DGLGraph``). We refer to this type of IDs
    as *homogeneous ID*.
    Users can use :func:`dgl.distributed.GraphPartitionBook.map_to_per_ntype`
    and :func:`dgl.distributed.GraphPartitionBook.map_to_per_etype`
    to identify their node/edge types and node/edge IDs of that type.

    For heterogeneous graphs, ``nodes`` can be a dictionary whose key is node type
    and the value is type-specific node IDs; ``nodes`` can also be a tensor of
    *homogeneous ID*.

    Parameters
    ----------
    g : DistGraph
        The distributed graph..
    nodes : tensor or dict
        Node IDs to sample neighbors from. If it's a dict, it should contain only
        one key-value pair to make this API consistent with dgl.sampling.sample_neighbors.
    fanout : int
        The number of edges to be sampled for each node.

        If -1 is given, all of the neighbors will be selected.
    edge_dir : str, optional
        Determines whether to sample inbound or outbound edges.

        Can take either ``in`` for inbound edges or ``out`` for outbound edges.
    prob : str, optional
        Feature name used as the (unnormalized) probabilities associated with each
        neighboring edge of a node.  The feature must have only one element for each
        edge.

        The features must be non-negative floats, and the sum of the features of
        inbound/outbound edges for every node must be positive (though they don't have
        to sum up to one).  Otherwise, the result will be undefined.
    replace : bool, optional
        If True, sample with replacement.

        When sampling with replacement, the sampled subgraph could have parallel edges.

        For sampling without replacement, if fanout > the number of neighbors, all the
        neighbors are sampled. If fanout == -1, all neighbors are collected.

    Returns
    -------
    DGLGraph
        A sampled subgraph containing only the sampled neighboring edges.  It is on CPU.
    """
    ...

def find_edges(g, edge_ids): # -> tuple[Unknown, Unknown]:
    """ Given an edge ID array, return the source and destination
    node ID array ``s`` and ``d`` from a distributed graph.
    ``s[i]`` and ``d[i]`` are source and destination node ID for
    edge ``eid[i]``.

    Parameters
    ----------
    g : DistGraph
        The distributed graph.
    edges : tensor
        The edge ID array.

    Returns
    -------
    tensor
        The source node ID array.
    tensor
        The destination node ID array.
    """
    ...

def in_subgraph(g, nodes): # -> DGLHeteroGraph:
    """Return the subgraph induced on the inbound edges of the given nodes.

    The subgraph keeps the same type schema and all the nodes are preserved regardless
    of whether they have an edge or not.

    Node/edge features are not preserved. The original IDs of
    the extracted edges are stored as the `dgl.EID` feature in the returned graph.

    For now, we only support the input graph with one node type and one edge type.


    Parameters
    ----------
    g : DistGraph
        The distributed graph structure.
    nodes : tensor or dict
        Node ids to sample neighbors from.

    Returns
    -------
    DGLGraph
        The subgraph.

        One can retrieve the mapping from subgraph edge ID to parent
        edge ID via ``dgl.EID`` edge features of the subgraph.
    """
    ...

def in_degrees(g, v): # -> None:
    '''Get in-degrees
    '''
    ...

def out_degrees(g, u): # -> None:
    '''Get out-degrees
    '''
    ...

