# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/graph/_graph_definition.pyx

"""GraphDefinition: explicit CUDA graph definition."""
from __future__ import annotations

from cuda.core._device import Device
from cuda.core._event import Event
from cuda.core._launch_config import LaunchConfig
from cuda.core._memory._buffer import Buffer
from cuda.core._module import Kernel
from cuda.core._utils.cuda_utils import driver
from cuda.core.graph._graph_builder import (Graph, GraphCompleteOptions,
                                            GraphDebugPrintOptions)
from cuda.core.graph._graph_node import GraphNode
from cuda.core.graph._subclasses import (AllocNode, ChildGraphNode, EmptyNode,
                                         EventRecordNode, EventWaitNode,
                                         FreeNode, HostCallbackNode,
                                         IfElseNode, IfNode, KernelNode,
                                         MemcpyNode, MemsetNode, SwitchNode,
                                         WhileNode)
from cuda.core.typing import GraphMemoryType


class GraphCondition:
    """A condition variable for conditional graph nodes.

    Created by :meth:`GraphDefinition.create_condition` (or
    :meth:`GraphBuilder.create_condition`) and passed to
    conditional-node builder methods (:meth:`~GraphDefinition.if_then`,
    :meth:`~GraphDefinition.if_else`, :meth:`~GraphDefinition.while_loop`,
    :meth:`~GraphDefinition.switch`). The underlying value is set at
    runtime by device code via ``cudaGraphSetConditional``.

    A :class:`GraphCondition` may be passed directly as a kernel
    argument to ``launch()``: the launcher unwraps it to the underlying
    ``CUgraphConditionalHandle`` value so device code can update the
    condition.
    """

    def __repr__(self) -> str:
        ...

    def __eq__(self, other: object) -> bool:
        ...

    def __hash__(self) -> int:
        ...

    @property
    def handle(self) -> driver.CUgraphConditionalHandle:
        """The raw CUgraphConditionalHandle as an int."""

class GraphDefinition:
    """A graph definition.

    A GraphDefinition is used to construct a graph explicitly by adding nodes
    and specifying dependencies. Once construction is complete, call
    instantiate() to obtain an executable Graph.
    """

    def __init__(self):
        """Create a new empty graph definition."""

    def __repr__(self) -> str:
        ...

    def __eq__(self, other: object) -> bool:
        ...

    def __hash__(self) -> int:
        ...

    @property
    def _entry(self) -> GraphNode:
        """Return the internal entry-point GraphNode (no dependencies)."""

    def allocate(self, size: int, *, device: Device | int | None=None, memory_type: GraphMemoryType=..., peer_access: list[Device | int] | None=None) -> AllocNode:
        """Add an entry-point memory allocation node (no dependencies).

        See :meth:`GraphNode.allocate` for full documentation.
        """

    def deallocate(self, dptr: int) -> FreeNode:
        """Add an entry-point memory free node (no dependencies).

        See :meth:`GraphNode.deallocate` for full documentation.
        """

    def memset(self, dst: Buffer | int, value, width: int, height: int=1, pitch: int=0, *, dst_owner=None) -> MemsetNode:
        """Add an entry-point memset node (no dependencies).

        See :meth:`GraphNode.memset` for full documentation.
        """

    def launch(self, config: LaunchConfig, kernel: Kernel, *args) -> KernelNode:
        """Add an entry-point kernel launch node (no dependencies).

        See :meth:`GraphNode.launch` for full documentation.
        """

    def empty(self) -> EmptyNode:
        """Add an entry-point empty node (no dependencies).

        Returns
        -------
        EmptyNode
            A new EmptyNode with no dependencies.
        """

    def join(self, *nodes: GraphNode) -> EmptyNode:
        """Create an empty node that depends on all given nodes.

        Parameters
        ----------
        *nodes : GraphNode
            Nodes to merge.

        Returns
        -------
        EmptyNode
            A new EmptyNode that depends on all input nodes.
        """

    def memcpy(self, dst: Buffer | int, src: Buffer | int, size: int, *, dst_owner=None, src_owner=None) -> MemcpyNode:
        """Add an entry-point memcpy node (no dependencies).

        See :meth:`GraphNode.memcpy` for full documentation.
        """

    def embed(self, child: GraphDefinition) -> ChildGraphNode:
        """Add an entry-point child graph node (no dependencies).

        See :meth:`GraphNode.embed` for full documentation.
        """

    def record(self, event: Event) -> EventRecordNode:
        """Add an entry-point event record node (no dependencies).

        See :meth:`GraphNode.record` for full documentation.
        """

    def wait(self, event: Event) -> EventWaitNode:
        """Add an entry-point event wait node (no dependencies).

        See :meth:`GraphNode.wait` for full documentation.
        """

    def callback(self, fn, *, user_data=None) -> HostCallbackNode:
        """Add an entry-point host callback node (no dependencies).

        See :meth:`GraphNode.callback` for full documentation.
        """

    def create_condition(self, default_value: int | None=None) -> GraphCondition:
        """Create a condition variable for use with conditional nodes.

        The returned :class:`GraphCondition` object is passed to conditional-node
        builder methods. Its value is controlled at runtime by device code
        via ``cudaGraphSetConditional``.

        Parameters
        ----------
        default_value : int, optional
            The default value to assign to the condition.
            If None, no default is assigned.

        Returns
        -------
        GraphCondition
            A condition variable for controlling conditional execution.
        """

    def if_then(self, condition: GraphCondition) -> IfNode:
        """Add an entry-point if-conditional node (no dependencies).

        See :meth:`GraphNode.if_then` for full documentation.
        """

    def if_else(self, condition: GraphCondition) -> IfElseNode:
        """Add an entry-point if-else conditional node (no dependencies).

        See :meth:`GraphNode.if_else` for full documentation.
        """

    def while_loop(self, condition: GraphCondition) -> WhileNode:
        """Add an entry-point while-loop conditional node (no dependencies).

        See :meth:`GraphNode.while_loop` for full documentation.
        """

    def switch(self, condition: GraphCondition, count: int) -> SwitchNode:
        """Add an entry-point switch conditional node (no dependencies).

        See :meth:`GraphNode.switch` for full documentation.
        """

    def instantiate(self, options: GraphCompleteOptions | None=None) -> Graph:
        """Instantiate the graph definition into an executable Graph.

        Parameters
        ----------
        options : :obj:`~graph.GraphCompleteOptions`, optional
            Customizable dataclass for graph instantiation options.

        Returns
        -------
        Graph
            An executable graph that can be launched on a stream.
        """

    def debug_dot_print(self, path: str, options: GraphDebugPrintOptions | None=None) -> None:
        """Write a GraphViz DOT representation of the graph to a file.

        Parameters
        ----------
        path : str
            File path for the DOT output.
        options : GraphDebugPrintOptions, optional
            Customizable options for the debug print.
        """

    def nodes(self) -> set[GraphNode]:
        """Return all nodes in the graph.

        Returns
        -------
        set of GraphNode
            All nodes in the graph.
        """

    def edges(self) -> set[tuple[GraphNode, GraphNode]]:
        """Return all edges in the graph as (from_node, to_node) pairs.

        Returns
        -------
        set of tuple
            Each element is a (from_node, to_node) pair representing
            a dependency edge in the graph.
        """

    @property
    def handle(self) -> driver.CUgraph:
        """Return the underlying driver CUgraph handle."""
__all__ = ['GraphCondition', 'GraphDefinition']