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

from collections.abc import Mapping
from ..base import DGLError

"""Internal utilities."""
class InconsistentDtypeException(DGLError):
    """Exception class for inconsistent dtype between graph and tensor"""
    def __init__(self, msg=..., *args, **kwargs) -> None:
        ...
    


class Index:
    """Index class that can be easily converted to list/tensor."""
    def __init__(self, data, dtype=...) -> None:
        ...
    
    def __iter__(self): # -> Generator[int, None, None]:
        ...
    
    def __len__(self): # -> Any | int:
        ...
    
    def __getitem__(self, i): # -> int:
        ...
    
    def tonumpy(self): # -> ndarray[Unknown, Unknown]:
        """Convert to a numpy ndarray."""
        ...
    
    def tousertensor(self, ctx=...):
        """Convert to user tensor (defined in `backend`)."""
        ...
    
    def todgltensor(self): # -> NDArray:
        """Convert to dgl.NDArray."""
        ...
    
    def slice_data(self): # -> slice | None:
        """Return the internal slice data.

        If this index is not initialized from slice, the return will be None.
        """
        ...
    
    def is_slice(self, start, stop): # -> bool:
        """Check if Index wraps a slice data with given start and stop"""
        ...
    
    def __getstate__(self): # -> tuple[slice, Unknown | str] | tuple[Unknown, Unknown | str]:
        ...
    
    def __setstate__(self, state): # -> None:
        ...
    
    def get_items(self, index): # -> Index:
        """Return values at given positions of an Index

        Parameters
        ----------
        index: utils.Index

        Returns
        -------
        utils.Index
            The values at the given position.
        """
        ...
    
    def set_items(self, index, value): # -> Index:
        """Set values at given positions of an Index. Set is not done in place,
        instead, a new Index object will be returned.

        Parameters
        ----------
        index: utils.Index
            Positions to set values
        value: int or utils.Index
            Values to set. If value is an integer, then all positions are set
            to the same value

        Returns
        -------
        utils.Index
            The new values.
        """
        ...
    
    def append_zeros(self, num): # -> Self@Index | Index:
        """Append zeros to an Index

        Parameters
        ----------
        num: int
            number of zeros to append
        """
        ...
    
    def nonzero(self): # -> Index:
        """Return the nonzero positions"""
        ...
    
    def has_nonzero(self):
        """Check if there is any nonzero value in this Index"""
        ...
    


def toindex(data, dtype=...): # -> Index:
    """Convert the given data to Index object.

    Parameters
    ----------
    data : index data
        Data to create the index.

    Returns
    -------
    Index
        The index object.

    See Also
    --------
    Index
    """
    ...

def zero_index(size, dtype=...): # -> Index:
    """Create a index with provided size initialized to zero

    Parameters
    ----------
    size: int
    """
    ...

def set_diff(ar1, ar2): # -> Index:
    """Find the set difference of two index arrays.
    Return the unique values in ar1 that are not in ar2.

    Parameters
    ----------
    ar1: utils.Index
        Input index array.

    ar2: utils.Index
        Input comparison index array.

    Returns
    -------
    setdiff:
        Array of values in ar1 that are not in ar2.
    """
    ...

class LazyDict(Mapping):
    """A readonly dictionary that does not materialize the storage."""
    def __init__(self, fn, keys) -> None:
        ...
    
    def __getitem__(self, key):
        ...
    
    def __contains__(self, key): # -> bool:
        ...
    
    def __iter__(self): # -> Iterator[_T@iter]:
        ...
    
    def __len__(self): # -> int:
        ...
    
    def keys(self): # -> Unknown:
        ...
    


class HybridDict(Mapping):
    """A readonly dictonary that merges several dict-like (python dict, LazyDict).

    If there are duplicate keys, early keys have priority over latter ones.
    """
    def __init__(self, *dict_like_list) -> None:
        ...
    
    def keys(self): # -> set[Unknown]:
        ...
    
    def __getitem__(self, key):
        ...
    
    def __contains__(self, key): # -> bool:
        ...
    
    def __iter__(self): # -> Iterator[Unknown]:
        ...
    
    def __len__(self): # -> int:
        ...
    


class ReadOnlyDict(Mapping):
    """A readonly dictionary wrapper."""
    def __init__(self, dict_like) -> None:
        ...
    
    def keys(self):
        ...
    
    def __getitem__(self, key):
        ...
    
    def __contains__(self, key): # -> bool:
        ...
    
    def __iter__(self): # -> Iterator[Unknown]:
        ...
    
    def __len__(self): # -> int:
        ...
    


def build_relabel_map(x, is_sorted=...): # -> tuple[Unknown, Unknown]:
    """Relabel the input ids to continuous ids that starts from zero.

    Ids are assigned new ids according to their ascending order.

    Examples
    --------
    >>> x = [1, 5, 3, 6]
    >>> n2o, o2n = build_relabel_map(x)
    >>> n2o
    [1, 3, 5, 6]
    >>> o2n
    [n/a, 0, n/a, 1, n/a, 2, 3]

    "n/a" will be filled with 0

    Parameters
    ----------
    x : Index
        The input ids.
    is_sorted : bool, default=False
        Whether the input has already been unique and sorted.

    Returns
    -------
    new_to_old : tensor
        The mapping from new id to old id.
    old_to_new : tensor
        The mapping from old id to new id. It is a vector of length MAX(x).
        One can use advanced indexing to convert an old id tensor to a
        new id tensor: new_id = old_to_new[old_id]
    """
    ...

def build_relabel_dict(x): # -> dict[Unknown, Unknown]:
    """Relabel the input ids to continuous ids that starts from zero.

    The new id follows the order of the given node id list.

    Parameters
    ----------
    x : list
      The input ids.

    Returns
    -------
    relabel_dict : dict
      Dict from old id to new id.
    """
    ...

class CtxCachedObject:
    """A wrapper to cache object generated by different context.

    Note: such wrapper may incur significant overhead if the wrapped object is very light.

    Parameters
    ----------
    generator : callable
        A callable function that can create the object given ctx as the only argument.
    """
    def __init__(self, generator) -> None:
        ...
    
    def __call__(self, ctx):
        ...
    


def cached_member(cache, prefix): # -> (func: Unknown) -> (self: Unknown, *args: Unknown, **kwargs: Unknown) -> Any:
    """A member function decorator to memorize the result.

    Note that the member function cannot support kwargs after being decorated.
    The member function must be functional. Otherwise, the behavior is undefined.

    Parameters
    ----------
    cache : str
        The cache name. The cache should be a dictionary attribute
        in the class object.
    prefix : str
        The key prefix to save the result of the function.
    """
    ...

def is_dict_like(obj): # -> bool:
    """Return true if the object can be treated as a dictionary."""
    ...

def reorder(dict_like, index): # -> dict[Unknown, Unknown]:
    """Reorder each column in the dict according to the index.

    Parameters
    ----------
    dict_like : dict of tensors
        The dict to be reordered.
    index : dgl.utils.Index
        The reorder index.
    """
    ...

def reorder_index(idx, order): # -> Index:
    """Reorder the idx according to the given order

    Parameters
    ----------
    idx : utils.Index
        The index to be reordered.
    order : utils.Index
        The order to follow.
    """
    ...

def is_iterable(obj): # -> bool:
    """Return true if the object is an iterable."""
    ...

def to_dgl_context(ctx): # -> DGLContext:
    """Convert a backend context to DGLContext"""
    ...

def to_nbits_int(tensor, nbits):
    """Change the dtype of integer tensor
    The dtype of returned tensor uses nbits, nbits can only be 32 or 64
    """
    ...

def make_invmap(array, use_numpy=...): # -> tuple[Unknown | list[Unknown], dict[Unknown, int], ndarray[Unknown, Unknown]]:
    """Find the unique elements of the array and return another array with indices
    to the array of unique elements."""
    ...

def expand_as_pair(input_, g=...): # -> tuple[Unknown, ...] | tuple[Mapping[Unknown, Unknown] | Unknown, dict[Unknown, Unknown] | Unknown] | tuple[Unknown, Unknown]:
    """Return a pair of same element if the input is not a pair.

    If the graph is a block, obtain the feature of destination nodes from the source nodes.

    Parameters
    ----------
    input_ : Tensor, dict[str, Tensor], or their pairs
        The input features
    g : DGLHeteroGraph or DGLGraph or None
        The graph.

        If None, skip checking if the graph is a block.

    Returns
    -------
    tuple[Tensor, Tensor] or tuple[dict[str, Tensor], dict[str, Tensor]]
        The features for input and output nodes
    """
    ...

def check_eq_shape(input_): # -> None:
    """If input_ is a pair of features, check if the feature shape of source
    nodes is equal to the feature shape of destination nodes.
    """
    ...

def retry_method_with_fix(fix_method): # -> (func: Unknown) -> (self: Unknown, *args: Unknown, **kwargs: Unknown) -> Unknown:
    """Decorator that executes a fix method before retrying again when the decorated method
    fails once with any exception.

    If the decorated method fails again, the execution fails with that exception.

    Notes
    -----
    This decorator only works on class methods, and the fix function must also be a class method.
    It would not work on functions.

    Parameters
    ----------
    fix_func : callable
        The fix method to execute.  It should not accept any arguments.  Its return values are
        ignored.
    """
    ...

def group_as_dict(pairs): # -> defaultdict[Unknown, list[Unknown]]:
    """Combines a list of key-value pairs to a dictionary of keys and value lists.

    Does not require the pairs to be sorted by keys.

    Parameters
    ----------
    pairs : iterable
        Iterable of key-value pairs

    Returns
    -------
    dict
        The dictionary of keys and value lists.
    """
    ...

class FlattenedDict:
    """Iterates over each item in a dictionary of groups.

    Parameters
    ----------
    groups : dict
        The item groups.

    Examples
    --------
    >>> groups = FlattenedDict({'a': [1, 3], 'b': [2, 5, 8], 'c': [7]})
    >>> list(groups)
    [('a', 1), ('a', 3), ('b', 2), ('b', 5), ('b', 8), ('c', 7)]
    >>> groups[2]
    ('b', 2)
    >>> len(groups)
    6
    """
    def __init__(self, groups) -> None:
        ...
    
    def __len__(self):
        """Return the total number of items."""
        ...
    
    def __iter__(self): # -> Generator[tuple[Unknown | int, Unknown], None, None]:
        """Return the iterator of all items with the key of its original group."""
        ...
    
    def __getitem__(self, idx): # -> tuple[Unknown | int, Unknown]:
        """Return the item at the given position with the key of its original group."""
        ...
    


def compensate(ids, origin_ids):
    """computing the compensate set of ids from origin_ids

    Note: ids should be a subset of origin_ids.
    Any of ids and origin_ids can be non-consecutive,
    and origin_ids should be sorted.

    Example:
    >>> ids = th.Tensor([0, 2, 4])
    >>> origin_ids = th.Tensor([0, 1, 2, 4, 5])
    >>> compensate(ids, origin_ids)
    th.Tensor([1, 5])
    """
    ...

def relabel(x): # -> tuple[Unknown, Unknown]:
    """Relabel the input ids to continuous ids that starts from zero.

    Ids are assigned new ids according to their ascending order.

    Examples
    --------
    >>> x = [1, 5, 3, 6]
    >>> n2o, o2n = build_relabel_map(x)
    >>> n2o
    [1, 3, 5, 6]
    >>> o2n
    [n/a, 0, n/a, 1, n/a, 2, 3]

    "n/a" will be filled with 0

    Parameters
    ----------
    x : Tensor
        ID tensor.

    Returns
    -------
    new_to_old : Tensor
        The mapping from new id to old id.
    old_to_new : Tensor
        The mapping from old id to new id. It is a vector of length MAX(x).
        One can use advanced indexing to convert an old id tensor to a
        new id tensor: new_id = old_to_new[old_id]
    """
    ...

def extract_node_subframes(graph, nodes, store_ids=...): # -> list[Unknown]:
    """Extract node features of the given nodes from :attr:`graph`
    and return them in frames.

    Note that this function does not perform actual tensor memory copy but using `Frame.subframe`
    to get the features. If :attr:`nodes` is None, it performs a shallow copy of the
    original node frames that only copies the dictionary structure but not the tensor
    contents.

    Parameters
    ----------
    graph : DGLGraph
        The graph to extract features from.
    nodes : list[Tensor] or None
        Node IDs. If not None, the list length must be equal to the number of node types
        in the graph. If None, the whole frame is shallow-copied.
    store_ids : bool
        If True, the returned frames will store :attr:`nodes` in the ``dgl.NID`` field
        unless it is None.

    Returns
    -------
    list[Frame]
        Extracted node frames.
    """
    ...

def extract_node_subframes_for_block(graph, srcnodes, dstnodes): # -> list[Unknown]:
    """Extract the input node features and output node features of the given nodes from
    :attr:`graph` and return them in frames ready for a block.

    Note that this function does not perform actual tensor memory copy but using `Frame.subframe`
    to get the features. If :attr:`srcnodes` or :attr:`dstnodes` is None, it performs a
    shallow copy of the original node frames that only copies the dictionary structure
    but not the tensor contents.

    Parameters
    ----------
    graph : DGLGraph
        The graph to extract features from.
    srcnodes : list[Tensor]
        Input node IDs. The list length must be equal to the number of node types
        in the graph. The returned frames store the node IDs in the ``dgl.NID`` field.
    dstnodes : list[Tensor]
        Output node IDs. The list length must be equal to the number of node types
        in the graph. The returned frames store the node IDs in the ``dgl.NID`` field.

    Returns
    -------
    list[Frame]
        Extracted node frames.
    """
    ...

def extract_edge_subframes(graph, edges, store_ids=...): # -> list[Unknown]:
    """Extract edge features of the given edges from :attr:`graph`
    and return them in frames.

    Note that this function does not perform actual tensor memory copy but using `Frame.subframe`
    to get the features. If :attr:`edges` is None, it performs a shallow copy of the
    original edge frames that only copies the dictionary structure but not the tensor
    contents.

    Parameters
    ----------
    graph : DGLGraph
        The graph to extract features from.
    edges : list[Tensor] or None
        Edge IDs. If not None, the list length must be equal to the number of edge types
        in the graph. If None, the whole frame is shallow-copied.
    store_ids : bool
        If True, the returned frames will store :attr:`edges` in the ``dgl.EID`` field
        unless it is None.

    Returns
    -------
    list[Frame]
        Extracted edge frames.
    """
    ...

def set_new_frames(graph, *, node_frames=..., edge_frames=...): # -> None:
    """Set the node and edge frames of a given graph to new ones.

    Parameters
    ----------
    graph : DGLGraph
        The graph whose node and edge frames are to be updated.
    node_frames : list[Frame], optional
        New node frames.

        Default is None, where the node frames are not updated.
    edge_frames : list[Frame], optional
        New edge frames

        Default is None, where the edge frames are not updated.
    """
    ...

def set_num_threads(num_threads): # -> None:
    """Set the number of OMP threads in the process.

    Parameters
    ----------
    num_threads : int
        The number of OMP threads in the process.
    """
    ...

def alias_func(func): # -> (*args: Unknown, **kwargs: Unknown) -> Unknown:
    """Return an alias function with proper docstring."""
    ...

