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

from ._ffi.object import ObjectBase, register_object

"""Module for heterogeneous graph index class definition."""
@register_object('graph.HeteroGraph')
class HeteroGraphIndex(ObjectBase):
    """HeteroGraph index object.

    Note
    ----
    Do not create GraphIndex directly.
    """
    def __new__(cls): # -> Self@HeteroGraphIndex:
        ...
    
    def __getstate__(self):
        """Issue: https://github.com/pytorch/pytorch/issues/32351
           Need to set the tensor created in the __getstate__ function
            as object attribute to avoid potential bugs
        """
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    @property
    def metagraph(self):
        """Meta graph

        Returns
        -------
        GraphIndex
            The meta graph.
        """
        ...
    
    def number_of_ntypes(self):
        """Return number of node types."""
        ...
    
    def number_of_etypes(self):
        """Return number of edge types."""
        ...
    
    def get_relation_graph(self, etype):
        """Get the unitgraph graph of the given edge/relation type.

        Parameters
        ----------
        etype : int
            The edge/relation type.

        Returns
        -------
        HeteroGraphIndex
            The unitgraph graph.
        """
        ...
    
    def flatten_relations(self, etypes):
        """Convert the list of requested unitgraph graphs into a single unitgraph
        graph.

        Parameters
        ----------
        etypes : list[int]
            The edge/relation types.

        Returns
        -------
        FlattenedHeteroGraph
            A flattened heterograph object
        """
        ...
    
    def add_nodes(self, ntype, num): # -> None:
        """Add nodes.

        Parameters
        ----------
        ntype : int
            Node type
        num : int
            Number of nodes to be added.
        """
        ...
    
    def add_edge(self, etype, u, v): # -> None:
        """Add one edge.

        Parameters
        ----------
        etype : int
            Edge type
        u : int
            The src node.
        v : int
            The dst node.
        """
        ...
    
    def add_edges(self, etype, u, v): # -> None:
        """Add many edges.

        Parameters
        ----------
        etype : int
            Edge type
        u : utils.Index
            The src nodes.
        v : utils.Index
            The dst nodes.
        """
        ...
    
    def clear(self): # -> None:
        """Clear the graph."""
        ...
    
    @property
    def dtype(self):
        """Return the data type of this graph index.

        Returns
        -------
        DGLDataType
            The data type of the graph.
        """
        ...
    
    @property
    def ctx(self):
        """Return the context of this graph index.

        Returns
        -------
        DGLContext
            The context of the graph.
        """
        ...
    
    def bits_needed(self, etype): # -> Literal[64, 32]:
        """Return the number of integer bits needed to represent the unitgraph graph.

        Parameters
        ----------
        etype : int
            The edge type.

        Returns
        -------
        int
            The number of bits needed.
        """
        ...
    
    def asbits(self, bits):
        """Transform the graph to a new one with the given number of bits storage.

        NOTE: this method only works for immutable graph index

        Parameters
        ----------
        bits : int
            The number of integer bits (32 or 64)

        Returns
        -------
        HeteroGraphIndex
            The graph index stored using the given number of bits.
        """
        ...
    
    def copy_to(self, ctx):
        """Copy this immutable graph index to the given device context.

        NOTE: this method only works for immutable graph index

        Parameters
        ----------
        ctx : DGLContext
            The target device context.

        Returns
        -------
        HeteroGraphIndex
            The graph index on the given device context.
        """
        ...
    
    def shared_memory(self, name, ntypes=..., etypes=..., formats=...):
        """Return a copy of this graph in shared memory

        Parameters
        ----------
        name : str
            The name of the shared memory.
        ntypes : list of str
            Name of node types
        etypes : list of str
            Name of edge types
        format : list of str
            Desired formats to be materialized.

        Returns
        -------
        HeteroGraphIndex
            The graph index in shared memory
        """
        ...
    
    def is_multigraph(self): # -> bool:
        """Return whether the graph is a multigraph
        The time cost will be O(E)

        Returns
        -------
        bool
            True if it is a multigraph, False otherwise.
        """
        ...
    
    def is_readonly(self): # -> bool:
        """Return whether the graph index is read-only.

        Returns
        -------
        bool
            True if it is a read-only graph, False otherwise.
        """
        ...
    
    def number_of_nodes(self, ntype):
        """Return the number of nodes.

        Parameters
        ----------
        ntype : int
            Node type

        Returns
        -------
        int
            The number of nodes
        """
        ...
    
    def number_of_edges(self, etype):
        """Return the number of edges.

        Parameters
        ----------
        etype : int
            Edge type

        Returns
        -------
        int
            The number of edges
        """
        ...
    
    def has_nodes(self, ntype, vids):
        """Return true if the nodes exist.

        Parameters
        ----------
        ntype : int
            Node type
        vid : Tensor
            Node IDs

        Returns
        -------
        Tensor
            0-1 array indicating existence
        """
        ...
    
    def has_edges_between(self, etype, u, v):
        """Return true if the edge exists.

        Parameters
        ----------
        etype : int
            Edge type
        u : Tensor
            Src node Ids.
        v : Tensor
            Dst node Ids.

        Returns
        -------
        Tensor
            0-1 array indicating existence
        """
        ...
    
    def predecessors(self, etype, v):
        """Return the predecessors of the node.

        Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.

        Parameters
        ----------
        etype : int
            Edge type
        v : int
            The node.

        Returns
        -------
        Tensor
            Array of predecessors
        """
        ...
    
    def successors(self, etype, v):
        """Return the successors of the node.

        Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.

        Parameters
        ----------
        etype : int
            Edge type
        v : int
            The node.

        Returns
        -------
        Tensor
            Array of successors
        """
        ...
    
    def edge_ids_all(self, etype, u, v): # -> tuple[Unknown, Unknown, Unknown]:
        """Return a triplet of arrays that contains the edge IDs.

        Parameters
        ----------
        etype : int
            Edge type
        u : Tensor
            The src nodes.
        v : Tensor
            The dst nodes.

        Returns
        -------
        Tensor
            The src nodes.
        Tensor
            The dst nodes.
        Tensor
            The edge ids.
        """
        ...
    
    def edge_ids_one(self, etype, u, v):
        """Return an arrays of edge IDs.

        Parameters
        ----------
        etype : int
            Edge type
        u : Tensor
            The src nodes.
        v : Tensor
            The dst nodes.

        Returns
        -------
        Tensor
            The edge ids.
        """
        ...
    
    def find_edges(self, etype, eid): # -> tuple[Unknown, Unknown, Unknown]:
        """Return a triplet of arrays that contains the edge IDs.

        Parameters
        ----------
        etype : int
            Edge type
        eid : Tensor
            Edge ids.

        Returns
        -------
        Tensor
            The src nodes.
        Tensor
            The dst nodes.
        Tensor
            The edge ids.
        """
        ...
    
    def in_edges(self, etype, v): # -> tuple[Unknown, Unknown, Unknown]:
        """Return the in edges of the node(s).

        Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.

        Parameters
        ----------
        etype : int
            Edge type
        v : Tensor
            Node IDs.

        Returns
        -------
        Tensor
            The src nodes.
        Tensor
            The dst nodes.
        Tensor
            The edge ids.
        """
        ...
    
    def out_edges(self, etype, v): # -> tuple[Unknown, Unknown, Unknown]:
        """Return the out edges of the node(s).

        Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.

        Parameters
        ----------
        etype : int
            Edge type
        v : Tensor
            Node IDs.

        Returns
        -------
        Tensor
            The src nodes.
        Tensor
            The dst nodes.
        Tensor
            The edge ids.
        """
        ...
    
    def edges(self, etype, order=...): # -> tuple[Unknown, Unknown, Unknown]:
        """Return all the edges

        Parameters
        ----------
        etype : int
            Edge type
        order : string
            The order of the returned edges. Currently support:

            - 'srcdst' : sorted by their src and dst ids.
            - 'eid'    : sorted by edge Ids.
            - None     : the arbitrary order.

        Returns
        -------
        Tensor
            The src nodes.
        Tensor
            The dst nodes.
        Tensor
            The edge ids.
        """
        ...
    
    def in_degrees(self, etype, v):
        """Return the in degrees of the nodes.

        Assume that node_type(v) == dst_type(etype). Thus, the ntype argument is omitted.

        Parameters
        ----------
        etype : int
            Edge type
        v : Tensor
            The nodes.

        Returns
        -------
        Tensor
            The in degree array.
        """
        ...
    
    def out_degrees(self, etype, v):
        """Return the out degrees of the nodes.

        Assume that node_type(v) == src_type(etype). Thus, the ntype argument is omitted.

        Parameters
        ----------
        etype : int
            Edge type
        v : Tensor
            The nodes.

        Returns
        -------
        Tensor
            The out degree array.
        """
        ...
    
    def adjacency_matrix(self, etype, transpose, ctx): # -> tuple[Unknown, Unknown]:
        """Return the adjacency matrix representation of this graph.

        By default, a row of returned adjacency matrix represents the source
        of an edge and the column represents the destination.

        When transpose is True, a row represents the destination and a column represents
        the source.

        Parameters
        ----------
        etype : int
            Edge type
        transpose : bool
            A flag to transpose the returned adjacency matrix.
        ctx : context
            The context of the returned matrix.

        Returns
        -------
        SparseTensor
            The adjacency matrix.
        Tensor
            A index for data shuffling due to sparse format change. Return None
            if shuffle is not required.
        """
        ...
    
    def adjacency_matrix_tensors(self, etype, transpose, fmt): # -> tuple[Unknown, Unknown, Unknown, Unknown, Unknown] | tuple[Unknown, Unknown, Unknown, Unknown]:
        """Return the adjacency matrix as a triplet of tensors.

        By default, a row of returned adjacency matrix represents the source
        of an edge and the column represents the destination.

        When transpose is True, a row represents the destination and a column represents
        the source.

        Parameters
        ----------
        etype : int
            Edge type
        transpose : bool
            A flag to transpose the returned adjacency matrix.
        fmt : str
            Indicates the format of returned adjacency matrix.

        Returns
        -------
        tuple[int, int, Tensor, Tensor] or tuple[int, int, Tensor, Tensor, Tensor]
            The number of rows and columns, followed by the adjacency matrix tensors
            whose data type and device are the same as those of the graph.

            If :attr:`fmt` is ``'coo'``, then the triplet will be
            the row array and column array of the COO representation.

            If :attr:`fmt` is ``'csr'``, then the triplet will be
            the index pointer array (``indptr``), indices array, and data array
            of the CSR representation.  The data array will contain the edge ID for
            each entry of the adjacency matrix.  If the data array is empty, then it is
            equivalent to a consecutive array from zero to the number of edges minus one.
        """
        ...
    
    def adjacency_matrix_scipy(self, etype, transpose, fmt, return_edge_ids=...):
        """Return the scipy adjacency matrix representation of this graph.

        By default, a row of returned adjacency matrix represents the destination
        of an edge and the column represents the source.

        When transpose is True, a row represents the source and a column represents
        a destination.

        Parameters
        ----------
        etype : int
            Edge type
        transpose : bool
            A flag to transpose the returned adjacency matrix.
        fmt : str
            Indicates the format of returned adjacency matrix.
        return_edge_ids : bool
            Indicates whether to return edge IDs or 1 as elements.

        Returns
        -------
        scipy.sparse.spmatrix
            The scipy representation of adjacency matrix.
        """
        ...
    
    def incidence_matrix(self, etype, typestr, ctx): # -> tuple[Unknown, Unknown]:
        """Return the incidence matrix representation of this graph.

        An incidence matrix is an n x m sparse matrix, where n is
        the number of nodes and m is the number of edges. Each nnz
        value indicating whether the edge is incident to the node
        or not.

        There are three types of an incidence matrix `I`:
        * "in":
          - I[v, e] = 1 if e is the in-edge of v (or v is the dst node of e);
          - I[v, e] = 0 otherwise.
        * "out":
          - I[v, e] = 1 if e is the out-edge of v (or v is the src node of e);
          - I[v, e] = 0 otherwise.
        * "both":
          - I[v, e] = 1 if e is the in-edge of v;
          - I[v, e] = -1 if e is the out-edge of v;
          - I[v, e] = 0 otherwise (including self-loop).

        Parameters
        ----------
        etype : int
            Edge type
        typestr : str
            Can be either "in", "out" or "both"
        ctx : context
            The context of returned incidence matrix.

        Returns
        -------
        SparseTensor
            The incidence matrix.
        utils.Index
            A index for data shuffling due to sparse format change. Return None
            if shuffle is not required.
        """
        ...
    
    def node_subgraph(self, induced_nodes, relabel_nodes):
        """Return the induced node subgraph.

        Parameters
        ----------
        induced_nodes : list of utils.Index
            Induced nodes. The length should be equal to the number of
            node types in this heterograph.
        relabel_nodes : bool
            If True, the extracted subgraph will only have the nodes in the specified node set
            and it will relabel the nodes in order.

        Returns
        -------
        SubgraphIndex
            The subgraph index.
        """
        ...
    
    def edge_subgraph(self, induced_edges, preserve_nodes):
        """Return the induced edge subgraph.

        Parameters
        ----------
        induced_edges : list of utils.Index
            Induced edges. The length should be equal to the number of
            edge types in this heterograph.
        preserve_nodes : bool
            Indicates whether to preserve all nodes or not.
            If true, keep the nodes which have no edge connected in the subgraph;
            If false, all nodes without edge connected to it would be removed.

        Returns
        -------
        SubgraphIndex
            The subgraph index.
        """
        ...
    
    def get_unitgraph(self, etype, ctx):
        """Create a unitgraph graph from given edge type and copy to the given device
        context.

        Note: this internal function is for DGL scheduler use only

        Parameters
        ----------
        etype : int
            If the graph index is a Bipartite graph index, this argument must be None.
            Otherwise, it represents the edge type.
        ctx : DGLContext
            The context of the returned graph.

        Returns
        -------
        HeteroGraphIndex
        """
        ...
    
    def get_csr_shuffle_order(self, etype): # -> tuple[Index, Index]:
        """Return the edge shuffling order when a coo graph is converted to csr format

        Parameters
        ----------
        etype : int
            The edge type

        Returns
        -------
        tuple of two utils.Index
            The first element of the tuple is the shuffle order for outward graph
            The second element of the tuple is the shuffle order for inward graph
        """
        ...
    
    def formats(self, formats=...): # -> dict[str, list[Unknown]]:
        """Get a graph index with the specified sparse format(s) or query
        for the usage status of sparse formats

        If the graph has multiple edge types, they will have the same
        sparse format.

        Parameters
        ----------
        formats : str or list of str or None

            * If formats is None, return the usage status of sparse formats
            * Otherwise, it can be ``'coo'``/``'csr'``/``'csc'`` or a sublist of
            them, specifying the sparse formats to use.

        Returns
        -------
        dict or GraphIndex

            * If formats is None, the result will be a dict recording the usage
              status of sparse formats.
            * Otherwise, a GraphIndex will be returned, which is a clone of the
              original graph with the specified sparse format(s) ``formats``.

        """
        ...
    
    def create_formats_(self):
        """Create all sparse matrices allowed for the graph."""
        ...
    
    def reverse(self):
        """Reverse the heterogeneous graph adjacency

        The node types and edge types are not changed.

        Returns
        -------
        A new graph index.
        """
        ...
    


@register_object('graph.HeteroSubgraph')
class HeteroSubgraphIndex(ObjectBase):
    """Hetero-subgraph data structure"""
    @property
    def graph(self):
        """The subgraph structure

        Returns
        -------
        HeteroGraphIndex
            The subgraph
        """
        ...
    
    @property
    def induced_nodes(self): # -> list[Unknown]:
        """Induced nodes for each node type. The return list
        length should be equal to the number of node types.

        Returns
        -------
        list of utils.Index
            Induced nodes
        """
        ...
    
    @property
    def induced_edges(self): # -> list[Unknown]:
        """Induced edges for each edge type. The return list
        length should be equal to the number of edge types.

        Returns
        -------
        list of utils.Index
            Induced edges
        """
        ...
    


def create_metagraph_index(ntypes, canonical_etypes): # -> tuple[Unknown, list[Unknown], list[Unknown], list[Unknown]]:
    """Return a GraphIndex instance for a metagraph given the node types and canonical
    edge types.

    This function will reorder the node types and canonical edge types.

    Parameters
    ----------
    ntypes : Iterable[str]
        The node types.
    canonical_etypes : Iterable[tuple[str, str, str]]
        The canonical edge types.

    Returns
    -------
    GraphIndex
        The index object for metagraph.
    list[str]
        The reordered node types for each node in the metagraph.
    list[str]
        The reordered edge types for each edge in the metagraph.
    list[tuple[str, str, str]]
        The reordered canonical edge types for each edge in the metagraph.
    """
    ...

def create_unitgraph_from_coo(num_ntypes, num_src, num_dst, row, col, formats, row_sorted=..., col_sorted=...):
    """Create a unitgraph graph index from COO format

    Parameters
    ----------
    num_ntypes : int
        Number of node types (must be 1 or 2).
    num_src : int
        Number of nodes in the src type.
    num_dst : int
        Number of nodes in the dst type.
    row : utils.Index
        Row index.
    col : utils.Index
        Col index.
    formats : list of str.
        Restrict the storage formats allowed for the unit graph.
    row_sorted : bool, optional
        Whether or not the rows of the COO are in ascending order.
    col_sorted : bool, optional
        Whether or not the columns of the COO are in ascending order within
        each row. This only has an effect when ``row_sorted`` is True.

    Returns
    -------
    HeteroGraphIndex
    """
    ...

def create_unitgraph_from_csr(num_ntypes, num_src, num_dst, indptr, indices, edge_ids, formats, transpose=...):
    """Create a unitgraph graph index from CSR format

    Parameters
    ----------
    num_ntypes : int
        Number of node types (must be 1 or 2).
    num_src : int
        Number of nodes in the src type.
    num_dst : int
        Number of nodes in the dst type.
    indptr : utils.Index
        CSR indptr.
    indices : utils.Index
        CSR indices.
    edge_ids : utils.Index
        Edge shuffle id.
    formats : str
        Restrict the storage formats allowed for the unit graph.
    transpose : bool, optional
        If True, treats the input matrix as CSC.

    Returns
    -------
    HeteroGraphIndex
    """
    ...

def create_heterograph_from_relations(metagraph, rel_graphs, num_nodes_per_type):
    """Create a heterograph from metagraph and graphs of every relation.

    Parameters
    ----------
    metagraph : GraphIndex
        Meta-graph.
    rel_graphs : list of HeteroGraphIndex
        Bipartite graph of each relation.
    num_nodes_per_type : utils.Index, optional
        Number of nodes per node type

    Returns
    -------
    HeteroGraphIndex
    """
    ...

def create_heterograph_from_shared_memory(name): # -> tuple[Unknown, list[Unknown], list[Unknown]]:
    """Create a heterograph from shared memory with the given name.

    Paramaters
    ----------
    name : str
        The name of the share memory

    Returns
    -------
    HeteroGraphIndex (in shared memory)
    ntypes : list of str
        Names of node types
    etypes : list of str
        Names of edge types
    """
    ...

def joint_union(metagraph, gidx_list):
    """Return a joint union of the input heterographs.

    Parameters
    ----------
    metagraph : GraphIndex
        Meta-graph.
    gidx_list : list of HeteroGraphIndex
        Heterographs to be joint_unioned.

    Returns
    -------
    HeteroGraphIndex
        joint_unioned Heterograph.
    """
    ...

def disjoint_union(metagraph, graphs):
    """Return a disjoint union of the input heterographs.

    Parameters
    ----------
    metagraph : GraphIndex
        Meta-graph.
    graphs : list of HeteroGraphIndex
        Heterographs to be batched.

    Returns
    -------
    HeteroGraphIndex
        Batched Heterograph.
    """
    ...

def disjoint_partition(graph, bnn_all_types, bne_all_types):
    """Partition the graph disjointly.

    Parameters
    ----------
    graph : HeteroGraphIndex
        The graph to be partitioned.
    bnn_all_types : list of list of int
        bnn_all_types[t] gives the number of nodes with t-th type in the batch.
    bne_all_types : list of list of int
        bne_all_types[t] gives the number of edges with t-th type in the batch.

    Returns
    --------
    list of HeteroGraphIndex
        Heterographs unbatched.
    """
    ...

def slice_gidx(graph, num_nodes, start_nid, num_edges, start_eid):
    """Slice a chunk of the graph.

    Parameters
    ----------
    graph : HeteroGraphIndex
        The batched graph to slice.
    num_nodes : utils.Index
        Number of nodes per node type in the result graph.
    start_nid : utils.Index
        Start node ID per node type in the result graph.
    num_edges : utils.Index
        Number of edges per edge type in the result graph.
    start_eid : utils.Index
        Start edge ID per edge type in the result graph.

    Returns
    -------
    HeteroGraphIndex
        The sliced graph.
    """
    ...

@register_object("graph.FlattenedHeteroGraph")
class FlattenedHeteroGraph(ObjectBase):
    """FlattenedHeteroGraph object class in C++ backend."""
    ...


@register_object("graph.HeteroPickleStates")
class HeteroPickleStates(ObjectBase):
    """Pickle states object class in C++ backend."""
    @property
    def version(self):
        """Version number

        Returns
        -------
        int
            version number
        """
        ...
    
    @property
    def meta(self): # -> bytearray:
        """Meta info

        Returns
        -------
        bytearray
            Serialized meta info
        """
        ...
    
    @property
    def arrays(self): # -> list[Unknown]:
        """Arrays representing the graph structure (COO or CSR)

        Returns
        -------
        list of dgl.ndarray.NDArray
            Arrays
        """
        ...
    
    def __getstate__(self): # -> tuple[Unknown, bytearray, list[Unknown]]:
        """Issue: https://github.com/pytorch/pytorch/issues/32351
           Need to set the tensor created in the __getstate__ function
            as object attribute to avoid potential bugs
        """
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    


