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

from abc import ABC

"""Define graph partition book."""
def get_shared_mem_partition_book(graph_name, graph_part): # -> RangePartitionBook | BasicPartitionBook | None:
    '''Get a graph partition book from shared memory.

    A graph partition book of a specific graph can be serialized to shared memory.
    We can reconstruct a graph partition book from shared memory.

    Parameters
    ----------
    graph_name : str
        The name of the graph.
    graph_part : DGLGraph
        The graph structure of a partition.

    Returns
    -------
    GraphPartitionBook
        A graph partition book for a particular partition.
    '''
    ...

class GraphPartitionBook(ABC):
    """ The base class of the graph partition book.

    For distributed training, a graph is partitioned into multiple parts and is loaded
    in multiple machines. The partition book contains all necessary information to locate
    nodes and edges in the cluster.

    The partition book contains various partition information, including

    * the number of partitions,
    * the partition ID that a node or edge belongs to,
    * the node IDs and the edge IDs that a partition has.
    * the local IDs of nodes and edges in a partition.

    Currently, there are two classes that implement ``GraphPartitionBook``:
    ``BasicGraphPartitionBook`` and ``RangePartitionBook``. ``BasicGraphPartitionBook``
    stores the mappings between every individual node/edge ID and partition ID on
    every machine, which usually consumes a lot of memory, while ``RangePartitionBook``
    calculates the mapping between node/edge IDs and partition IDs based on some small
    metadata because nodes/edges have been relabeled to have IDs in the same partition
    fall in a contiguous ID range. ``RangePartitionBook`` is usually a preferred way to
    provide mappings between node/edge IDs and partition IDs.

    A graph partition book is constructed automatically when a graph is partitioned.
    When a graph partition is loaded, a graph partition book is loaded as well.
    Please see :py:meth:`~dgl.distributed.partition.partition_graph`,
    :py:meth:`~dgl.distributed.partition.load_partition` and
    :py:meth:`~dgl.distributed.partition.load_partition_book` for more details.
    """
    def shared_memory(self, graph_name): # -> None:
        """Move the partition book to shared memory.

        Parameters
        ----------
        graph_name : str
            The graph name. This name will be used to read the partition book from shared
            memory in another process.
        """
        ...
    
    def num_partitions(self): # -> None:
        """Return the number of partitions.

        Returns
        -------
        int
            number of partitions
        """
        ...
    
    def metadata(self): # -> None:
        """Return the partition meta data.

        The meta data includes:

        * The machine ID.
        * Number of nodes and edges of each partition.

        Examples
        --------
        >>> print(g.get_partition_book().metadata())
        >>> [{'machine_id' : 0, 'num_nodes' : 3000, 'num_edges' : 5000},
        ...  {'machine_id' : 1, 'num_nodes' : 2000, 'num_edges' : 4888},
        ...  ...]

        Returns
        -------
        list[dict[str, any]]
            Meta data of each partition.
        """
        ...
    
    def nid2partid(self, nids, ntype): # -> None:
        """From global node IDs to partition IDs

        Parameters
        ----------
        nids : tensor
            global node IDs
        ntype : str
            The node type

        Returns
        -------
        tensor
            partition IDs
        """
        ...
    
    def eid2partid(self, eids, etype): # -> None:
        """From global edge IDs to partition IDs

        Parameters
        ----------
        eids : tensor
            global edge IDs
        etype : str
            The edge type

        Returns
        -------
        tensor
            partition IDs
        """
        ...
    
    def partid2nids(self, partid, ntype): # -> None:
        """From partition id to global node IDs

        Parameters
        ----------
        partid : int
            partition id
        ntype : str
            The node type

        Returns
        -------
        tensor
            node IDs
        """
        ...
    
    def partid2eids(self, partid, etype): # -> None:
        """From partition id to global edge IDs

        Parameters
        ----------
        partid : int
            partition id
        etype : str
            The edge type

        Returns
        -------
        tensor
            edge IDs
        """
        ...
    
    def nid2localnid(self, nids, partid, ntype): # -> None:
        """Get local node IDs within the given partition.

        Parameters
        ----------
        nids : tensor
            global node IDs
        partid : int
            partition ID
        ntype : str
            The node type

        Returns
        -------
        tensor
             local node IDs
        """
        ...
    
    def eid2localeid(self, eids, partid, etype): # -> None:
        """Get the local edge ids within the given partition.

        Parameters
        ----------
        eids : tensor
            global edge IDs
        partid : int
            partition ID
        etype : str
            The edge type

        Returns
        -------
        tensor
             local edge IDs
        """
        ...
    
    @property
    def partid(self): # -> None:
        """Get the current partition ID

        Return
        ------
        int
            The partition ID of current machine
        """
        ...
    
    @property
    def ntypes(self): # -> None:
        """Get the list of node types
        """
        ...
    
    @property
    def etypes(self): # -> None:
        """Get the list of edge types
        """
        ...
    
    def map_to_per_ntype(self, ids): # -> None:
        """Map homogeneous node IDs to type-wise IDs and node types.

        Parameters
        ----------
        ids : tensor
            Homogeneous node IDs.

        Returns
        -------
        (tensor, tensor)
            node type IDs and type-wise node IDs.
        """
        ...
    
    def map_to_per_etype(self, ids): # -> None:
        """Map homogeneous edge IDs to type-wise IDs and edge types.

        Parameters
        ----------
        ids : tensor
            Homogeneous edge IDs.

        Returns
        -------
        (tensor, tensor)
            edge type IDs and type-wise edge IDs.
        """
        ...
    
    def map_to_homo_nid(self, ids, ntype): # -> None:
        """Map type-wise node IDs and type IDs to homogeneous node IDs.

        Parameters
        ----------
        ids : tensor
            Type-wise node Ids
        ntype : str
            node type

        Returns
        -------
        Tensor
            Homogeneous node IDs.
        """
        ...
    
    def map_to_homo_eid(self, ids, etype): # -> None:
        """Map type-wise edge IDs and type IDs to homogeneous edge IDs.

        Parameters
        ----------
        ids : tensor
            Type-wise edge Ids
        etype : str
            edge type

        Returns
        -------
        Tensor
            Homogeneous edge IDs.
        """
        ...
    


class BasicPartitionBook(GraphPartitionBook):
    """This provides the most flexible way to store parition information.

    The partition book maintains the mapping of every single node IDs and edge IDs to
    partition IDs. This is very flexible at the coast of large memory consumption.
    On a large graph, the mapping consumes significant memory and this partition book
    is not recommended.

    Parameters
    ----------
    part_id : int
        partition ID of current partition book
    num_parts : int
        number of total partitions
    node_map : tensor
        global node ID mapping to partition ID
    edge_map : tensor
        global edge ID mapping to partition ID
    part_graph : DGLGraph
        The graph partition structure.
    """
    def __init__(self, part_id, num_parts, node_map, edge_map, part_graph) -> None:
        ...
    
    def shared_memory(self, graph_name): # -> None:
        """Move data to shared memory.
        """
        ...
    
    def num_partitions(self): # -> int:
        """Return the number of partitions.
        """
        ...
    
    def metadata(self): # -> list[Unknown]:
        """Return the partition meta data.
        """
        ...
    
    def map_to_per_ntype(self, ids): # -> tuple[Unknown, Unknown]:
        """Map global homogeneous node IDs to node type IDs.
        Returns
            type_ids, per_type_ids
        """
        ...
    
    def map_to_per_etype(self, ids): # -> tuple[Unknown, Unknown]:
        """Map global homogeneous edge IDs to edge type IDs.
        Returns
            type_ids, per_type_ids
        """
        ...
    
    def map_to_homo_nid(self, ids, ntype):
        """Map per-node-type IDs to global node IDs in the homogeneous format.
        """
        ...
    
    def map_to_homo_eid(self, ids, etype):
        """Map per-edge-type IDs to global edge IDs in the homoenegeous format.
        """
        ...
    
    def nid2partid(self, nids, ntype=...):
        """From global node IDs to partition IDs
        """
        ...
    
    def eid2partid(self, eids, etype=...):
        """From global edge IDs to partition IDs
        """
        ...
    
    def partid2nids(self, partid, ntype=...):
        """From partition id to global node IDs
        """
        ...
    
    def partid2eids(self, partid, etype=...):
        """From partition id to global edge IDs
        """
        ...
    
    def nid2localnid(self, nids, partid, ntype=...):
        """Get local node IDs within the given partition.
        """
        ...
    
    def eid2localeid(self, eids, partid, etype=...):
        """Get the local edge ids within the given partition.
        """
        ...
    
    @property
    def partid(self): # -> int:
        """Get the current partition ID
        """
        ...
    
    @property
    def ntypes(self): # -> list[str]:
        """Get the list of node types
        """
        ...
    
    @property
    def etypes(self): # -> list[str]:
        """Get the list of edge types
        """
        ...
    


class RangePartitionBook(GraphPartitionBook):
    """This partition book supports more efficient storage of partition information.

    This partition book is used if the nodes and edges of a graph partition are assigned
    with contiguous IDs. It uses very small amount of memory to store the partition
    information.

    Parameters
    ----------
    part_id : int
        partition ID of current partition book
    num_parts : int
        number of total partitions
    node_map : dict[str, Tensor]
        Global node ID ranges within partitions for each node type. The key is the node type
        name in string. The value is a tensor of shape :math:`(K, 2)`, where :math:`K` is
        the number of partitions. Each row has two integers: the starting and the ending IDs
        for a particular node type in a partition. For example, all nodes of type ``"T"`` in
        partition ``i`` has ID range ``node_map["T"][i][0]`` to ``node_map["T"][i][1]``.
    edge_map : dict[str, Tensor]
        Global edge ID ranges within partitions for each edge type. The key is the edge type
        name in string. The value is a tensor of shape :math:`(K, 2)`, where :math:`K` is
        the number of partitions. Each row has two integers: the starting and the ending IDs
        for a particular edge type in a partition. For example, all edges of type ``"T"`` in
        partition ``i`` has ID range ``edge_map["T"][i][0]`` to ``edge_map["T"][i][1]``.
    ntypes : dict[str, int]
        map ntype strings to ntype IDs.
    etypes : dict[str, int]
        map etype strings to etype IDs.
    """
    def __init__(self, part_id, num_parts, node_map, edge_map, ntypes, etypes) -> None:
        ...
    
    def shared_memory(self, graph_name): # -> None:
        """Move data to shared memory.
        """
        ...
    
    def num_partitions(self): # -> Unknown:
        """Return the number of partitions.
        """
        ...
    
    def metadata(self): # -> list[Unknown]:
        """Return the partition meta data.
        """
        ...
    
    def map_to_per_ntype(self, ids): # -> tuple[Unknown, Unknown]:
        """Map global homogeneous node IDs to node type IDs.
        Returns
            type_ids, per_type_ids
        """
        ...
    
    def map_to_per_etype(self, ids): # -> tuple[Unknown, Unknown]:
        """Map global homogeneous edge IDs to edge type IDs.
        Returns
            type_ids, per_type_ids
        """
        ...
    
    def map_to_homo_nid(self, ids, ntype):
        """Map per-node-type IDs to global node IDs in the homogeneous format.
        """
        ...
    
    def map_to_homo_eid(self, ids, etype):
        """Map per-edge-type IDs to global edge IDs in the homoenegeous format.
        """
        ...
    
    def nid2partid(self, nids, ntype=...):
        """From global node IDs to partition IDs
        """
        ...
    
    def eid2partid(self, eids, etype=...):
        """From global edge IDs to partition IDs
        """
        ...
    
    def partid2nids(self, partid, ntype=...):
        """From partition ID to global node IDs
        """
        ...
    
    def partid2eids(self, partid, etype=...):
        """From partition ID to global edge IDs
        """
        ...
    
    def nid2localnid(self, nids, partid, ntype=...):
        """Get local node IDs within the given partition.
        """
        ...
    
    def eid2localeid(self, eids, partid, etype=...):
        """Get the local edge IDs within the given partition.
        """
        ...
    
    @property
    def partid(self): # -> Unknown:
        """Get the current partition ID.
        """
        ...
    
    @property
    def ntypes(self): # -> list[None]:
        """Get the list of node types
        """
        ...
    
    @property
    def etypes(self): # -> list[None]:
        """Get the list of edge types
        """
        ...
    


NODE_PART_POLICY = ...
EDGE_PART_POLICY = ...
class PartitionPolicy:
    """This defines a partition policy for a distributed tensor or distributed embedding.

    When DGL shards tensors and stores them in a cluster of machines, it requires
    partition policies that map rows of the tensors to machines in the cluster.

    Although an arbitrary partition policy can be defined, DGL currently supports
    two partition policies for mapping nodes and edges to machines. To define a partition
    policy from a graph partition book, users need to specify the policy name ('node' or 'edge').

    Parameters
    ----------
    policy_str : str
        Partition policy name, e.g., 'edge:_E' or 'node:_N'.
    partition_book : GraphPartitionBook
        A graph partition book
    """
    def __init__(self, policy_str, partition_book) -> None:
        ...
    
    @property
    def policy_str(self): # -> str:
        """Get the policy name

        Returns
        -------
        str
            The name of the partition policy.
        """
        ...
    
    @property
    def part_id(self):
        """Get partition ID

        Returns
        -------
        int
            The partition ID
        """
        ...
    
    @property
    def partition_book(self):
        """Get partition book

        Returns
        -------
        GraphPartitionBook
            The graph partition book
        """
        ...
    
    def get_data_name(self, name): # -> HeteroDataName:
        """Get HeteroDataName
        """
        ...
    
    def to_local(self, id_tensor):
        """Mapping global ID to local ID.

        Parameters
        ----------
        id_tensor : tensor
            Gloabl ID tensor

        Return
        ------
        tensor
            local ID tensor
        """
        ...
    
    def to_partid(self, id_tensor):
        """Mapping global ID to partition ID.

        Parameters
        ----------
        id_tensor : tensor
            Global ID tensor

        Return
        ------
        tensor
            partition ID
        """
        ...
    
    def get_part_size(self): # -> int:
        """Get data size of current partition.

        Returns
        -------
        int
            data size
        """
        ...
    
    def get_size(self):
        """Get the full size of the data.

        Returns
        -------
        int
            data size
        """
        ...
    


class NodePartitionPolicy(PartitionPolicy):
    '''Partition policy for nodes.
    '''
    def __init__(self, partition_book, ntype=...) -> None:
        ...
    


class EdgePartitionPolicy(PartitionPolicy):
    '''Partition policy for edges.
    '''
    def __init__(self, partition_book, etype=...) -> None:
        ...
    


class HeteroDataName:
    ''' The data name in a heterogeneous graph.

    A unique data name has three components:
    * indicate it's node data or edge data.
    * indicate the node/edge type.
    * the name of the data.

    Parameters
    ----------
    is_node : bool
        Indicate whether it's node data or edge data.
    entity_type : str
        The type of the node/edge.
    data_name : str
        The name of the data.
    '''
    def __init__(self, is_node, entity_type, data_name) -> None:
        ...
    
    def is_node(self): # -> bool:
        ''' Is this the name of node data
        '''
        ...
    
    def is_edge(self): # -> bool:
        ''' Is this the name of edge data
        '''
        ...
    
    def get_type(self): # -> str:
        ''' The type of the node/edge.
        This is only meaningful in a heterogeneous graph.
        In homogeneous graph, type is '_N' for a node and '_E' for an edge.
        '''
        ...
    
    def get_name(self): # -> Unknown:
        ''' The name of the data.
        '''
        ...
    
    def __str__(self) -> str:
        ''' The full name of the data.

        The full name is used as the key in the KVStore.
        '''
        ...
    


def parse_hetero_data_name(name): # -> HeteroDataName:
    '''Parse data name and create HeteroDataName.

    The data name has a specialized format. We can parse the name to determine if
    it's node data or edge data, node/edge type and its actual name. The data name
    has three fields and they are separated by ":".

    Parameters
    ----------
    name : str
        The data name

    Returns
    -------
    HeteroDataName
    '''
    ...

