# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# coding: utf-8
"""Context management API of mxnet."""
from __future__ import absolute_import
import threading
import warnings
import ctypes
from .base import classproperty, with_metaclass, _MXClassPropertyMetaClass
from .base import _LIB
from .base import check_call


class Context(with_metaclass(_MXClassPropertyMetaClass, object)):
    """Constructs a context.

    MXNet can run operations on CPU and different GPUs.
    A context describes the device type and ID on which computation should be carried on.

    One can use mx.cpu and mx.gpu for short.

    See also
    ----------
    `How to run MXNet on multiple CPU/GPUs <http://mxnet.incubator.apache.org/api/faq/distributed_training>`
    for more details.

    Parameters
    ----------
    device_type : {'cpu', 'gpu'} or Context.
        String representing the device type.

    device_id : int (default=0)
        The device id of the device, needed for GPU.

    Note
    ----
    Context can also be used as a way to change the default context.

    Examples
    --------
    >>> # array on cpu
    >>> cpu_array = mx.nd.ones((2, 3))
    >>> # switch default context to GPU(2)
    >>> with mx.Context(mx.gpu(2)):
    ...     gpu_array = mx.nd.ones((2, 3))
    >>> gpu_array.context
    gpu(2)

    One can also explicitly specify the context when creating an array.

    >>> gpu_array = mx.nd.ones((2, 3), mx.gpu(1))
    >>> gpu_array.context
    gpu(1)
    """
    # static class variable
    _default_ctx = threading.local()
    devtype2str = {1: 'cpu', 2: 'gpu', 3: 'cpu_pinned', 5: 'cpu_shared'}
    devstr2type = {'cpu': 1, 'gpu': 2, 'cpu_pinned': 3, 'cpu_shared': 5}
    def __init__(self, device_type, device_id=0):
        if isinstance(device_type, Context):
            self.device_typeid = device_type.device_typeid
            self.device_id = device_type.device_id
        else:
            self.device_typeid = Context.devstr2type[device_type]
            self.device_id = device_id
        self._old_ctx = None

    @property
    def device_type(self):
        """Returns the device type of current context.

        Examples
        -------
        >>> mx.context.current_context().device_type
        'cpu'
        >>> mx.current_context().device_type
        'cpu'

        Returns
        -------
        device_type : str
        """
        return Context.devtype2str[self.device_typeid]

    def __hash__(self):
        """Compute hash value of context for dictionary lookup"""
        return hash((self.device_typeid, self.device_id))

    def __eq__(self, other):
        """Compares two contexts. Two contexts are equal if they
        have the same device type and device id.
        """
        return isinstance(other, Context) and \
            self.device_typeid == other.device_typeid and \
            self.device_id == other.device_id

    def __str__(self):
        return '%s(%d)' % (self.device_type, self.device_id)

    def __repr__(self):
        return self.__str__()

    def __enter__(self):
        if not hasattr(Context._default_ctx, "value"):
            Context._default_ctx.value = Context('cpu', 0)
        self._old_ctx = Context._default_ctx.value
        Context._default_ctx.value = self
        return self

    def __exit__(self, ptype, value, trace):
        Context._default_ctx.value = self._old_ctx

    #pylint: disable=no-self-argument
    @classproperty
    def default_ctx(cls):
        warnings.warn("Context.default_ctx has been deprecated. "
                      "Please use Context.current_context() instead. "
                      "Please use test_utils.set_default_context to set a default context",
                      DeprecationWarning)
        if not hasattr(Context._default_ctx, "value"):
            cls._default_ctx.value = Context('cpu', 0)
        return cls._default_ctx.value

    @default_ctx.setter
    def default_ctx(cls, val):
        warnings.warn("Context.default_ctx has been deprecated. "
                      "Please use Context.current_context() instead. "
                      "Please use test_utils.set_default_context to set a default context",
                      DeprecationWarning)
        cls._default_ctx.value = val
    #pylint: enable=no-self-argument

    def empty_cache(self):
        """Empties the memory cache for the current contexts device.

        MXNet utilizes a memory pool to avoid excessive allocations.
        Calling empty_cache will empty the memory pool of the contexts
        device. This will only free the memory of the unreferenced data.

        Examples
        -------
        >>> ctx = mx.gpu(0)
        >>> arr = mx.nd.ones((200,200), ctx=ctx)
        >>> del arr
        >>> ctx.empty_cache() # forces release of memory allocated for arr
        """
        dev_type = ctypes.c_int(self.device_typeid)
        dev_id = ctypes.c_int(self.device_id)
        check_call(_LIB.MXStorageEmptyCache(dev_type, dev_id))

# initialize the default context in Context
Context._default_ctx.value = Context('cpu', 0)


def cpu(device_id=0):
    """Returns a CPU context.

    This function is a short cut for ``Context('cpu', device_id)``.
    For most operations, when no context is specified, the default context is `cpu()`.

    Examples
    ----------
    >>> with mx.cpu():
    ...     cpu_array = mx.nd.ones((2, 3))
    >>> cpu_array.context
    cpu(0)
    >>> cpu_array = mx.nd.ones((2, 3), ctx=mx.cpu())
    >>> cpu_array.context
    cpu(0)

    Parameters
    ----------
    device_id : int, optional
        The device id of the device. `device_id` is not needed for CPU.
        This is included to make interface compatible with GPU.

    Returns
    -------
    context : Context
        The corresponding CPU context.
    """
    return Context('cpu', device_id)


def cpu_pinned(device_id=0):
    """Returns a CPU pinned memory context. Copying from CPU pinned memory to GPU
    is faster than from normal CPU memory.

    This function is a short cut for ``Context('cpu_pinned', device_id)``.

    Examples
    ----------
    >>> with mx.cpu_pinned():
    ...     cpu_array = mx.nd.ones((2, 3))
    >>> cpu_array.context
    cpu_pinned(0)
    >>> cpu_array = mx.nd.ones((2, 3), ctx=mx.cpu_pinned())
    >>> cpu_array.context
    cpu_pinned(0)

    Parameters
    ----------
    device_id : int, optional
        The device id of the device. `device_id` is not needed for CPU.
        This is included to make interface compatible with GPU.

    Returns
    -------
    context : Context
        The corresponding CPU pinned memory context.
    """
    return Context('cpu_pinned', device_id)


def gpu(device_id=0):
    """Returns a GPU context.

    This function is a short cut for Context('gpu', device_id).
    The K GPUs on a node are typically numbered as 0,...,K-1.

    Examples
    ----------
    >>> cpu_array = mx.nd.ones((2, 3))
    >>> cpu_array.context
    cpu(0)
    >>> with mx.gpu(1):
    ...     gpu_array = mx.nd.ones((2, 3))
    >>> gpu_array.context
    gpu(1)
    >>> gpu_array = mx.nd.ones((2, 3), ctx=mx.gpu(1))
    >>> gpu_array.context
    gpu(1)

    Parameters
    ----------
    device_id : int, optional
        The device id of the device, needed for GPU.

    Returns
    -------
    context : Context
        The corresponding GPU context.
    """
    return Context('gpu', device_id)


def num_gpus():
    """Query CUDA for the number of GPUs present.

    Raises
    ------
    Will raise an exception on any CUDA error.

    Returns
    -------
    count : int
        The number of GPUs.

    """
    count = ctypes.c_int()
    check_call(_LIB.MXGetGPUCount(ctypes.byref(count)))
    return count.value


def gpu_memory_info(device_id=0):
    """Query CUDA for the free and total bytes of GPU global memory.

    Parameters
    ----------
    device_id : int, optional
        The device id of the GPU device.

    Raises
    ------
    Will raise an exception on any CUDA error.

    Returns
    -------
    (free, total) : (int, int)
    """
    free = ctypes.c_uint64()
    total = ctypes.c_uint64()
    dev_id = ctypes.c_int(device_id)
    check_call(_LIB.MXGetGPUMemoryInformation64(dev_id, ctypes.byref(free), ctypes.byref(total)))
    return (free.value, total.value)


def current_context():
    """Returns the current context.

    By default, `mx.cpu()` is used for all the computations
    and it can be overridden by using `with mx.Context(x)` statement where
    x can be cpu(device_id) or gpu(device_id).

    Examples
    -------
    >>> mx.current_context()
    cpu(0)
    >>> with mx.Context('gpu', 1):  # Context changed in `with` block.
    ...    mx.current_context()  # Computation done here will be on gpu(1).
    ...
    gpu(1)
    >>> mx.current_context() # Back to default context.
    cpu(0)

    Returns
    -------
    default_ctx : Context
    """
    if not hasattr(Context._default_ctx, "value"):
        Context._default_ctx.value = Context('cpu', 0)
    return Context._default_ctx.value