# mypy: ignore-errors

r"""
This package adds support for JIT compilation for CUDA Streams and events,
This is similar to API's available in the eager mode
:ref:`cuda-semantics` has more details about working with CUDA.
"""

import torch
from typing import Optional, Any
from torch import device as _device

def get_current_device_index() -> int:
    r"""Checks if there are CUDA devices available and
    returns the device index of the current default CUDA device.
    Returns -1 in case there are no CUDA devices available.

    Arguments: ``None``
    """
    if torch.cuda.device_count() > 0:
        return torch.cuda._current_device()
    return -1

def get_device_index(device: Optional[_device] = None, optional: bool = False, allow_cpu: bool = False) -> int:
    r"""Gets the device index from :attr:`device`, which can be a torch.device
    object, a Python integer, or ``None``.

    If :attr:`device` is a torch.device object, returns the device index if it
    is a CUDA device. Note that for a CUDA device without a specified index,
    , this will return the current default CUDA device if :attr:`optional` is ``True``.
    If :attr:`allow_cpu` is ``True``,CPU devices will be accepted and ``-1`` will be
    returned in this case.

    If :attr:`device` is a Python integer, it is returned as is.

    If :attr:`device` is ``None``, this will return the current default CUDA
    device if :attr:`optional` is ``True``.
    """
    if device is None:
        if optional:
            return get_current_device_index()
        else:
            raise ValueError('Expected a torch.device with a specified index '
                             f'or an integer, but got: {device}')
    device_index = -1
    if isinstance(device, str):
        device = torch.device(device)

    if isinstance(device, torch.device):
        if not allow_cpu and device.type == 'cpu':
            raise ValueError(f'Expected a non cpu device, but got: {device}')
        device_index = -1 if device.type == 'cpu' else torch.cuda.device_index(device)

    if isinstance(device, int):
        device_index = device

    return device_index

class device(object):
    r"""Context-manager that changes the selected device.
    This is similar to device (torch.device or int), but has been
    introduced for JIT compatibility.
    Arguments:
        device (torch.device or int): device index to select. It's a no-op if
            this argument is a negative integer or ``None``.
    """
    def __init__(self, device: Optional[_device]):
        self.idx = -1
        self.prev_idx = -1
        self.device = device

    def __enter__(self):
        self.idx = get_device_index(self.device, optional=True)

        if self.idx == -1:
            return
        self.prev_idx = torch.cuda._current_device()

        if self.prev_idx != self.idx:
            torch.cuda._set_device(self.idx)

    def __exit__(self, type: Any, value: Any, traceback: Any):
        if self.prev_idx != self.idx:
            torch.cuda._set_device(self.prev_idx)

class StreamContext(object):
    r"""Context-manager that selects a given stream.
    All CUDA kernels queued within its context will be enqueued on a selected
    stream.
    Arguments:
        StreamContext (Stream): selected stream. This manager is a no-op if it's
            ``None``.
    .. note:: Streams are per-device. If the selected stream is not on the
        current device, this function will also change the current device to
        match the stream.
    """
    cur_stream : Optional['torch.classes.cuda.Stream']

    def __init__(self, stream: Optional['torch.classes.cuda.Stream']):
        self.idx = -1
        self.stream = stream
        # Initialize the below streams to default stream on the current device
        self.device_index = get_current_device_index()
        self.src_prev_stream = torch.cuda.default_stream(self.device_index)
        self.dst_prev_stream = torch.cuda.default_stream(self.device_index)

    def __enter__(self):
        self.idx = get_device_index(device=None, optional=True)
        # If there is no CUDA device available, return
        if self.idx == -1:
            return

        # Local cur_stream variable for type refinement
        cur_stream = self.stream
        # Return if stream is None
        if cur_stream is None:
            return
        self.src_prev_stream = torch.cuda.current_stream(self.idx)
        # If the stream is not on the current device, then change the device
        # and set the current stream on the device
        if self.src_prev_stream.device_index() != cur_stream.device_index():
            with device(cur_stream.device()):
                self.dst_prev_stream = torch.cuda.current_stream(cur_stream.device_index())
            torch.cuda._set_device(cur_stream.device_index())
        torch.cuda.set_stream(cur_stream)

    def __exit__(self, type: Any, value: Any, traceback: Any):
        # Local cur_stream variable for type refinement
        cur_stream = self.stream
        # If stream is None or no CUDA device available, return
        if cur_stream is None or self.idx == -1:
            return
        # If the stream was not on the current device, restore the previous stream on
        # the destination device and also reset the current device to the previous device.
        # Set the current stream on the device to the src_prev_stream
        if self.src_prev_stream.device_index() != cur_stream.device_index():
            torch.cuda.set_stream(self.dst_prev_stream)
            torch.cuda._set_device(self.idx)
        torch.cuda.set_stream(self.src_prev_stream)

def stream(stream: Optional['torch.classes.cuda.Stream']) -> StreamContext:
    r"""Wrapper around the Context-manager that selects a given stream.
    All CUDA kernels queued within its context will be enqueued on a selected
    stream.
    Arguments:
        stream (Stream): selected stream. This manager is a no-op if it's
            ``None``.
    """
    return StreamContext(stream)

def Stream(device: int = -1, priority: int = 0) -> 'torch.classes.cuda.Stream':
    r"""Wrapper around a CUDA stream.
    A CUDA stream is a linear sequence of execution that belongs to a specific
    device, independent from other streams.  See :ref:`cuda-semantics` for
    details.
    Arguments:
        device(int, optional): a device on which to allocate
            the stream. If :attr:`device` is ``None`` (default) or a negative
            integer, this will use the current device.
        priority(int, optional): priority of the stream. Can be either
            -1 (high priority) or 0 (low priority). By default, streams have
            priority 0.
    .. note:: Although CUDA versions >= 11 support more than two levels of
        priorities, in PyTorch, we only support two levels of priorities.
    """
    return torch.classes.cuda.Stream(device, priority)

def Event(enable_timing: bool = False, blocking: bool = False, interprocess: bool = False) -> 'torch.classes.cuda.Event':
    r"""Wrapper around a CUDA event.
    CUDA events are synchronization markers that can be used to monitor the
    device's progress, to accurately measure timing, and to synchronize CUDA
    streams.
    Arguments:
        enable_timing (bool, optional): indicates if the event should measure time
            (default: ``False``)
        blocking (bool, optional): if ``True``, :meth:`wait` will be blocking (default: ``False``)
        interprocess (bool): if ``True``, the event can be shared between processes
            (default: ``False``)
    .. _CUDA Event Documentation:
       https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EVENT.html
    """
    return torch.classes.cuda.Event(enable_timing, blocking, interprocess)